!> This module defines a data structure to store process local work assignment
module MW_localwork
   use MW_electrode, only: MW_electrode_t
   use MW_ion, only: MW_ion_t
   use MW_constants, only: pair_block_size
   implicit none
   private

   ! public types
   public :: MW_localwork_t

   ! public subroutines
   public :: define_type
   public :: void_type
   public :: print_type
   public :: setup_ewald_local_work
   public :: setup_pair_atom2atom_local_work
   public :: setup_pair_ion2ion_local_work
   public :: setup_pair_ion2ion_local_work_ms
   public :: setup_pair_ion2atom_local_work
   public :: setup_pair_ion2atom_local_work_ms
   public :: setup_pair_atom2ion_local_work
   public :: setup_atoms_local_work
   public :: setup_ions_local_work
   public :: setup_molecules_local_work
   public :: update_diag_block_boundaries
   public :: update_tri_block_boundaries
   public :: update_other_block_boundaries

   ! A data structure to hold comm, rank and group for a process
   type MW_localwork_t

      ! Parallel distribution of computation of Ewald reciprocal space sum
      integer :: ewald_num_mode_global = 0
      integer :: ewald_num_mode_local = 0
      integer :: ewald_local_mode_offset = 0
      integer :: ewald_kstart_x = 0
      integer :: ewald_kstart_y = 0
      integer :: ewald_kstart_z = 0

      ! Atom 2 Atom interactions
      ! ------------------------
      ! Parallel distribution of computation of pair interaction
      integer, allocatable :: pair_atom2atom_num_block_diag_local(:)
      integer, allocatable :: pair_atom2atom_num_block_full_local(:,:)
      integer, allocatable :: pair_atom2atom_diag_iblock_offset(:)
      integer, allocatable :: pair_atom2atom_full_iblock_offset(:,:)

      ! Ion 2 Ion interactions
      ! ----------------------
      ! Parallel distribution of computation of pair interaction
      integer, allocatable :: pair_ion2ion_num_block_diag_local(:)
      integer, allocatable :: pair_ion2ion_num_block_full_local(:,:)
      integer, allocatable :: pair_ion2ion_diag_iblock_offset(:)
      integer, allocatable :: pair_ion2ion_full_iblock_offset(:,:)

      integer :: pair_ion2ion_num_block_diag_local_ms
      integer :: pair_ion2ion_num_block_full_local_ms
      integer :: pair_ion2ion_diag_iblock_offset_ms
      integer :: pair_ion2ion_full_iblock_offset_ms

      ! Ion 2 Atoms interactions
      ! ------------------------
      ! Parallel distribution of computation of pair interaction
      integer, allocatable :: pair_ion2atom_num_block_full_local(:,:)
      integer, allocatable :: pair_ion2atom_full_iblock_offset(:,:)

      integer :: pair_ion2atom_num_block_full_local_ms
      integer :: pair_ion2atom_full_iblock_offset_ms

      ! Atom 2 Ion interactions
      ! ------------------------
      ! Parallel distribution of computation of pair interaction
      integer, allocatable :: pair_atom2ion_num_block_full_local(:,:)
      integer, allocatable :: pair_atom2ion_full_iblock_offset(:,:)

      ! Parallel distribution of atoms
      integer :: atoms_num_global = 0
      integer :: atoms_num_local = 0
      integer, allocatable :: offset_atoms(:) !< offset atoms for local work per type
      integer, allocatable :: count_atoms(:)  !< count atoms for local work per type

      ! Parallel distribution of ions
      integer :: ions_num_global = 0
      integer :: ions_num_local = 0
      integer, allocatable :: offset_ions(:) !< offset ions for local work per type
      integer, allocatable :: count_ions(:) !< offset ions for local work per type

      ! Parallel distribution of molecules
      integer :: molecules_num_global = 0
      integer :: molecules_num_local = 0
      integer, allocatable :: molecules_istart(:) !< start index for each molecule type
      integer, allocatable :: molecules_iend(:) !< end index for each molecule type

   end type MW_localwork_t

