module MW_dipoles_and_electrodes
   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_electrode, only: MW_electrode_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

contains

   !================================================================================
   ! Evaluate b, the potential on electrode atoms due to melt ions
   subroutine setup_b(system, algorithms)
      use MW_coulomb, only: &
         MW_coulomb_qmelt2Qelec_potential => qmelt2Qelec_potential ,&
         MW_coulomb_Qelec2mumelt_electricfield => Qelec2mumelt_electricfield ,&
         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, &
         MW_external_field_Efield_elec_potential => Efield_elec_potential, &
         MW_external_field_Dfield_elec_potentialb => Dfield_elec_potentialb
      implicit none
      ! Parameters
      ! ----------
      type(MW_system_t), intent(inout) :: system
      type(MW_algorithms_t), intent(inout) :: algorithms

      ! Locals
      ! ------
      integer :: num_ions, num_atoms
      integer :: xyz_now, dipoles_now, q_atoms_now
      integer :: i, itype, ix, iy, iz, ei
      real(wp) :: alpha
      real(wp), dimension(system%num_ions,3) :: efield, efield2, efield3, efield4
      real(wp) :: Vi
      real(wp), dimension(system%num_atoms,4) :: pot_atoms
      real(wp), dimension(system%num_atoms) :: pot_field

      num_ions = system%num_ions
      num_atoms = system%num_atoms
      xyz_now = system%xyz_ions_step(1)
      dipoles_now = system%dipoles_ions_step(1)
      q_atoms_now = system%q_atoms_step(1)

      efield=0.0_wp
      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.0_wp
      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.0_wp
      if (system%electrode_constant_charge) then
         call MW_coulomb_Qelec2mumelt_electricfield(system%num_PBC, system%localwork, system%ewald, system%box, &
            system%ions, system%xyz_ions(:,:,xyz_now), system%electrodes, system%xyz_atoms, &
            system%q_atoms(:,q_atoms_now), system%type_electrode_atoms(:), efield3)
      endif

      efield4=0.0_wp
      if (system%field%field_type > 0) then
         call MW_external_field_field2mumelt_electricfieldb(system%field, system%ions,&
            system%polarization_field, efield4)
      endif

      efield=efield+efield2+efield3+efield4
      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

      if (.not.system%electrode_constant_charge) then
         pot_atoms(:,:)=0.0_wp
         call MW_coulomb_qmelt2Qelec_potential(system%num_PBC, system%localwork, &
                 system%ewald, system%box, system%ions, system%xyz_ions(:,:,xyz_now), system%type_ions, &
            system%electrodes, system%xyz_atoms, pot_atoms)

         pot_field=0.0_wp
         if (system%field%field_type == FIELD_TYPE_ELECTRIC) then
            call MW_external_field_Efield_elec_potential(system%field, &
               system%electrodes, system%xyz_atoms, pot_field)
         else if (system%field%field_type == FIELD_TYPE_DISPLACEMENT) then
            call MW_external_field_Dfield_elec_potentialb(system%field, &
               system%electrodes, system%xyz_atoms, system%polarization_field, pot_field)
         endif

         do ei = 1, system%num_elec                ! loop over each electrode wall
            Vi = system%electrodes(ei)%V           ! potential applied on the electrode
            if (algorithms%maze%calc_elec) then
               do i = system%electrodes(ei)%offset+1, system%electrodes(ei)%offset + system%electrodes(ei)%count
                  algorithms%maze%const_constr(3*num_ions + i) = (-Vi + &
                     (pot_atoms(i,1) + pot_atoms(i,2) + pot_atoms(i,3) + pot_field(i)))
               end do
            else if (algorithms%matrix_inversion%calc_elec) then
               do i = system%electrodes(ei)%offset+1, system%electrodes(ei)%offset + system%electrodes(ei)%count
                  algorithms%matrix_inversion%b(3*num_ions + i) = Vi - &
                     (pot_atoms(i,1) + pot_atoms(i,2) + pot_atoms(i,3) + pot_field(i))
               end do
            else
               do i = system%electrodes(ei)%offset+1, system%electrodes(ei)%offset + system%electrodes(ei)%count
                  algorithms%cg%b(3*num_ions + i) = Vi - &
                     (pot_atoms(i,1) + pot_atoms(i,2) + pot_atoms(i,3) + pot_field(i))
               end do
            end if
            if (system%electrode_charge_neutrality.and.algorithms%maze%calc_elec) then
               if (abs(system%total_ions_charge) > 1.0e-12_wp .and. system%global_charge_neutrality) then
                  algorithms%maze%const_constr(3*num_ions + num_atoms + 1) = -system%total_ions_charge
               else
                  algorithms%maze%const_constr(3*num_ions + num_atoms + 1) = 0.0_wp
               end if
            end if
         end do

      end if
   end subroutine setup_b

   !================================================================================
   ! Compute A*x, the potential on electrode atoms due to electrode atoms
   !
   subroutine apply_A(system, x, y)
      use MW_coulomb, only: &
         MW_coulomb_mumelt2Qelec_potential => mumelt2Qelec_potential, &
         MW_coulomb_Qelec2Qelec_potential => Qelec2Qelec_potential, &
         MW_coulomb_mumelt2mumelt_electricfield => mumelt2mumelt_electricfield, &
         MW_coulomb_Qelec2mumelt_electricfield => Qelec2mumelt_electricfield, &
         MW_coulomb_fix_molecule_mumelt2mumelt_electricfield => fix_molecule_mumelt2mumelt_electricfield
      use MW_external_field, only: &
         MW_external_field_field2mumelt_electricfieldA => field2mumelt_electricfieldA, &
         MW_external_field_Dfield_elec_potentialA => Dfield_elec_potentialA
      use MW_thomas_fermi, only: &
         MW_thomas_fermi_kinetic_elec_potential => kinetic_elec_potential
      implicit none
      ! Parameters
      ! ----------
      type(MW_system_t), intent(inout) :: system
      real(wp), intent(in) :: x(:)
      real(wp), intent(out) :: y(:)

      ! Locals
      ! ------
      integer :: num_ions, num_atoms
      integer :: i, itype, ix, iy, iz
      integer :: xyz_now, q_atoms_now
      real(wp) :: alpha
      real(wp), dimension(system%num_ions,3) :: efield, efield2, efield3, efield4
      real(wp), dimension(system%num_ions,3) :: dipoles
      real(wp), dimension(system%num_atoms,4) :: pot_atoms, pot_atoms2
      real(wp), dimension(system%num_atoms) :: q_atoms, pot_field, pot_quantum

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

      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)
      end do
      if (.not.system%electrode_constant_charge) then
         do i=1,num_atoms
            q_atoms(i) = x(3*num_ions+i)
         end do
      end if

      efield=0.0_wp
      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)
      end if

      efield3=0.0_wp
      if (.not.system%electrode_constant_charge) then
         call MW_coulomb_Qelec2mumelt_electricfield(system%num_PBC, system%localwork, system%ewald, system%box, &
            system%ions, system%xyz_ions(:,:,xyz_now), system%electrodes, system%xyz_atoms, &
            q_atoms, system%type_electrode_atoms(:), efield3)
      end if

      efield4=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, q_atoms, system%ions, dipoles, system%polarization_field, efield4)
      end if

      efield=efield+efield2+efield3+efield4
      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
            end do
         end if
      end do

      if (.not.system%electrode_constant_charge) then
         pot_atoms=0.0_wp
         call MW_coulomb_mumelt2Qelec_potential(system%num_PBC, system%localwork, system%ewald, system%box, &
               system%ions, system%xyz_ions(:,:,xyz_now), system%electrodes, system%xyz_atoms, dipoles, &
               system%type_electrode_atoms, pot_atoms)

         pot_atoms2=0.0_wp
         call MW_coulomb_Qelec2Qelec_potential(system%num_PBC, system%localwork, system%ewald, system%box, &
               system%electrodes, system%xyz_atoms, q_atoms, pot_atoms2)

         pot_field=0.0_wp
         ! Compute the contribution due to the external electric displacement
         if (system%field%field_type == FIELD_TYPE_DISPLACEMENT) then
            call MW_external_field_Dfield_elec_potentialA(system%box, system%field, &
                  system%electrodes, system%xyz_atoms, q_atoms, system%polarization_field, pot_field)
         endif

         ! TO DO: What about Thomas-Fermi?
         !   Compute the quantum corrections to A * x in the Thomas-Fermi framework
         pot_quantum=0.0_wp
         call MW_thomas_fermi_kinetic_elec_potential(system%thomas_fermi, &
            system%electrodes, q_atoms, pot_quantum)

         pot_atoms=pot_atoms+pot_atoms2
         do i = 1, num_atoms
            y(3*num_ions+i) = (pot_atoms(i,1) + pot_atoms(i,2) &
               +pot_atoms(i,3) + pot_atoms(i,4) + pot_field(i) + pot_quantum(i))
         end do
      end if   
   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, &
         MW_coulomb_gradQelec_electricfield => gradQelec_electricfield, &
         MW_coulomb_gradQelec_potential => gradQelec_potential
      use MW_errors, only: &
         MW_errors_allocate_error => allocate_error
      use MW_external_field, only: &
         MW_external_field_Dfield_gradQelec_potential => Dfield_gradQelec_potential
      use MW_thomas_fermi, only: &
         MW_thomas_fermi_kinetic_gradQelec_potential => kinetic_gradQelec_potential
      implicit none

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

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

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

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

      call MW_coulomb_gradmumelt_electricfield(system%num_PBC, system%localwork, system%ewald, &
      system%box, system%ions, system%xyz_ions(:,:,xyz_now), hessian(1:3*num_ions,1:3*num_ions))
      dip_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), dip_fix_mol_hessian(:,:))
      endif
      hessian(1:3*num_ions,1:3*num_ions) = hessian(1:3*num_ions,1:3*num_ions) + dip_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)
         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/system%ions(itype)%polarizability
            hessian(iy,iy) = hessian(iy,iy) + 1.0_wp/system%ions(itype)%polarizability
            hessian(iz,iz) = hessian(iz,iz) + 1.0_wp/system%ions(itype)%polarizability
         end do
      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

      call MW_coulomb_gradQelec_electricfield(system%num_PBC, system%localwork, system%ewald, system%box, &
         system%ions, system%xyz_ions(:,:,xyz_now), system%electrodes, system%xyz_atoms, &
         hessian(1:3*num_ions,3*num_ions+1:3*num_ions+num_atoms))

      do i = 1, num_atoms
         do j = 1, num_ions

            jx = j
            jy = j + num_ions
            jz = j + num_ions*2

            hessian(3*num_ions+i,jx) = hessian(jx,3*num_ions+i)
            hessian(3*num_ions+i,jy) = hessian(jy,3*num_ions+i)
            hessian(3*num_ions+i,jz) = hessian(jz,3*num_ions+i)

         end do
      end do

      call MW_coulomb_gradQelec_potential(system%num_PBC, system%localwork, system%ewald, system%box, &
         system%electrodes, system%xyz_atoms, hessian(3*num_ions+1:3*num_ions+num_atoms,3*num_ions+1:3*num_ions+num_atoms))

     ! Compute the contribution due to the external electric displacement
     if (system%field%field_type == FIELD_TYPE_DISPLACEMENT) then
         call MW_external_field_Dfield_gradQelec_potential(system%box, system%field, &
             system%electrodes, system%xyz_atoms, system%q_atoms(:, system%q_atoms_step(1)), system%polarization_field, &
             hessian(3*num_ions+1:3*num_ions+num_atoms,3*num_ions+1:3*num_ions+num_atoms))
     endif

     call MW_thomas_fermi_kinetic_gradQelec_potential(system%thomas_fermi, &
          system%electrodes, hessian(3*num_ions+1:3*num_ions+num_atoms,3*num_ions+1:3*num_ions+num_atoms))

     do i = 1, num_atoms
         do j = 1, i-1
            hessian(3*num_ions+i,3*num_ions+j) = hessian(3*num_ions+j,3*num_ions+i)
         end do
     end do

   end subroutine compute_hessian

   ! Parallelized routine for computing matrix vector products invloving dipoles and electrodes together
   ! The matrix has to be Symmetric
   ! To be implemented
   subroutine SMatrix_vector_product_dipelec(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(:)

   end subroutine SMatrix_vector_product_dipelec

   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, n_atoms
      real(wp) :: alpha

      n_ions = system%num_ions
      n_atoms = system%num_atoms

      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

      if (.not.system%electrode_constant_charge) then
         call MW_coulomb_diag_gradQelec_potential(system%num_PBC, system%localwork, system%ewald, system%box, &
            system%electrodes, algorithms%cg%jacobi_precond(3*n_ions+1:3*n_ions+n_atoms))
      end if

   end subroutine compute_jacobi_precond

   !================================================================================
   !> Compute charges on electrode
   !!
   !! The charge 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_ions_dipoles, only: &
         MW_ions_dipoles_predictor => predictor
      use MW_electrode_charge, only: &
         MW_electrode_charge_predictor => predictor
      use MW_cg, only: &
         MW_cg_solve_preconditioned => solve_preconditioned, &
         MW_cg_solve => solve
      use MW_maze, only: &
         MW_maze_compute_x0 => compute_x0, &
         MW_maze_update_history => update_history, &
         MW_maze_solve_invert => solve_invert, &
         MW_maze_solve_block_iterate => solve_block_iterate, &
         MW_maze_solve_CG_store => solve_CG_store, &
         MW_maze_solve_CG_nostore => solve_CG_nostore, &
         MW_maze_solve_shake => solve_shake
      use MW_matrix_inversion, only: &
         MW_matrix_inversion_solve => solve
      use MW_external_field, only: &
         MW_external_field_elec_polarization => elec_polarization, &
         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 :: i, itype
      integer :: num_ions, num_atoms
      integer :: q_atoms_tmdt, q_atoms_t, q_atoms_tpdt
      integer :: dipoles_ions_tpdt, dipoles_ions_t, dipoles_ions_tmdt
      real(wp) :: total_charges, average_charges, alpha
      real(wp), dimension(3*system%num_ions+system%num_atoms) :: x
      real(wp), dimension(3*system%num_ions+system%num_atoms) :: xt, xtmdt

      num_ions = system%num_ions
      num_atoms = system%num_atoms

      dipoles_ions_tpdt = system%dipoles_ions_step(1)
      dipoles_ions_t = system%dipoles_ions_step(2)
      dipoles_ions_tmdt = system%dipoles_ions_step(3)

      q_atoms_tpdt = system%q_atoms_step(1)
      q_atoms_t = system%q_atoms_step(2)
      q_atoms_tmdt = system%q_atoms_step(3)

      if (algorithms%constant_charge) then

         ! Dipoles and charges are set equal to previous step (no predictor available for dipoles yet).
         if (system%dipoles_use_predictor) then
            call MW_ions_dipoles_predictor(system%dipoles_ions_step, system%ions, system%dipoles_ions)
            do itype = 1, system%num_elec
               do i = system%electrodes(itype)%offset+1, system%electrodes(itype)%offset + system%electrodes(itype)%count
                  system%q_atoms(i, q_atoms_tpdt) = system%q_atoms(i, q_atoms_t)
               end do
            end do
         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
            do itype = 1, system%num_elec
               do i = system%electrodes(itype)%offset+1, system%electrodes(itype)%offset + system%electrodes(itype)%count
                  system%q_atoms(i, q_atoms_tpdt) = system%q_atoms(i, q_atoms_t)
               end do
            end do
         end if

         x(           1:1*num_ions)=system%dipoles_ions(:,1,dipoles_ions_tpdt)
         x(  num_ions+1:2*num_ions)=system%dipoles_ions(:,2,dipoles_ions_tpdt)
         x(2*num_ions+1:3*num_ions)=system%dipoles_ions(:,3,dipoles_ions_tpdt)

         x(3*num_ions+1:3*num_ions+num_atoms)=system%q_atoms(:,q_atoms_tpdt)

         ! The conjugate gradient algorithm is executed on dipoles [x(1:3*num_ions)] only.
         call setup_b(system, algorithms)
         call MW_cg_solve_preconditioned(algorithms%cg, system, apply_A, x(1:3*num_ions))

      else

         if (algorithms%maze%calc_dip .and. algorithms%maze%calc_elec) then

            x(           1:1*num_ions)=system%dipoles_ions(:,1,dipoles_ions_tpdt)
            x(  num_ions+1:2*num_ions)=system%dipoles_ions(:,2,dipoles_ions_tpdt)
            x(2*num_ions+1:3*num_ions)=system%dipoles_ions(:,3,dipoles_ions_tpdt)
            xt(           1:1*num_ions)=system%dipoles_ions(:,1,dipoles_ions_t)
            xt(  num_ions+1:2*num_ions)=system%dipoles_ions(:,2,dipoles_ions_t)
            xt(2*num_ions+1:3*num_ions)=system%dipoles_ions(:,3,dipoles_ions_t)
            xtmdt(           1:1*num_ions)=system%dipoles_ions(:,1,dipoles_ions_tmdt)
            xtmdt(  num_ions+1:2*num_ions)=system%dipoles_ions(:,2,dipoles_ions_tmdt)
            xtmdt(2*num_ions+1:3*num_ions)=system%dipoles_ions(:,3,dipoles_ions_tmdt)

            x(3*num_ions+1:3*num_ions+num_atoms)=system%q_atoms(:,q_atoms_tpdt)
            xt(3*num_ions+1:3*num_ions+num_atoms)=system%q_atoms(:,q_atoms_t)
            xtmdt(3*num_ions+1:3*num_ions+num_atoms)=system%q_atoms(:,q_atoms_tmdt)

!            call setup_b(system, algorithms)
!            if ((.not.algorithms%maze%constant_constr_grad).or.(.not.algorithms%maze%CG_nostore)) then
!               call MW_maze_update_history(algorithms%maze) !Only updates gradient of constraints
!            end if
!            call MW_maze_compute_x0(algorithms%maze, 1.0_wp,&
!                 system%q_atoms(:, q_atoms_tmdt), system%q_atoms(:, q_atoms_t), system%q_atoms(:, q_atoms_tpdt))
!            if (algorithms%maze%invert) then
!               call MW_maze_solve_invert(algorithms%maze, system, &
!                    compute_hessian, SMatrix_vector_product_elec, system%q_atoms(:, q_atoms_tpdt))
!            else if (algorithms%maze%shake) then
!               call MW_maze_solve_shake(algorithms%maze, system, &
!                    compute_hessian, system%q_atoms(:, q_atoms_tpdt))
!            else if (algorithms%maze%block_iterate) then
!               call MW_maze_solve_block_iterate(algorithms%maze, system, &
!                    compute_hessian, system%q_atoms(:, q_atoms_tpdt))
!            else if (algorithms%maze%CG_store_prec.or.algorithms%maze%CG_store_noprec) then
!               call MW_maze_solve_CG_store(algorithms%maze, system, &
!                    compute_hessian, SMatrix_vector_product_elec, &
!                    vector_vector_inner_product_elec, system%q_atoms(:, q_atoms_tpdt))
!            else if (algorithms%maze%CG_nostore) then
!               call MW_maze_solve_CG_nostore(algorithms%maze, system, &
!                    apply_A, vector_vector_inner_product_elec, system%q_atoms(:, q_atoms_tpdt))
!            end if
         else if (algorithms%matrix_inversion%calc_dip .and. algorithms%matrix_inversion%calc_elec) then

            x(           1:1*num_ions)=system%dipoles_ions(1:num_ions,1,dipoles_ions_tpdt)
            x(  num_ions+1:2*num_ions)=system%dipoles_ions(1:num_ions,2,dipoles_ions_tpdt)
            x(2*num_ions+1:3*num_ions)=system%dipoles_ions(1:num_ions,3,dipoles_ions_tpdt)

            x(3*num_ions+1:3*num_ions+num_atoms)=system%q_atoms(1:num_atoms,q_atoms_tpdt)

            call setup_b(system, algorithms)
            call MW_matrix_inversion_solve(algorithms%matrix_inversion, system, compute_hessian, SMatrix_vector_product_dipelec, x)

            system%dipoles_ions(1:num_ions,1,dipoles_ions_tpdt) = x(           1:1*num_ions)
            system%dipoles_ions(1:num_ions,2,dipoles_ions_tpdt) = x(  num_ions+1:2*num_ions)
            system%dipoles_ions(1:num_ions,3,dipoles_ions_tpdt) = x(2*num_ions+1:3*num_ions)

            system%q_atoms(1:num_atoms, q_atoms_tpdt) = x(3*num_ions+1:3*num_ions+num_atoms)

            ! update the contribution of the electrode to the polarization
            call MW_external_field_elec_polarization(system%box, system%electrodes(:), system%xyz_atoms(:,:), &
               system%q_atoms(:, q_atoms_tpdt), system%polarization_field(:,:))
            call MW_external_field_dipole_polarization(system%box, system%ions, &
               system%dipoles_ions(:,:,dipoles_ions_tpdt), system%polarization_field(:,:))

         else

            if (system%dipoles_use_predictor .and. system%electrode_use_predictor) then
               call MW_ions_dipoles_predictor(system%dipoles_ions_step, system%ions, system%dipoles_ions)
               call MW_electrode_charge_predictor(system%q_atoms_step, system%electrodes, system%q_atoms)
            else if (system%dipoles_use_predictor) then
               call MW_ions_dipoles_predictor(system%dipoles_ions_step, system%ions, system%dipoles_ions)
               do itype = 1, system%num_elec
                  do i = system%electrodes(itype)%offset+1, system%electrodes(itype)%offset + system%electrodes(itype)%count
                     system%q_atoms(i, q_atoms_tpdt) = system%q_atoms(i, q_atoms_t)
                  end do
               end do
            else if (system%electrode_use_predictor) then
               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
               call MW_electrode_charge_predictor(system%q_atoms_step, system%electrodes, system%q_atoms)
            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
               do itype = 1, system%num_elec
                  do i = system%electrodes(itype)%offset+1, system%electrodes(itype)%offset + system%electrodes(itype)%count
                     system%q_atoms(i, q_atoms_tpdt) = system%q_atoms(i, q_atoms_t)
                  end do
               end do
            end if
            x(           1:1*num_ions)=system%dipoles_ions(1:num_ions,1,dipoles_ions_tpdt)
            x(  num_ions+1:2*num_ions)=system%dipoles_ions(1:num_ions,2,dipoles_ions_tpdt)
            x(2*num_ions+1:3*num_ions)=system%dipoles_ions(1:num_ions,3,dipoles_ions_tpdt)

            x(3*num_ions+1:3*num_ions+num_atoms)=system%q_atoms(1:num_atoms,q_atoms_tpdt)

            ! Enforce charge neutrality constraint on initial vector
            if (system%electrode_charge_neutrality) then
               total_charges = sum(x(3*num_ions+1:3*num_ions+num_atoms))
               average_charges = total_charges / real(num_atoms,wp)
               if (system%global_charge_neutrality) then
                  average_charges = average_charges + system%total_ions_charge / real(system%num_atoms,wp)
               end if
               x(3*num_ions+1:3*num_ions+num_atoms) = x(3*num_ions+1:3*num_ions+num_atoms) - average_charges
            end if

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

            ! compute \chi = V_i - \frac{\partial U_{es}}{\partial q_i}
            ! cg%b stores V_i - \frac{\partial U_{es}}{\partial q_i} for melt->elec interactions
            ! apply_A(system, system%q_atoms(:, q_atoms_now), cg%Ad) computes \frac{\partial U_{es}}{\partial q_i} for elec<->elec interactions
            ! chi(:) is stored in cg%Ad(:)
            if (system%electrode_charge_neutrality) then
               call apply_A(system, x, algorithms%cg%Ad)
               system%chi = sum(algorithms%cg%b(3*num_ions+1:3*num_ions+num_atoms) &
                     - algorithms%cg%Ad(3*num_ions+1:3*num_ions+num_atoms)) / real(num_atoms,wp)
            end if
         end if
      end if

      system%dipoles_ions(1:num_ions,1,dipoles_ions_tpdt) = x(           1:1*num_ions)
      system%dipoles_ions(1:num_ions,2,dipoles_ions_tpdt) = x(  num_ions+1:2*num_ions)
      system%dipoles_ions(1:num_ions,3,dipoles_ions_tpdt) = x(2*num_ions+1:3*num_ions)

      system%q_atoms(1:num_atoms, q_atoms_tpdt) = x(3*num_ions+1:3*num_ions+num_atoms)

      ! update the contribution of the electrode to the polarization
      call MW_external_field_elec_polarization(system%box, system%electrodes(:), system%xyz_atoms(:,:), &
         system%q_atoms(:, q_atoms_tpdt), system%polarization_field(:,:))
      call MW_external_field_dipole_polarization(system%box, system%ions, &
         system%dipoles_ions(:,:,dipoles_ions_tpdt), system%polarization_field(:,:))

   end subroutine compute

end module MW_dipoles_and_electrodes
