module MW_ions_dipoles
#ifndef MW_SERIAL
   use MPI
#endif
   use MW_parallel, only: MW_COMM_WORLD
   use MW_stdio
   use MW_kinds, only: wp
   use MW_algorithms, only: MW_algorithms_t
   use MW_cg, only: MW_cg_t
   use MW_system, only: MW_system_t
   use MW_ion, only: MW_ion_t
   use MW_external_field, only: MW_external_field_t, &
         FIELD_TYPE_ELECTRIC, FIELD_TYPE_DISPLACEMENT

   implicit none
   private

   ! Public subroutines
   ! ------------------
   public :: compute_jacobi_precond
   public :: compute
   public :: predictor

contains

   !================================================================================
   ! Evaluate b, the dipole due to point charges only
   subroutine setup_b(system, algorithms)
      use MW_coulomb, only: &
      MW_coulomb_qmelt2mumelt_electricfield => qmelt2mumelt_electricfield, &
      MW_coulomb_fix_molecule_qmelt2mumelt_electricfield => fix_molecule_qmelt2mumelt_electricfield
      use MW_external_field, only: &
      MW_external_field_field2mumelt_electricfieldb => field2mumelt_electricfieldb
      implicit none
      ! Parameters
      ! ----------
      type(MW_system_t), intent(inout) :: system
      type(MW_algorithms_t), intent(inout) :: algorithms

      ! Locals
      ! ------
      integer :: num_ions
      integer :: xyz_now, qatoms_now
      integer :: i, itype, ix, iy, iz
      real(wp) :: alpha
      real(wp), dimension(system%num_ions,3) :: efield, efield2, efield3


      num_ions = system%num_ions
      xyz_now = system%xyz_ions_step(1)
      qatoms_now = system%q_atoms_step(1)

      call MW_coulomb_qmelt2mumelt_electricfield(system%num_PBC, system%localwork, system%ewald, system%box, &
              system%ions, system%xyz_ions(:,:,xyz_now), system%type_ions(:), efield, system%tt)

      efield2=0.d0
      if(system%num_molecule_types > 0)then
         call MW_coulomb_fix_molecule_qmelt2mumelt_electricfield(system%num_PBC, system%localwork, &
         system%molecules, system%box, system%ions, system%xyz_ions(:,:,xyz_now), efield2, system%tt)
      endif

      efield3=0.d0
      if (system%field%field_type > 0) then
         call MW_external_field_field2mumelt_electricfieldb(system%field, system%ions,&
         system%polarization_field, efield3)
      endif

      efield=efield+efield2+efield3

      do itype = 1,size(system%ions,1)
         alpha = system%ions(itype)%polarizability
         if (alpha .ne. 0.0_wp) then
            if(algorithms%maze%calc_dip) then
               do i=system%ions(itype)%offset+1,system%ions(itype)%offset+system%ions(itype)%count
                  ix = i
                  iy = i + num_ions
                  iz = i + num_ions*2
                  algorithms%maze%const_constr(ix) = -efield(i,1)
                  algorithms%maze%const_constr(iy) = -efield(i,2)
                  algorithms%maze%const_constr(iz) = -efield(i,3)
               end do
            else if (algorithms%matrix_inversion%calc_dip) then
               do i=system%ions(itype)%offset+1,system%ions(itype)%offset+system%ions(itype)%count
                  ix = i
                  iy = i + num_ions
                  iz = i + num_ions*2
                  algorithms%matrix_inversion%b(ix) = efield(i,1)
                  algorithms%matrix_inversion%b(iy) = efield(i,2)
                  algorithms%matrix_inversion%b(iz) = efield(i,3)
               end do
            else
               do i=system%ions(itype)%offset+1,system%ions(itype)%offset+system%ions(itype)%count
                  ix = i
                  iy = i + num_ions
                  iz = i + num_ions*2
                  algorithms%cg%b(ix) = efield(i,1)
                  algorithms%cg%b(iy) = efield(i,2)
                  algorithms%cg%b(iz) = efield(i,3)
               end do
            end if
         end if
      end do

   end subroutine setup_b

   !================================================================================
   ! Compute A*x
   !
   subroutine apply_A(system, x, y)
      use MW_coulomb, only: &
      MW_coulomb_mumelt2mumelt_electricfield => mumelt2mumelt_electricfield, &
      MW_coulomb_fix_molecule_mumelt2mumelt_electricfield => fix_molecule_mumelt2mumelt_electricfield
      use MW_external_field, only: &
      MW_external_field_field2mumelt_electricfieldA => field2mumelt_electricfieldA
      implicit none
      ! Parameters
      ! ----------
      type(MW_system_t), intent(inout) :: system
      real(wp), intent(in) :: x(:)
      real(wp), intent(out) :: y(:)

      ! Locals
      ! ------
      integer :: num_ions
      integer :: i,itype,ix,iy,iz
      integer :: q_atoms_now,xyz_now
      real(wp) :: alpha
      real(wp), dimension(system%num_ions,3) :: efield,efield2,efield3
      real(wp), dimension(system%num_ions,3) :: dipoles

      num_ions = system%num_ions

      do i=1,num_ions
         ix = i
         iy = i + num_ions
         iz = i + num_ions*2
         dipoles(i,1) = x(ix)
         dipoles(i,2) = x(iy)
         dipoles(i,3) = x(iz)
      enddo

      q_atoms_now = system%q_atoms_step(1)
      xyz_now = system%xyz_ions_step(1)


      call MW_coulomb_mumelt2mumelt_electricfield(system%num_PBC, system%localwork, system%ewald, system%box, &
      system%ions, system%xyz_ions(:,:,xyz_now), efield, dipoles)

      efield2=0.0_wp
      if(system%num_molecule_types > 0) then
         call MW_coulomb_fix_molecule_mumelt2mumelt_electricfield(system%num_PBC, system%localwork, &
         system%molecules, system%box, system%ions, system%xyz_ions(:,:,xyz_now), efield2, dipoles)
      endif

      efield3=0.0_wp
      if (system%field%field_type == FIELD_TYPE_DISPLACEMENT) then
         call MW_external_field_field2mumelt_electricfieldA(system%box, system%field, system%electrodes, &
         system%xyz_atoms, system%q_atoms(:,q_atoms_now), system%ions, &
         system%dipoles_ions(:,:,1), system%polarization_field, efield3)
         !Are we sure this is system%dipoles and not dipoles?
         !We are minimizing over dipoles and they change at each iteration and we are not updating the system variable.
      endif

      efield=efield+efield2+efield3
      do itype = 1,size(system%ions,1)
         alpha = system%ions(itype)%polarizability
         if (alpha .ne. 0.0_wp) then
            do i=system%ions(itype)%offset+1,system%ions(itype)%offset+system%ions(itype)%count
               ix = i
               iy = i + num_ions
               iz = i + num_ions*2
               y(ix) = -efield(i,1) + dipoles(i,1)/alpha
               y(iy) = -efield(i,2) + dipoles(i,2)/alpha
               y(iz) = -efield(i,3) + dipoles(i,3)/alpha
            enddo
         end if
      end do

   end subroutine apply_A

   subroutine compute_hessian(system, hessian)
      use MW_coulomb, only: &
         MW_coulomb_gradmumelt_electricfield => gradmumelt_electricfield, &
         MW_coulomb_fix_molecule_gradmumelt_electricfield => fix_molecule_gradmumelt_electricfield
      use MW_errors, only: &
         MW_errors_allocate_error => allocate_error
      implicit none

      ! Parameters Input/Output
      ! -----------------------
      type(MW_system_t),  intent(inout) :: system
      real(wp), intent(out) :: hessian(:,:)

      !Locals
      !------
      integer  :: num_ions
      integer  :: xyz_now
      integer  :: i, itype, j, ix, iy, iz, jx, jy, jz
      real(wp) :: alpha
      real(wp), allocatable :: fix_mol_hessian(:,:)
      integer :: ierr

      allocate(fix_mol_hessian(3*system%num_ions,3*system%num_ions), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("compute_hessian", "ions_dipoles.F90", ierr)
      end if

      num_ions = system%num_ions
      xyz_now = system%xyz_ions_step(1)

      call MW_coulomb_gradmumelt_electricfield(system%num_PBC, system%localwork, system%ewald, &
      system%box, system%ions, system%xyz_ions(:,:,xyz_now), hessian(:,:))
      fix_mol_hessian(:,:) = 0.0_wp
      if(system%num_molecule_types > 0) then
         call MW_coulomb_fix_molecule_gradmumelt_electricfield(system%num_PBC, system%localwork, &
         system%molecules, system%box, &
         system%ions, system%xyz_ions(:,:,xyz_now), fix_mol_hessian(:,:))
      endif
      hessian(:,:) = hessian(:,:) + fix_mol_hessian(:,:)
      ! Computing polarization energy term in gradient of constraints for dipoles
      ! Not included in MW_coulomb_melt_gradmu_electricfield
      do itype = 1, size(system%ions,1)
         alpha = system%ions(itype)%polarizability
         if (alpha .ne. 0.0_wp) then
            do i = system%ions(itype)%offset+1, system%ions(itype)%offset+system%ions(itype)%count
               ix = i
               iy = i + num_ions
               iz = i + num_ions*2
               hessian(ix,ix) = hessian(ix,ix) + 1.0_wp/alpha
               hessian(iy,iy) = hessian(iy,iy) + 1.0_wp/alpha
               hessian(iz,iz) = hessian(iz,iz) + 1.0_wp/alpha
            end do
         end if
      end do

      do i = 1, num_ions
         ix = i
         iy = i + num_ions
         iz = i + num_ions*2
         do j = 1, i
            jx = j
            jy = j + num_ions
            jz = j + num_ions*2
            !xx
            hessian(ix,jx) = hessian(jx,ix)
            !xy
            hessian(iy,jx) = hessian(jy,ix)
            hessian(jx,iy) = hessian(jy,ix)
            hessian(ix,jy) = hessian(jx,iy)
            !xz
            hessian(iz,jx) = hessian(jz,ix)
            hessian(jx,iz) = hessian(jz,ix)
            hessian(ix,jz) = hessian(jx,iz)
            !yy
            hessian(iy,jy) = hessian(jy,iy)
            !yz
            hessian(iz,jy) = hessian(jz,iy)
            hessian(jy,iz) = hessian(jz,iy)
            hessian(iy,jz) = hessian(jy,iz)
            !zz
            hessian(iz,jz) = hessian(jz,iz)
         end do
      end do

   end subroutine compute_hessian

   ! Parallelized routine for computing matrix vector products invloving dipoles
   ! The matrix has to be Symmetric
   subroutine SMatrix_vector_product_dip(system, matrix, vector, matvec_product)
      implicit none

      ! Parameters Input
      ! ----------------
      type(MW_system_t), intent(in) :: system  !< system parameters
      real(wp),          intent(in) :: vector(:)
      real(wp),          intent(in) :: matrix(:,:)

      ! Parameters Output
      ! -----------------
      real(wp),          intent(out) :: matvec_product(:)

      !Locals
      !---------------
      integer  :: num_ion_types, num_ions
      integer  :: iblock, iblock_offset, istart, iend, jstart, jend, itype, jtype
      integer  :: num_block_diag, num_block_full
      integer  :: count_n, offset_n, count_m, offset_m, count, ierr
      integer  :: i, j, ix, iy, iz, jx, jy, jz
      real(wp) :: matvec_product_ix, matvec_product_iy, matvec_product_iz
      real(wp) :: vector_ix, vector_iy, vector_iz

      num_ion_types = size(system%ions,1)
      num_ions = size(matrix,1)/3

      matvec_product(:) = 0.0_wp
      do itype = 1, num_ion_types
         count_n = system%ions(itype)%count
         offset_n = system%ions(itype)%offset
         do jtype = 1, itype - 1
            count_m = system%ions(jtype)%count
            offset_m = system%ions(jtype)%offset

            ! Number of blocks with full ion2ion interactions
            num_block_full = system%localwork%pair_ion2ion_num_block_full_local(itype, jtype)
            iblock_offset = system%localwork%pair_ion2ion_full_iblock_offset(itype, jtype)
            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
               offset_n, count_n, offset_m, count_m, &
               istart, iend, jstart, jend)
               do i = istart, iend
                  ix = i
                  iy = i + num_ions
                  iz = i + num_ions*2
                  matvec_product_ix = 0.0_wp
                  matvec_product_iy = 0.0_wp
                  matvec_product_iz = 0.0_wp
                  vector_ix = vector(ix)
                  vector_iy = vector(iy)
                  vector_iz = vector(iz)
                  do j = jstart, jend
                     jx = j
                     jy = j + num_ions
                     jz = j + num_ions*2
                     matvec_product_ix = matvec_product_ix &
                     + matrix(jx,ix)*vector(jx) &
                     + matrix(jx,iy)*vector(jy) &
                     + matrix(jx,iz)*vector(jz)
                     matvec_product_iy = matvec_product_iy &
                     + matrix(jy,ix)*vector(jx) &
                     + matrix(jy,iy)*vector(jy) &
                     + matrix(jy,iz)*vector(jz)
                     matvec_product_iz = matvec_product_iz &
                     + matrix(jz,ix)*vector(jx) &
                     + matrix(jz,iy)*vector(jy) &
                     + matrix(jz,iz)*vector(jz)

                     matvec_product(jx) = matvec_product(jx) &
                     + matrix(jx,ix)*vector_ix &
                     + matrix(jx,iy)*vector_iy &
                     + matrix(jx,iz)*vector_iz
                     matvec_product(jy) = matvec_product(jy) &
                     + matrix(jy,ix)*vector_ix &
                     + matrix(jy,iy)*vector_iy &
                     + matrix(jy,iz)*vector_iz
                     matvec_product(jz) = matvec_product(jz) &
                     + matrix(jz,ix)*vector_ix &
                     + matrix(jz,iy)*vector_iy &
                     + matrix(jz,iz)*vector_iz
                  end do
                  matvec_product(ix) = matvec_product(ix) + matvec_product_ix
                  matvec_product(iy) = matvec_product(iy) + matvec_product_iy
                  matvec_product(iz) = matvec_product(iz) + matvec_product_iz
               end do
            end do
         end do

         ! Blocks on the diagonal compute only half of the pair interactions
         num_block_diag = system%localwork%pair_ion2ion_num_block_diag_local(itype)
         iblock_offset = system%localwork%pair_ion2ion_diag_iblock_offset(itype)
         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock_offset+iblock, &
            offset_n, count_n, istart, iend)
            do i = istart, iend
               ix = i
               iy = i + num_ions
               iz = i + num_ions*2
               matvec_product_ix = 0.0_wp
               matvec_product_iy = 0.0_wp
               matvec_product_iz = 0.0_wp
               do j = istart, iend
                  jx = j
                  jy = j + num_ions
                  jz = j + num_ions*2
                  matvec_product_ix = matvec_product_ix &
                  + matrix(jx,ix)*vector(jx) &
                  + matrix(jx,iy)*vector(jy) &
                  + matrix(jx,iz)*vector(jz)
                  matvec_product_iy = matvec_product_iy &
                  + matrix(jy,ix)*vector(jx) &
                  + matrix(jy,iy)*vector(jy) &
                  + matrix(jy,iz)*vector(jz)
                  matvec_product_iz = matvec_product_iz &
                  + matrix(jz,ix)*vector(jx) &
                  + matrix(jz,iy)*vector(jy) &
                  + matrix(jz,iz)*vector(jz)
               end do
               matvec_product(ix) = matvec_product(ix) + matvec_product_ix
               matvec_product(iy) = matvec_product(iy) + matvec_product_iy
               matvec_product(iz) = matvec_product(iz) + matvec_product_iz
            end do
         end do

         ! Number of blocks with full ion2ion interactions
         num_block_full = system%localwork%pair_ion2ion_num_block_full_local(itype,itype)
         iblock_offset = system%localwork%pair_ion2ion_full_iblock_offset(itype,itype)
         do iblock = 1, num_block_full
            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
            istart, iend, jstart, jend)
            do i = istart, iend
               ix = i
               iy = i + num_ions
               iz = i + num_ions*2
               matvec_product_ix = 0.0_wp
               matvec_product_iy = 0.0_wp
               matvec_product_iz = 0.0_wp
               vector_ix = vector(ix)
               vector_iy = vector(iy)
               vector_iz = vector(iz)
               do j = jstart, jend
                  jx = j
                  jy = j + num_ions
                  jz = j + num_ions*2
                  matvec_product_ix = matvec_product_ix &
                  + matrix(jx,ix)*vector(jx) &
                  + matrix(jx,iy)*vector(jy) &
                  + matrix(jx,iz)*vector(jz)
                  matvec_product_iy = matvec_product_iy &
                  + matrix(jy,ix)*vector(jx) &
                  + matrix(jy,iy)*vector(jy) &
                  + matrix(jy,iz)*vector(jz)
                  matvec_product_iz = matvec_product_iz &
                  + matrix(jz,ix)*vector(jx) &
                  + matrix(jz,iy)*vector(jy) &
                  + matrix(jz,iz)*vector(jz)

                  matvec_product(jx) = matvec_product(jx) &
                  + matrix(jx,ix)*vector_ix &
                  + matrix(jx,iy)*vector_iy &
                  + matrix(jx,iz)*vector_iz
                  matvec_product(jy) = matvec_product(jy) &
                  + matrix(jy,ix)*vector_ix &
                  + matrix(jy,iy)*vector_iy &
                  + matrix(jy,iz)*vector_iz
                  matvec_product(jz) = matvec_product(jz) &
                  + matrix(jz,ix)*vector_ix &
                  + matrix(jz,iy)*vector_iy &
                  + matrix(jz,iz)*vector_iz
               end do
               matvec_product(ix) = matvec_product(ix) + matvec_product_ix
               matvec_product(iy) = matvec_product(iy) + matvec_product_iy
               matvec_product(iz) = matvec_product(iz) + matvec_product_iz
            end do
         end do
      end do