contains
   !================================================================================
   ! Define a parallel data structure
   subroutine define_type(this)
      implicit none
      type(MW_localwork_t), intent(inout) :: this

      ! Parallel distribution of computation of Ewald reciprocal space sum
      this%ewald_num_mode_global = 0
      this%ewald_num_mode_local = 0
      this%ewald_local_mode_offset = 0
      this%ewald_kstart_x = 0
      this%ewald_kstart_y = 0
      this%ewald_kstart_z = 0

      ! Parallel distribution of atoms
      this%atoms_num_global = 0
      this%atoms_num_local = 0

      ! Parallel distribution of ions
      this%ions_num_global = 0
      this%ions_num_local = 0

      ! Parallel distribution of molecules
      this%molecules_num_global = 0
      this%molecules_num_local = 0

   end subroutine define_type

   !================================================================================
   ! Void a parallel data structure
   subroutine void_type(this)
      use MW_errors, only: MW_errors_deallocate_error => deallocate_error
      implicit none
      type(MW_localwork_t), intent(inout) :: this
      integer :: ierr

      if (allocated(this%offset_ions)) then
         deallocate(this%offset_ions, this%count_ions, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      if (allocated(this%molecules_istart)) then
         deallocate(this%molecules_istart, this%molecules_iend, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      if (allocated(this%offset_atoms)) then
         deallocate(this%offset_atoms, this%count_atoms, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      ! Parallel distribution of computation of Ewald reciprocal space sum
      this%ewald_num_mode_global = 0
      this%ewald_num_mode_local = 0
      this%ewald_local_mode_offset = 0
      this%ewald_kstart_x = 0
      this%ewald_kstart_y = 0
      this%ewald_kstart_z = 0

      ! Parallel distribution of computation of pair interaction
      if (allocated(this%pair_atom2atom_diag_iblock_offset)) then
         deallocate(this%pair_atom2atom_diag_iblock_offset, &
               this%pair_atom2atom_full_iblock_offset, &
               this%pair_atom2atom_num_block_diag_local, &
               this%pair_atom2atom_num_block_full_local, &
               stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      if (allocated(this%pair_ion2ion_diag_iblock_offset)) then
         deallocate(this%pair_ion2ion_diag_iblock_offset, &
               this%pair_ion2ion_full_iblock_offset, &
               this%pair_ion2ion_num_block_diag_local, &
               this%pair_ion2ion_num_block_full_local, &
               stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      if (allocated(this%pair_ion2atom_full_iblock_offset)) then
         deallocate(this%pair_ion2atom_full_iblock_offset, &
               this%pair_ion2atom_num_block_full_local, &
               stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      if (allocated(this%pair_atom2ion_full_iblock_offset)) then
         deallocate(this%pair_atom2ion_full_iblock_offset, &
               this%pair_atom2ion_num_block_full_local, &
               stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "localwork.f90", ierr)
         end if
      end if

      ! Parallel distribution of atoms
      this%atoms_num_global = 0
      this%atoms_num_local = 0

      ! Parallel distribution of ions
      this%ions_num_global = 0
      this%ions_num_local = 0

      ! Parallel distribution of molecules
      this%molecules_num_global = 0
      this%molecules_num_local = 0

   end subroutine void_type

   ! ================================================================================
   ! Print type
   subroutine print_type(this, ounit)
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(in) :: this
      integer, intent(in) :: ounit

      ! Local
      ! -----
      integer :: itype, jtype

      write(ounit,'("|localwork| k-modes assignment")')
      write(ounit,'("|localwork|    number of global k-modes: ", i8)') this%ewald_num_mode_global
      write(ounit,'("|localwork|    number of local  k-modes: ", i8)') this%ewald_num_mode_local
      write(ounit,'("|localwork|    offset of local  k-modes: ", i6)') this%ewald_local_mode_offset
      write(ounit,'("|localwork|    start k-mode index      : ", 3(i6,:,", "))') &
            this%ewald_kstart_x, this%ewald_kstart_y, this%ewald_kstart_z
      write(ounit,*)
      write(ounit,'("|localwork| electrode to electrode pair diagonal block assignment:")')
      write(ounit,'("|localwork| itype |  count | offset |")')
      do itype = 1, size(this%pair_atom2atom_num_block_diag_local,1)
         write(ounit,'("|localwork| ",i5," | ",i6," | ",i6," |")') &
              itype, this%pair_atom2atom_num_block_diag_local(itype), &
              this%pair_atom2atom_diag_iblock_offset(itype)
      end do
      write(ounit,*)
      write(ounit,'("|localwork| electrode to electrode pair block assignment:")')
      write(ounit,'("|localwork| itype | jtype |  count | offset |")')
      do itype = 1, size(this%pair_atom2atom_num_block_full_local,2)
         do jtype = 1, size(this%pair_atom2atom_num_block_full_local,1)
            write(ounit,'("|localwork| ",i5," | ",i5," | ",i6," | ",i6," |")') &
                 itype, jtype, this%pair_atom2atom_num_block_full_local(jtype,itype), &
                 this%pair_atom2atom_full_iblock_offset(jtype,itype)
         end do
      end do

      write(ounit,*)
      write(ounit,'("|localwork| electrolyte to electrolyte pair diagonal block assignment:")')
      write(ounit,'("|localwork| itype |  count | offset |")')
      do itype = 1, size(this%pair_ion2ion_num_block_diag_local,1)
         write(ounit,'("|localwork| ",i5," | ",i6," | ",i6," |")') &
              itype, this%pair_ion2ion_num_block_diag_local(itype), &
              this%pair_ion2ion_diag_iblock_offset(itype)
      end do
      write(ounit,*)
      write(ounit,'("|localwork| electrolyte to electrolyte pair block assignment:")')
      write(ounit,'("|localwork| itype | jtype |  count | offset |")')
      do itype = 1, size(this%pair_ion2ion_num_block_full_local,2)
         do jtype = 1, size(this%pair_ion2ion_num_block_full_local,1)
            write(ounit,'("|localwork| ",i5," | ",i5," | ",i6," | ",i6," |")') &
                 itype, jtype, this%pair_ion2ion_num_block_full_local(jtype,itype), &
                 this%pair_ion2ion_full_iblock_offset(jtype,itype)
         end do
      end do

      write(ounit,*)
      write(ounit,'("|localwork| electrode to electrolyte pair block assignment:")')
      write(ounit,'("|localwork| itype | jtype |  count | offset |")')
      do itype = 1, size(this%pair_atom2ion_num_block_full_local,2)
         do jtype = 1, size(this%pair_atom2ion_num_block_full_local,1)
            write(ounit,'("|localwork| ",i5," | ",i5," | ",i6," | ",i6," |")') &
                 itype, jtype, this%pair_atom2ion_num_block_full_local(jtype,itype), &
                 this%pair_atom2ion_full_iblock_offset(jtype,itype)
         end do
      end do

      write(ounit,*)
      write(ounit,'("|localwork| electrolyte to electrode pair block assignment:")')
      write(ounit,'("|localwork| itype | jtype |  count | offset |")')
      do itype = 1, size(this%pair_ion2atom_num_block_full_local,2)
         do jtype = 1, size(this%pair_ion2atom_num_block_full_local,1)
            write(ounit,'("|localwork| ",i5," | ",i5," | ",i6," | ",i6," |")') &
                 itype, jtype, this%pair_ion2atom_num_block_full_local(jtype,itype), &
                 this%pair_ion2atom_full_iblock_offset(jtype,itype)
         end do
      end do

      write(ounit,*)
      write(ounit,'("|localwork| electrode vector assignment")')
      write(ounit,'("|localwork|    number of global electrode atoms: ",i6)') &
            this%atoms_num_global
      write(ounit,'("|localwork|    number of local electrode atoms: ",i6)') &
            this%atoms_num_local
      write(ounit,'("|localwork|    electrode atoms local offset:     ",100(i6,:,", "))') &
            this%offset_atoms
      write(ounit,'("|localwork|    electrode atoms local count:      ",100(i6,:,", "))') &
            this%count_atoms

      write(ounit,*)
      write(ounit,'("|localwork| electrolyte vector assignment")')
      write(ounit,'("|localwork|    number of global electrolyte atoms: ",i6)') &
            this%ions_num_global
      write(ounit,'("|localwork|    number of local electrolyte atoms: ",i6)') &
            this%ions_num_local
      write(ounit,'("|localwork|    electrolyte atoms local offset:     ",100(i6,", ",:))') &
            this%offset_ions
      write(ounit,'("|localwork|    electrolyte atoms local count :     ",100(i6,", ",:))') &
            this%count_ions

      if (size(this%molecules_istart,1) > 0) then
         write(ounit,*)
         write(ounit,'("|localwork| molecules vector assignment")')
         write(ounit,'("|localwork|    number of global molecules: ",i6)') &
               this%molecules_num_global
         write(ounit,'("|localwork|    number of local molecules: ",i6)') &
               this%molecules_num_local
         write(ounit,'("|localwork|    molecules local istart:     ",100(i6,:,", "))') &
               this%molecules_istart
         write(ounit,'("|localwork|    molecules local iend:       ",100(i6,:,", "))') &
               this%molecules_iend
      end if

   end subroutine print_type

   !================================================================================
   ! Setup parallel distribution of work in Ewald reciprocal sum
   ! The sum over k-modes is symmetric
   !   -> we loop only on half of them and mulitply by 2.0
   ! Total number of modes
   !    (2*kmax_x+1)*(2*kmax_y+1)*nz
   ! Num modes in the exectude globally
   !    (2*kmax_x*kmax_y + kmax_x + kmax_y)*nz
   ! Num modes executed by the local process (approx.)
   !    (2*kmax_x*kmax_y + kmax_x + kmax_y)*nz / num_rank
   subroutine setup_ewald_local_work(this, ewald, num_pbc, box, comm_size, comm_rank)
      use MW_kinds, only: wp
      use MW_constants, only: twopi
      use MW_ewald, only: MW_ewald_t, &
            MW_ewald_update_kmode_index => update_kmode_index
      use MW_box, only: MW_box_t
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_ewald_t), intent(in) :: ewald
      integer, intent(in) :: num_pbc
      type(MW_box_t), intent(in) :: box
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank
      ! Local
      ! -----
      integer :: imode, l, m, n, kstart_x, kstart_y, kstart_z
      integer :: num_modes_within_cutoff
      integer :: num_modes_global, num_modes_global_within_cutoff
      integer :: num_modes_local, num_modes_local_within_cutoff
      integer :: num_modes_global_before, imode_start
      real(wp) :: box_lengthx_rec, box_lengthy_rec
      real(wp) :: k(3), knorm2

      num_modes_local = 0
      kstart_x = 0
      kstart_y = 0
      kstart_z = 0

      select case (num_pbc)
      case (2)
         num_modes_global = ewald%kmax_x*(2*ewald%kmax_y+1)*(2*ewald%kmax_z+1) + &
              ewald%kmax_y*(2*ewald%kmax_z+1)
      case (3)
         num_modes_global = ewald%kmax_x*(2*ewald%kmax_y+1)*(2*ewald%kmax_z+1) + &
              ewald%kmax_y*(2*ewald%kmax_z+1) + ewald%kmax_z
      case default
         num_modes_global = 0
      end select
      
      box_lengthx_rec = 1.0_wp / box%length(1)
      box_lengthy_rec = 1.0_wp / box%length(2)
      k(:) = 0.0_wp

      ! Loop once through all modes to determine the number of effective modes
      num_modes_within_cutoff = 0

      select case (num_pbc)
      case (2)
         l = 0
         m = 1
         n = -ewald%kmax_z
      case (3)
         l = 0
         m = 0
         n = 1
      case default
         l = 0
         m = 0
         n = 0
      end select

      do imode = 1, num_modes_global
         ! k(1) = l * twopi / box%length(1)
         k(1) = (real(l,wp) * twopi) * box_lengthx_rec
         ! k(2) = m * twopi / box%length(2)
         k(2) = (real(m,wp) * twopi) * box_lengthy_rec
         ! k(3) = zpoint(n)
         k(3) = real(sign(1,n),wp) * ewald%zpoint(abs(n))

         ! Norm square of the k-mode vector
         knorm2 = k(1)*k(1) + k(2)*k(2) + k(3)*k(3)
         if (knorm2 <= ewald%knorm2_max) then
            num_modes_within_cutoff = num_modes_within_cutoff + 1
         end if
         ! Update k-mode parameters
         call MW_ewald_update_kmode_index(ewald, l, m, n)
      end do
      num_modes_global_within_cutoff = num_modes_within_cutoff

      ! Balance computational effort
      num_modes_local_within_cutoff = num_modes_global_within_cutoff / comm_size
      if (comm_rank < mod(num_modes_global_within_cutoff, comm_size)) then
         num_modes_local_within_cutoff = num_modes_local_within_cutoff + 1
      end if

      ! Determine the number of modes to skip before assigning them to this rank
      num_modes_global_before = comm_rank * (num_modes_global_within_cutoff / comm_size) &
            + min(comm_rank, mod(num_modes_global_within_cutoff, comm_size))

      ! Loop a second time to determine kstart_x, kstart_y, kstart_z
      select case (num_pbc)
      case (2)
         l = 0
         m = 1
         n = -ewald%kmax_z
      case (3)
         l = 0
         m = 0
         n = 1
      case default
         l = 0
         m = 0
         n = 0
      end select

      k(:) = 0.0_wp
      num_modes_within_cutoff = 0
      imode_start = 0
      do imode = 1, num_modes_global

         ! k(1) = l * twopi / box%length(1)
         k(1) = (real(l,wp) * twopi) * box_lengthx_rec
         ! k(2) = m * twopi / box%length(2)
         k(2) = (real(m,wp) * twopi) * box_lengthy_rec
         ! k(3) = zpoint(n)
         k(3) = real(sign(1,n),wp) * ewald%zpoint(abs(n))

         ! Norm square of the k-mode vector
         knorm2 = k(1)*k(1) + k(2)*k(2) + k(3)*k(3)
         if (knorm2 <= ewald%knorm2_max) then
            ! Store starting values for (l,m,n) when num_modes_global_before is reached
            if (num_modes_within_cutoff == num_modes_global_before) then
               if (imode_start == 0) then
                  imode_start = imode
                  kstart_x = l
                  kstart_y = m
                  kstart_z = n
               end if
            end if
            num_modes_within_cutoff = num_modes_within_cutoff + 1

            ! Stop searching when enough modes have been accounted for
            if (num_modes_within_cutoff == num_modes_global_before + num_modes_local_within_cutoff) then
               exit
            end if
         end if
         ! Update k-mode parameters
         call MW_ewald_update_kmode_index(ewald, l, m, n)
      end do
      num_modes_local = imode - imode_start + 1

      this%ewald_num_mode_global = num_modes_global
      this%ewald_num_mode_local = num_modes_local
      this%ewald_local_mode_offset = imode_start - 1
      this%ewald_kstart_x = kstart_x
      this%ewald_kstart_y = kstart_y
      this%ewald_kstart_z = kstart_z

   end subroutine setup_ewald_local_work

   !================================================================================
   ! Setup parallel distribution of work in atom 2 atom pair interactions
   !
   ! Interaction are symmetric and the work is divided into block
   subroutine setup_pair_atom2atom_local_work(this, electrodes, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_electrode_t), intent(in) :: electrodes(:)
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Local
      ! -----
      integer :: num_block_global, num_block_local
      integer :: iblock_offset
      integer :: pseudo_rank, rank_cut
      integer :: num_types, itype, jtype
      integer :: ierr
      integer :: num_block_rows, num_block_columns
      num_block_global = 0
      num_block_local = 0
      iblock_offset = 0
      rank_cut = 0

      num_types = size(electrodes,1)

      allocate( &
            this%pair_atom2atom_num_block_diag_local(num_types), &
            this%pair_atom2atom_diag_iblock_offset(num_types), &
            this%pair_atom2atom_num_block_full_local(num_types, num_types), &
            this%pair_atom2atom_full_iblock_offset(num_types, num_types), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_pair_atom2atom_local_work", "localwork.f90", ierr)
      end if
      this%pair_atom2atom_num_block_diag_local(:) = 0
      this%pair_atom2atom_diag_iblock_offset(:) = 0
      this%pair_atom2atom_num_block_full_local(:,:) = 0
      this%pair_atom2atom_full_iblock_offset(:,:) = 0

      ! Distribute diagonal blocks
      pseudo_rank = comm_rank
      do itype = 1, num_types
         if (electrodes(itype)%count > 0) then
            num_block_global = (electrodes(itype)%count - 1) / pair_block_size + 1
            num_block_local = num_block_global / comm_size
            rank_cut = mod(num_block_global, comm_size)
            if (pseudo_rank < rank_cut) then
               num_block_local = num_block_local + 1
            end if
            if (num_block_local > 0) then
               iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                     min(pseudo_rank, rank_cut)
            else
               iblock_offset = 0
            end if

            this%pair_atom2atom_num_block_diag_local(itype) = num_block_local
            this%pair_atom2atom_diag_iblock_offset(itype) = iblock_offset
            pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
         end if
      end do

      ! Distribute off-diagonal blocks
      ! Shift rank such that ranks with less diagonal block get more full block
      ! in priority
      !
      ! ...
      ! rank_cut-1 -> size-1
      ! rank_cut   -> 0
      ! rank_cut+1 -> 1
      ! ...

      do itype = 1, num_types
         if (electrodes(itype)%count > 0) then
            num_block_rows = (electrodes(itype)%count - 1) / pair_block_size + 1
            do jtype = 1, itype-1
               if (electrodes(jtype)%count > 0) then
                  num_block_columns = (electrodes(jtype)%count - 1) / pair_block_size + 1
                  num_block_global = num_block_rows * num_block_columns

                  num_block_local = num_block_global / comm_size
                  rank_cut = mod(num_block_global, comm_size)
                  if (pseudo_rank < rank_cut) then
                     num_block_local = num_block_local + 1
                  end if
                  if (num_block_local > 0) then
                     iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                           min(pseudo_rank, rank_cut)
                  else
                     iblock_offset = 0
                  end if

                  this%pair_atom2atom_num_block_full_local(itype,jtype) = num_block_local
                  this%pair_atom2atom_full_iblock_offset(itype,jtype) = iblock_offset
                  pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
               end if
            end do

            num_block_global = num_block_rows * (num_block_rows-1) / 2
            num_block_local = num_block_global / comm_size
            rank_cut = mod(num_block_global, comm_size)
            if (pseudo_rank < rank_cut) then
               num_block_local = num_block_local + 1
            end if
            if (num_block_local > 0) then
               iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                     min(pseudo_rank, rank_cut)
            else
               iblock_offset = 0
            end if

            this%pair_atom2atom_num_block_full_local(itype,itype) = num_block_local
            this%pair_atom2atom_full_iblock_offset(itype,itype) = iblock_offset
            pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
         end if
      end do
   end subroutine setup_pair_atom2atom_local_work

   !================================================================================
   ! Setup parallel distribution of work in atom 2 atom pair interactions
   !
   ! Interaction are symmetric and the work is divided into block
   subroutine setup_pair_ion2ion_local_work(this, ions, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_ion_t), intent(in) :: ions(:)
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Local
      ! -----
      integer :: num_block_global, num_block_local
      integer :: num_block_rows, num_block_columns
      integer :: iblock_offset
      integer :: pseudo_rank, rank_cut
      integer :: num_types, itype, jtype
      integer :: ierr
      num_block_global = 0
      num_block_local = 0
      iblock_offset = 0
      rank_cut = 0

      num_types = size(ions,1)

      allocate( &
            this%pair_ion2ion_num_block_diag_local(num_types), &
            this%pair_ion2ion_diag_iblock_offset(num_types), &
            this%pair_ion2ion_num_block_full_local(num_types, num_types), &
            this%pair_ion2ion_full_iblock_offset(num_types, num_types), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_pair_ion2ion_local_work", "localwork.f90", ierr)
      end if
      this%pair_ion2ion_num_block_diag_local(:) = 0
      this%pair_ion2ion_diag_iblock_offset(:) = 0
      this%pair_ion2ion_num_block_full_local(:,:) = 0
      this%pair_ion2ion_full_iblock_offset(:,:) = 0

      ! Distribute diagonal blocks
      pseudo_rank = comm_rank
      do itype = 1, num_types
         if (ions(itype)%count > 0) then
            num_block_global = (ions(itype)%count - 1) / pair_block_size + 1
            num_block_local = num_block_global / comm_size
            rank_cut = mod(num_block_global, comm_size)
            if (pseudo_rank < rank_cut) then
               num_block_local = num_block_local + 1
            end if
            if (num_block_local > 0) then
               iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                     min(pseudo_rank, rank_cut)
            else
               iblock_offset = 0
            end if

            this%pair_ion2ion_num_block_diag_local(itype) = num_block_local
            this%pair_ion2ion_diag_iblock_offset(itype) = iblock_offset
            pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
         end if
      end do

      ! Distribute off-diagonal blocks
      ! Shift rank such that ranks with less diagonal block get more full block
      ! in priority
      !
      ! ...
      ! rank_cut-1 -> size-1
      ! rank_cut   -> 0
      ! rank_cut+1 -> 1
      ! ...

      do itype = 1, num_types
         if (ions(itype)%count > 0) then
            num_block_rows = (ions(itype)%count - 1) / pair_block_size + 1
            do jtype = 1, itype-1
               if (ions(jtype)%count > 0) then
                  num_block_columns = (ions(jtype)%count - 1) / pair_block_size + 1
                  num_block_global = num_block_rows * num_block_columns
                  num_block_local = num_block_global / comm_size
                  rank_cut = mod(num_block_global, comm_size)
                  if (pseudo_rank < rank_cut) then
                     num_block_local = num_block_local + 1
                  end if
                  if (num_block_local > 0) then
                     iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                           min(pseudo_rank, rank_cut)
                  else
                     iblock_offset = 0
                  end if
                  this%pair_ion2ion_num_block_full_local(itype,jtype) = num_block_local
                  this%pair_ion2ion_full_iblock_offset(itype,jtype) = iblock_offset
                  pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
               end if
            end do

            ! Triangle of blocks
            num_block_global = num_block_rows * (num_block_rows-1) / 2
            num_block_local = num_block_global / comm_size
            rank_cut = mod(num_block_global, comm_size)
            if (pseudo_rank < rank_cut) then
               num_block_local = num_block_local + 1
            end if
            if (num_block_local > 0) then
               iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                     min(pseudo_rank, rank_cut)
            else
               iblock_offset = 0
            end if

            this%pair_ion2ion_num_block_full_local(itype,itype) = num_block_local
            this%pair_ion2ion_full_iblock_offset(itype,itype) = iblock_offset
            pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
         end if
      end do

   end subroutine setup_pair_ion2ion_local_work

   !================================================================================
   ! Setup parallel distribution of work in atom 2 atom pair interactions
   !
   ! Interaction are symmetric and the work is divided into block
   subroutine setup_pair_ion2ion_local_work_ms(this, ions, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_ion_t), intent(in) :: ions(:)
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Local
      ! -----
      integer :: num_block_global, num_block_local
      integer :: num_block_rows
      integer :: iblock_offset
      integer :: pseudo_rank, rank_cut
      integer :: num_types, itype, num_ions
      num_block_global = 0
      num_block_local = 0
      iblock_offset = 0
      rank_cut = 0

      num_ions = 0
      num_types = size(ions,1)
      do itype = 1, num_types
         num_ions = num_ions+ions(itype)%count
      enddo

      this%pair_ion2ion_num_block_diag_local_ms = 0
      this%pair_ion2ion_diag_iblock_offset_ms = 0
      this%pair_ion2ion_num_block_full_local_ms = 0
      this%pair_ion2ion_full_iblock_offset_ms = 0

      ! Distribute diagonal blocks
      pseudo_rank = comm_rank
      num_block_global = (num_ions - 1) / pair_block_size + 1
      num_block_local = num_block_global / comm_size
      rank_cut = mod(num_block_global, comm_size)
      if (pseudo_rank < rank_cut) then
         num_block_local = num_block_local + 1
      end if
      if (num_block_local > 0) then
         iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
               min(pseudo_rank, rank_cut)
      else
         iblock_offset = 0
      end if

      this%pair_ion2ion_num_block_diag_local_ms = num_block_local
      this%pair_ion2ion_diag_iblock_offset_ms = iblock_offset
      pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)

      ! Distribute off-diagonal blocks
      ! Shift rank such that ranks with less diagonal block get more full block
      ! in priority
      !
      ! ...
      ! rank_cut-1 -> size-1
      ! rank_cut   -> 0
      ! rank_cut+1 -> 1
      ! ...

      num_block_rows = (num_ions - 1) / pair_block_size + 1
      ! Triangle of blocks
      num_block_global = num_block_rows * (num_block_rows-1) / 2
      num_block_local = num_block_global / comm_size
      rank_cut = mod(num_block_global, comm_size)
      if (pseudo_rank < rank_cut) then
         num_block_local = num_block_local + 1
      end if
      if (num_block_local > 0) then
         iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
               min(pseudo_rank, rank_cut)
      else
         iblock_offset = 0
      end if

      this%pair_ion2ion_num_block_full_local_ms = num_block_local
      this%pair_ion2ion_full_iblock_offset_ms = iblock_offset
      pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)

   end subroutine setup_pair_ion2ion_local_work_ms

   !================================================================================
   ! Setup parallel distribution of work in ion 2 atom pair interactions
   !
   ! Interaction are symmetric and the work is divided into block
   subroutine setup_pair_ion2atom_local_work(this, electrodes, ions, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_electrode_t), intent(in) :: electrodes(:)
      type(MW_ion_t), intent(in) :: ions(:)
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Local
      ! -----
      integer :: num_block_global, num_block_local
      integer :: num_block_rows, num_block_columns
      integer :: iblock_offset
      integer :: pseudo_rank, rank_cut
      integer :: num_itypes, num_jtypes, itype, jtype
      integer :: ierr
      num_block_global = 0
      num_block_local = 0
      iblock_offset = 0
      rank_cut = 0

      num_itypes = size(electrodes,1)
      num_jtypes = size(ions,1)

      allocate( &
            this%pair_ion2atom_num_block_full_local(num_itypes, num_jtypes), &
            this%pair_ion2atom_full_iblock_offset(num_itypes, num_jtypes), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_pair_ion2atom_local_work", "localwork.f90", ierr)
      end if
      this%pair_ion2atom_num_block_full_local(:,:) = 0
      this%pair_ion2atom_full_iblock_offset(:,:) = 0

      ! Distribute off-diagonal blocks
      ! Shift rank such that ranks with less diagonal block get more full block
      ! in priority
      !
      ! ...
      ! rank_cut-1 -> size-1
      ! rank_cut   -> 0
      ! rank_cut+1 -> 1
      ! ...
      pseudo_rank = comm_rank
      do itype = 1, num_itypes
         if (electrodes(itype)%count > 0) then
            num_block_rows = (electrodes(itype)%count - 1) / pair_block_size + 1
            do jtype = 1, num_jtypes
               if (ions(jtype)%count > 0) then
                  num_block_columns = (ions(jtype)%count - 1) / pair_block_size + 1
                  num_block_global = num_block_rows * num_block_columns

                  num_block_local = num_block_global / comm_size
                  rank_cut = mod(num_block_global, comm_size)
                  if (pseudo_rank < rank_cut) then
                     num_block_local = num_block_local + 1
                  end if
                  if (num_block_local > 0) then
                     iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                           min(pseudo_rank, rank_cut)
                  else
                     iblock_offset = 0
                  end if

                  this%pair_ion2atom_num_block_full_local(itype,jtype) = num_block_local
                  this%pair_ion2atom_full_iblock_offset(itype,jtype) = iblock_offset
                  pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
               end if
            end do
         end if
      end do

   end subroutine setup_pair_ion2atom_local_work

   !================================================================================
   ! Setup parallel distribution of work in ion 2 atom pair interactions
   !
   ! Interaction are symmetric and the work is divided into block 
   ! MS modifications for GPUs
   subroutine setup_pair_ion2atom_local_work_ms(this, electrodes, ions, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_electrode_t), intent(in) :: electrodes(:)
      type(MW_ion_t), intent(in) :: ions(:)
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Local
      ! -----
      integer :: num_block_global, num_block_local
      integer :: num_block_rows, num_block_columns
      integer :: iblock_offset
      integer :: pseudo_rank, rank_cut
      integer :: num_itypes, num_jtypes, itype, jtype
      integer :: num_ions, num_elec_atoms
      num_block_global = 0
      num_block_local = 0
      iblock_offset = 0
      rank_cut = 0

      num_itypes = size(electrodes,1)
      num_jtypes = size(ions,1)

      num_elec_atoms = 0
      do itype = 1, num_itypes
         num_elec_atoms = num_elec_atoms + electrodes(itype)%count
      enddo

      num_ions = 0
      do jtype = 1, num_jtypes
         num_ions = num_ions + ions(jtype)%count
      enddo

      this%pair_ion2atom_num_block_full_local_ms = 0
      this%pair_ion2atom_full_iblock_offset_ms = 0

      ! Distribute off-diagonal blocks
      ! Shift rank such that ranks with less diagonal block get more full block
      ! in priority
      !
      ! ...
      ! rank_cut-1 -> size-1
      ! rank_cut   -> 0
      ! rank_cut+1 -> 1
      ! ...
      pseudo_rank = comm_rank
      num_block_rows = (num_elec_atoms - 1) / pair_block_size + 1
      num_block_columns = (num_ions - 1) / pair_block_size + 1
      num_block_global = num_block_rows * num_block_columns

      num_block_local = num_block_global / comm_size
      rank_cut = mod(num_block_global, comm_size)
      if (pseudo_rank < rank_cut) then
         num_block_local = num_block_local + 1
      end if
      if (num_block_local > 0) then
         iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
               min(pseudo_rank, rank_cut)
      else
         iblock_offset = 0
      end if

      this%pair_ion2atom_num_block_full_local_ms = num_block_local
      this%pair_ion2atom_full_iblock_offset_ms = iblock_offset
      pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)

   end subroutine setup_pair_ion2atom_local_work_ms
   !================================================================================
   ! Setup parallel distribution of work in ion 2 atom pair interactions
   !
   ! Interaction are symmetric and the work is divided into block
   subroutine setup_pair_atom2ion_local_work(this, ions, electrodes, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_ion_t), intent(in) :: ions(:)
      type(MW_electrode_t), intent(in) :: electrodes(:)
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Local
      ! -----
      integer :: num_block_global, num_block_local
      integer :: num_block_rows, num_block_columns
      integer :: iblock_offset
      integer :: pseudo_rank, rank_cut
      integer :: num_itypes, num_jtypes, itype, jtype
      integer :: ierr
      num_block_global = 0
      num_block_local = 0
      iblock_offset = 0
      rank_cut = 0

      num_itypes = size(ions,1)
      num_jtypes = size(electrodes,1)

      allocate( &
            this%pair_atom2ion_num_block_full_local(num_itypes, num_jtypes), &
            this%pair_atom2ion_full_iblock_offset(num_itypes, num_jtypes), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_pair_atom2ion_local_work", "localwork.f90", ierr)
      end if
      this%pair_atom2ion_num_block_full_local(:,:) = 0
      this%pair_atom2ion_full_iblock_offset(:,:) = 0

      ! Distribute off-diagonal blocks
      ! Shift rank such that ranks with less diagonal block get more full block
      ! in priority
      !
      ! ...
      ! rank_cut-1 -> size-1
      ! rank_cut   -> 0
      ! rank_cut+1 -> 1
      ! ...
      pseudo_rank = comm_rank
      do itype = 1, num_itypes
         if (ions(itype)%count> 0) then
            num_block_rows = (ions(itype)%count - 1) / pair_block_size + 1

            do jtype = 1, num_jtypes
               if (electrodes(jtype)%count > 0) then
                  num_block_columns = (electrodes(jtype)%count - 1) / pair_block_size + 1
                  num_block_global = num_block_rows * num_block_columns

                  num_block_local = num_block_global / comm_size
                  rank_cut = mod(num_block_global, comm_size)
                  if (pseudo_rank < rank_cut) then
                     num_block_local = num_block_local + 1
                  end if
                  if (num_block_local > 0) then
                     iblock_offset = pseudo_rank * (num_block_global / comm_size) + &
                           min(pseudo_rank, rank_cut)
                  else
                     iblock_offset = 0
                  end if
                  this%pair_atom2ion_num_block_full_local(itype,jtype) = num_block_local
                  this%pair_atom2ion_full_iblock_offset(itype,jtype) = iblock_offset
                  pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
               end if
            end do
         end if
      end do

   end subroutine setup_pair_atom2ion_local_work

   ! ================================================================================
   ! Setup parallel work on ions vector
   !
   subroutine setup_ions_local_work(this, ions, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameter
      ! ---------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_ion_t), intent(in) :: ions(:) !< global ions offset by ion_types
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Locals
      ! ------
      integer :: pseudo_rank, rank_cut
      integer :: num_global, num_local
      integer :: num_itype_global, num_itype_local
      integer :: num_ion_types, itype
      integer :: ierr

      num_ion_types = size(ions,1)

      allocate(this%offset_ions(num_ion_types), this%count_ions(num_ion_types), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_ions_local_work", "localwork.f90", ierr)
      end if
      this%offset_ions(:) = 0
      this%count_ions(:) = 0

      pseudo_rank = comm_rank
      num_global = 0
      num_local = 0
      do itype = 1, num_ion_types
         num_itype_global = ions(itype)%count
         num_itype_local = num_itype_global / comm_size
         rank_cut = mod(num_itype_global, comm_size)
         if (pseudo_rank < rank_cut) then
            num_itype_local = num_itype_local + 1
         end if
         this%count_ions(itype) = num_itype_local
         if (num_itype_local > 0) then
            this%offset_ions(itype) = pseudo_rank * (num_itype_global / comm_size) + &
                  min(pseudo_rank, rank_cut) + ions(itype)%offset
         end if
         num_global = num_global + num_itype_global
         num_local = num_local + num_itype_local
         pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
      end do
      this%ions_num_global = num_global
      this%ions_num_local = num_local

   end subroutine setup_ions_local_work

   ! ================================================================================
   ! Setup parallel work on atoms vector
   !
   subroutine setup_atoms_local_work(this, electrodes, comm_size, comm_rank)
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameter
      ! ---------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_electrode_t), intent(in) :: electrodes(:) !< global atoms offset by atom_types
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Locals
      ! ------
      integer :: pseudo_rank, rank_cut
      integer :: num_global, num_local, num_itype_global, num_itype_local
      integer :: num_atom_types, itype
      integer :: ierr

      num_atom_types = size(electrodes,1)

      allocate(this%offset_atoms(num_atom_types), this%count_atoms(num_atom_types), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_atoms_local_work", "localwork.f90", ierr)
      end if
      this%offset_atoms(:) = 0
      this%count_atoms(:) = 0

      pseudo_rank = comm_rank
      num_global = 0
      num_local = 0
      do itype = 1, num_atom_types
         num_itype_global = electrodes(itype)%count
         num_itype_local = num_itype_global / comm_size
         rank_cut = mod(num_itype_global, comm_size)
         if (pseudo_rank < rank_cut) then
            num_itype_local = num_itype_local + 1
         end if
         this%count_atoms(itype) = num_itype_local
         if (num_itype_local > 0) then
            this%offset_atoms(itype) = pseudo_rank * (num_itype_global / comm_size) + &
                  min(pseudo_rank, rank_cut) + electrodes(itype)%offset
         end if
         num_global = num_global + num_itype_global
         num_local = num_local + num_itype_local
         pseudo_rank = mod(pseudo_rank + comm_size - rank_cut, comm_size)
      end do
      this%atoms_num_global = num_global
      this%atoms_num_local = num_local

   end subroutine setup_atoms_local_work

   ! ================================================================================
   ! Setup parallel work on arrays of molecules
   !
   subroutine setup_molecules_local_work(this, molecules, comm_size, comm_rank)
      use MW_molecule, only: MW_molecule_t
      use MW_errors, only: MW_errors_allocate_error => allocate_error
      implicit none
      ! Parameter
      ! ---------
      type(MW_localwork_t), intent(inout) :: this
      type(MW_molecule_t), intent(in) :: molecules(:) !< Molecule definition
      integer, intent(in) :: comm_size
      integer, intent(in) :: comm_rank

      ! Locals
      ! ------
      integer :: pseudo_rank, rank_cut
      integer :: num_global, num_local, istart, iend
      integer :: num_molecule_types, num_molecules, itype
      integer :: ierr

      num_molecule_types = size(molecules,1)
      num_molecules = 0
      do itype = 1, num_molecule_types
         num_molecules = num_molecules + molecules(itype)%n
      end do

      allocate(this%molecules_istart(num_molecule_types), &
            this%molecules_iend(num_molecule_types), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("setup_molecules_local_work", "localwork.f90", ierr)
      end if
      this%molecules_istart(:) = 0
      this%molecules_iend(:) = 0


      pseudo_rank = comm_rank
      call setup_vector_local_work(comm_size, pseudo_rank, num_molecules, &
         rank_cut, num_global, num_local, istart)

      this%molecules_num_global = num_global
      this%molecules_num_local = num_local

      iend = istart + num_local - 1
      num_molecules = 0
      do itype = 1, num_molecule_types
         this%molecules_istart(itype) = max(1, istart - num_molecules)
         this%molecules_iend(itype) = min(molecules(itype)%n, iend - num_molecules)
         num_molecules = num_molecules + molecules(itype)%n
      end do

   end subroutine setup_molecules_local_work

   ! ================================================================================
   ! Setup parallel distribution of work on 1D vector
   !
   subroutine setup_vector_local_work(comm_size, comm_rank, n, &
         rank_cut, num_global, num_local, istart)
      implicit none
      ! Parameters
      ! ----------
      integer, intent(in) :: comm_size         !< parallel size
      integer, intent(in) :: comm_rank         !< parallel rank
      integer, intent(in) :: n            !< Global vector size

      integer, intent(inout) :: rank_cut  !< Index of cutoff rank (rank<rank_cut have 1 more block rank>=rank_cut)
      integer, intent(inout) :: num_global !< Global vector size
      integer, intent(inout) :: num_local  !< Local vector size
      integer, intent(inout) :: istart     !< Beginning of local vector in global vector

      num_global = n
      num_local = n / comm_size
      rank_cut = mod(num_global, comm_size)
      if (comm_rank < rank_cut) then
         num_local = num_local + 1
      end if

      istart = comm_rank*(num_global / comm_size) + min(comm_rank, rank_cut) + 1
   end subroutine setup_vector_local_work

   include 'update_diag_block_boundaries.inc'
   include 'update_tri_block_boundaries.inc'
   include 'update_other_block_boundaries.inc'

end module MW_localwork