#ifndef MW_SERIAL
      count = size(matvec_product,1)
      call MPI_Allreduce(MPI_IN_PLACE, matvec_product, count, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD,  ierr)
#endif
   end subroutine SMatrix_vector_product_dip

   ! Parallelized routine for computing matrix vector products invloving dipoles
   ! The matrix is General
   subroutine GMatrix_vector_product_dip(system, matrix, vector, matvec_product)
      implicit none

      ! Parameters Input
      ! ----------------
      type(MW_system_t), intent(in) :: system  !< system parameters
      real(wp),          intent(in) :: vector(:)
      real(wp),          intent(in) :: matrix(:,:)

      ! Parameters Output
      ! -----------------
      real(wp),          intent(out) :: matvec_product(:)

      !Locals
      !---------------
      integer  :: num_ion_types, num_ions
      integer  :: iblock, iblock_offset, istart, iend, jstart, jend, itype, jtype
      integer  :: num_block_diag, num_block_full
      integer  :: count_n, offset_n, count_m, offset_m, count, ierr
      integer  :: i, j, ix, iy, iz, jx, jy, jz
      real(wp) :: matvec_product_ix, matvec_product_iy, matvec_product_iz
      real(wp) :: vector_ix, vector_iy, vector_iz

      num_ion_types = size(system%ions,1)
      num_ions = size(matrix,1)/3

      matvec_product(:) = 0.0_wp
      do itype = 1, num_ion_types
         count_n = system%ions(itype)%count
         offset_n = system%ions(itype)%offset
         do jtype = 1, itype - 1
            count_m = system%ions(jtype)%count
            offset_m = system%ions(jtype)%offset

            ! Number of blocks with full ion2ion interactions
            num_block_full = system%localwork%pair_ion2ion_num_block_full_local(itype, jtype)
            iblock_offset = system%localwork%pair_ion2ion_full_iblock_offset(itype, jtype)
            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
               offset_n, count_n, offset_m, count_m, &
               istart, iend, jstart, jend)
               do i = istart, iend
                  ix = i
                  iy = i + num_ions
                  iz = i + num_ions*2
                  matvec_product_ix = 0.0_wp
                  matvec_product_iy = 0.0_wp
                  matvec_product_iz = 0.0_wp
                  vector_ix = vector(ix)
                  vector_iy = vector(iy)
                  vector_iz = vector(iz)
                  do j = jstart, jend
                     jx = j
                     jy = j + num_ions
                     jz = j + num_ions*2
                     matvec_product_ix = matvec_product_ix &
                     + matrix(ix,jx)*vector(jx) &
                     + matrix(ix,jy)*vector(jy) &
                     + matrix(ix,jz)*vector(jz)
                     matvec_product_iy = matvec_product_iy &
                     + matrix(iy,jx)*vector(jx) &
                     + matrix(iy,jy)*vector(jy) &
                     + matrix(iy,jz)*vector(jz)
                     matvec_product_iz = matvec_product_iz &
                     + matrix(iz,jx)*vector(jx) &
                     + matrix(iz,jy)*vector(jy) &
                     + matrix(iz,jz)*vector(jz)

                     matvec_product(jx) = matvec_product(jx) &
                     + matrix(jx,ix)*vector_ix &
                     + matrix(jx,iy)*vector_iy &
                     + matrix(jx,iz)*vector_iz
                     matvec_product(jy) = matvec_product(jy) &
                     + matrix(jy,ix)*vector_ix &
                     + matrix(jy,iy)*vector_iy &
                     + matrix(jy,iz)*vector_iz
                     matvec_product(jz) = matvec_product(2*num_ions+j) &
                     + matrix(jz,ix)*vector_ix &
                     + matrix(jz,iy)*vector_iy &
                     + matrix(jz,iz)*vector_iz
                  end do
                  matvec_product(ix) = matvec_product(ix) + matvec_product_ix
                  matvec_product(iy) = matvec_product(iy) + matvec_product_iy
                  matvec_product(iz) = matvec_product(iz) + matvec_product_iz
               end do
            end do
         end do

         ! Blocks on the diagonal compute only half of the pair interactions
         num_block_diag = system%localwork%pair_ion2ion_num_block_diag_local(itype)
         iblock_offset = system%localwork%pair_ion2ion_diag_iblock_offset(itype)
         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock_offset+iblock, &
            offset_n, count_n, istart, iend)
            do i = istart, iend
               ix = i
               iy = i + num_ions
               iz = i + num_ions*2
               matvec_product_ix = 0.0_wp
               matvec_product_iy = 0.0_wp
               matvec_product_iz = 0.0_wp
               do j = istart, iend
                  jx = j
                  jy = j + num_ions
                  jz = j + num_ions*2
                  matvec_product_ix = matvec_product_ix &
                  + matrix(ix,jx)*vector(jx) &
                  + matrix(ix,jy)*vector(jy) &
                  + matrix(ix,jz)*vector(jz)
                  matvec_product_iy = matvec_product_iy &
                  + matrix(iy,jx)*vector(jx) &
                  + matrix(iy,jy)*vector(jy) &
                  + matrix(iy,jz)*vector(jz)
                  matvec_product_iz = matvec_product_iz &
                  + matrix(iz,jx)*vector(jx) &
                  + matrix(iz,jy)*vector(jy) &
                  + matrix(iz,jz)*vector(jz)
               end do
               matvec_product(ix) = matvec_product(ix) + matvec_product_ix
               matvec_product(iy) = matvec_product(iy) + matvec_product_iy
               matvec_product(iz) = matvec_product(iz) + matvec_product_iz
            end do
         end do

         ! Number of blocks with full ion2ion interactions
         num_block_full = system%localwork%pair_ion2ion_num_block_full_local(itype,itype)
         iblock_offset = system%localwork%pair_ion2ion_full_iblock_offset(itype,itype)
         do iblock = 1, num_block_full
            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
            istart, iend, jstart, jend)
            do i = istart, iend
               ix = i
               iy = i + num_ions
               iz = i + num_ions*2
               matvec_product_ix = 0.0_wp
               matvec_product_iy = 0.0_wp
               matvec_product_iz = 0.0_wp
               vector_ix = vector(ix)
               vector_iy = vector(iy)
               vector_iz = vector(iz)
               do j = jstart, jend
                  jx = j
                  jy = j + num_ions
                  jz = j + num_ions*2
                  matvec_product_ix = matvec_product_ix &
                  + matrix(ix,jx)*vector(jx) &
                  + matrix(ix,jy)*vector(jy) &
                  + matrix(ix,jz)*vector(jz)
                  matvec_product_iy = matvec_product_iy &
                  + matrix(iy,jx)*vector(jx) &
                  + matrix(iy,jy)*vector(jy) &
                  + matrix(iy,jz)*vector(jz)
                  matvec_product_iz = matvec_product_iz &
                  + matrix(iz,jx)*vector(jx) &
                  + matrix(iz,jy)*vector(jy) &
                  + matrix(iz,jz)*vector(jz)

                  matvec_product(jx) = matvec_product(jx) &
                  + matrix(jx,ix)*vector_ix &
                  + matrix(jx,iy)*vector_iy &
                  + matrix(jx,iz)*vector_iz
                  matvec_product(jy) = matvec_product(jy) &
                  + matrix(jy,ix)*vector_ix &
                  + matrix(jy,iy)*vector_iy &
                  + matrix(jy,iz)*vector_iz
                  matvec_product(jz) = matvec_product(2*num_ions+j) &
                  + matrix(jz,ix)*vector_ix &
                  + matrix(jz,iy)*vector_iy &
                  + matrix(jz,iz)*vector_iz
               end do
               matvec_product(ix) = matvec_product(ix) + matvec_product_ix
               matvec_product(iy) = matvec_product(iy) + matvec_product_iy
               matvec_product(iz) = matvec_product(iz) + matvec_product_iz
            end do
         end do
      end do

#ifndef MW_SERIAL
      count = size(matvec_product,1)
      call MPI_Allreduce(MPI_IN_PLACE, matvec_product, count, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD,  ierr)
#endif
   end subroutine GMatrix_vector_product_dip

   subroutine compute_jacobi_precond(system, algorithms)
      use MW_coulomb, only: &
      MW_coulomb_diag_gradmumelt_electricfield => diag_gradmumelt_electricfield, &
      MW_coulomb_diag_gradQelec_potential => diag_gradQelec_potential
      implicit none

      type(MW_system_t), intent(inout) :: system
      type(MW_algorithms_t), intent(inout) :: algorithms

      integer :: i, itype, ix, iy, iz
      integer :: n_ions
      real(wp) :: alpha

      n_ions = system%num_ions

      call MW_coulomb_diag_gradmumelt_electricfield(system%num_PBC, system%localwork, system%ewald, system%box, &
         system%ions, algorithms%cg%jacobi_precond(1:3*n_ions))
      do itype = 1,size(system%ions,1)
         alpha = system%ions(itype)%polarizability
         if (alpha .ne. 0.0_wp) then
            do i=system%ions(itype)%offset+1,system%ions(itype)%offset+system%ions(itype)%count
               ix = i
               iy = i + n_ions
               iz = i + n_ions*2
               algorithms%cg%jacobi_precond(ix) = algorithms%cg%jacobi_precond(ix) + 1.0_wp/system%ions(itype)%polarizability
               algorithms%cg%jacobi_precond(iy) = algorithms%cg%jacobi_precond(iy) + 1.0_wp/system%ions(itype)%polarizability
               algorithms%cg%jacobi_precond(iz) = algorithms%cg%jacobi_precond(iz) + 1.0_wp/system%ions(itype)%polarizability
            end do
         end if
      end do
   end subroutine compute_jacobi_precond

   !================================================================================
   !> Compute dipoles on electrode
   !!
   !! The dipole is evaluated by minimizing the Coulomb energy of the system
   !!
   !! The energy can be cast into E = Q*AQ - Q*b + c
   !!
   !! where Q  is the vector of charges on the electrodes
   !!       b  is the potential felt by the electrode due to the melt ions
   !!       AQ is the potential felt by the electrode du to other electrode atoms
   !!       c  is the energy due to melt ions interactions
   !!       *  denotes the transpose of a vector/matrix
   !!
   !! By construction A is a matrix symmetric positive definite
   !! The preconditioned conjugate gradient algorithm is used to find a solution to
   !! the system
   !!
   subroutine compute(system, algorithms)
      use MW_cg, only: &
         MW_cg_solve_preconditioned => solve_preconditioned, &
         MW_cg_solve => solve
      use MW_maze, only: &
            MW_maze_update_history => update_history, &
            MW_maze_compute_x0 => compute_x0, &
            MW_maze_solve_invert => solve_invert, &
            MW_maze_solve_shake => solve_shake, &
            MW_maze_solve_block_iterate => solve_block_iterate
      use MW_matrix_inversion, only: &
            MW_matrix_inversion_solve => solve
      use MW_ion, only: MW_ion_t
      use MW_external_field, only: &
            MW_external_field_dipole_polarization => dipole_polarization
      implicit none

      ! Parameters in
      ! -------------
      type(MW_system_t), intent(inout) :: system !< system parameters
      type(MW_algorithms_t), intent(inout) :: algorithms !< cg algorithm parameters

      ! Local
      ! -----
      integer :: num_ions
      integer :: i, ix, iy, iz, itype
      integer :: dipoles_ions_tpdt,dipoles_ions_t,dipoles_ions_tmdt
      real(wp) :: alpha
      real(wp), dimension(3*system%num_ions) :: x
      real(wp), dimension(3*system%num_ions) :: xt, xtmdt


      num_ions = system%num_ions
      dipoles_ions_tpdt = system%dipoles_ions_step(1)
      dipoles_ions_t = system%dipoles_ions_step(2)
      dipoles_ions_tmdt = 0
      if(algorithms%maze%calc_dip) dipoles_ions_tmdt = system%dipoles_ions_step(3)

      call setup_b(system, algorithms)

      if (algorithms%maze%calc_dip) then

         do i=1,num_ions
            ix = i
            iy = i + num_ions
            iz = i + num_ions*2
            x(ix)=system%dipoles_ions(i,1,dipoles_ions_tpdt)
            x(iy)=system%dipoles_ions(i,2,dipoles_ions_tpdt)
            x(iz)=system%dipoles_ions(i,3,dipoles_ions_tpdt)
            xt(ix)=system%dipoles_ions(i,1,dipoles_ions_t)
            xt(iy)=system%dipoles_ions(i,2,dipoles_ions_t)
            xt(iz)=system%dipoles_ions(i,3,dipoles_ions_t)
            xtmdt(ix)=system%dipoles_ions(i,1,dipoles_ions_tmdt)
            xtmdt(iy)=system%dipoles_ions(i,2,dipoles_ions_tmdt)
            xtmdt(iz)=system%dipoles_ions(i,3,dipoles_ions_tmdt)
         end do

         if (.not.algorithms%maze%constant_constr_grad) then
           call MW_maze_update_history(algorithms%maze) !Only updates gradient of constraints
         end if
         call MW_maze_compute_x0(algorithms%maze, system, 1.0_wp, xtmdt, xt, x)
         if (algorithms%maze%invert) then
            call MW_maze_solve_invert(algorithms%maze, system, compute_hessian, GMatrix_vector_product_dip, x)
         else if (algorithms%maze%shake) then
            call MW_maze_solve_shake(algorithms%maze, system, compute_hessian, x)
         else
            call MW_maze_solve_block_iterate(algorithms%maze, system, compute_hessian, x)
         end if

      else if (algorithms%matrix_inversion%calc_dip) then

         do i=1,num_ions
            ix = i
            iy = i + num_ions
            iz = i + num_ions*2
            x(ix)=system%dipoles_ions(i,1,dipoles_ions_tpdt)
            x(iy)=system%dipoles_ions(i,2,dipoles_ions_tpdt)
            x(iz)=system%dipoles_ions(i,3,dipoles_ions_tpdt)
         end do

         call MW_matrix_inversion_solve(algorithms%matrix_inversion, system, compute_hessian, GMatrix_vector_product_dip, x)

      else

         !When using cg the iterations start from the dipoles at previous timestep
         if (system%dipoles_use_predictor) then
            call predictor(system%dipoles_ions_step, system%ions, system%dipoles_ions)
         else
            do itype = 1, system%num_ion_types
               alpha = system%ions(itype)%polarizability
               if (alpha .ne. 0.0_wp) then
                  do i = system%ions(itype)%offset+1, system%ions(itype)%offset + system%ions(itype)%count
                    system%dipoles_ions(i, 1, dipoles_ions_tpdt) = system%dipoles_ions(i, 1, dipoles_ions_t)
                    system%dipoles_ions(i, 2, dipoles_ions_tpdt) = system%dipoles_ions(i, 2, dipoles_ions_t)
                    system%dipoles_ions(i, 3, dipoles_ions_tpdt) = system%dipoles_ions(i, 3, dipoles_ions_t)
                  end do
               end if
            end do
         end if

         do i=1,num_ions
            ix = i
            iy = i + num_ions
            iz = i + num_ions*2
            x(ix)=system%dipoles_ions(i,1,dipoles_ions_tpdt)
            x(iy)=system%dipoles_ions(i,2,dipoles_ions_tpdt)
            x(iz)=system%dipoles_ions(i,3,dipoles_ions_tpdt)
         end do

         call MW_cg_solve_preconditioned(algorithms%cg, system, apply_A, x)
      endif

      !Update ions at t+dt after minimization
      do i=1,num_ions
         ix = i
         iy = i + num_ions
         iz = i + num_ions*2
         system%dipoles_ions(i,1,dipoles_ions_tpdt)=x(ix)
         system%dipoles_ions(i,2,dipoles_ions_tpdt)=x(iy)
         system%dipoles_ions(i,3,dipoles_ions_tpdt)=x(iz)
      end do

      call MW_external_field_dipole_polarization(system%box, system%ions, &
         system%dipoles_ions(:,:,dipoles_ions_tpdt), system%polarization_field(:,:))
   end subroutine compute

   ! ================================================================================
   !> Predictor for ion dipoles
   subroutine predictor(dipoles_ions_step, ions, dipoles_ions)
     implicit none
     ! Parameters
     ! ----------
     integer, intent(in) :: dipoles_ions_step(:)  !< time step index
     type(MW_ion_t), intent(in) :: ions(:) !< ion parameters
     real(wp), intent(inout) :: dipoles_ions(:,:,:) !< atom charges

     ! Local
     ! -----
     integer :: dipole_now, dipole_m1, dipole_m2, dipole_m3, dipole_m4, dipole_m5, dipole_m6
     integer :: num_ion_types
     integer :: itype, i
     real(wp) :: alpha

     real(wp), parameter :: B1 =  22.0_wp /  7.0_wp
     real(wp), parameter :: B2 = -55.0_wp / 14.0_wp
     real(wp), parameter :: B3 =  55.0_wp / 21.0_wp
     real(wp), parameter :: B4 = -22.0_wp / 21.0_wp
     real(wp), parameter :: B5 =   5.0_wp / 21.0_wp
     real(wp), parameter :: B6 =  -1.0_wp / 42.0_wp


     num_ion_types = size(ions,1)

     dipole_now = dipoles_ions_step(1)
     dipole_m1 = dipoles_ions_step(2)
     dipole_m2 = dipoles_ions_step(3)
     dipole_m3 = dipoles_ions_step(4)
     dipole_m4 = dipoles_ions_step(5)
     dipole_m5 = dipoles_ions_step(6)
     dipole_m6 = dipoles_ions_step(7)


     !$acc data &
     !$acc present(ions(:), dipoles_ions(:,:,:))
     !$acc parallel
     !$acc loop
     do itype = 1, num_ion_types
        alpha = ions(itype)%polarizability
        if (alpha .ne. 0.0_wp) then
           !$acc loop
           do i = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count
             dipoles_ions(i, 1, dipole_now) = B1 * dipoles_ions(i, 1, dipole_m1) + B2 * dipoles_ions(i, 1, dipole_m2) &
                  + B3 * dipoles_ions(i, 1, dipole_m3) + B4 * dipoles_ions(i, 1, dipole_m4) + B5 * dipoles_ions(i, 1, dipole_m5) &
                  + B6 * dipoles_ions(i, 1, dipole_m6)
             dipoles_ions(i, 2, dipole_now) = B1 * dipoles_ions(i, 2, dipole_m1) + B2 * dipoles_ions(i, 2, dipole_m2) &
                  + B3 * dipoles_ions(i, 2, dipole_m3) + B4 * dipoles_ions(i, 2, dipole_m4) + B5 * dipoles_ions(i, 2, dipole_m5) &
                  + B6 * dipoles_ions(i, 2, dipole_m6)
             dipoles_ions(i, 3, dipole_now) = B1 * dipoles_ions(i, 3, dipole_m1) + B2 * dipoles_ions(i, 3, dipole_m2) &
                  + B3 * dipoles_ions(i, 3, dipole_m3) + B4 * dipoles_ions(i, 3, dipole_m4) + B5 * dipoles_ions(i, 3, dipole_m5) &
                  + B6 * dipoles_ions(i, 3, dipole_m6)
           end do
           !$acc end loop
        end if
     end do
     !$acc end loop
     !$acc end parallel
     !$acc end data
 end subroutine predictor

   ! ================================================================================
   include 'update_diag_block_boundaries.inc'
   include 'update_tri_block_boundaries.inc'
   include 'update_other_block_boundaries.inc'

end module MW_ions_dipoles
