module MW_electrode_charge
#ifndef MW_SERIAL
   use MPI
#endif
   use MW_parallel, only: MW_COMM_WORLD
   use MW_stdio
   use MW_timers
   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_electrode, only: MW_electrode_t
   use MW_external_field, only: MW_external_field_t, &
        FIELD_TYPE_ELECTRIC, FIELD_TYPE_DISPLACEMENT
   use MW_thomas_fermi, only: MW_thomas_fermi_t
   use MW_ion, only: MW_ion_t

  implicit none
  private

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

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
     use MW_external_field, only: &
         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 :: xyz_now
    integer :: ei, i
    real(wp) :: Vi

    !$acc data &
    !$acc present(system, system%electrodes(:), system%coulomb_potential_atoms(:,:), &
    !$acc         algorithms, algorithms%cg, algorithms%cg%b(:), &
    !$acc        algorithms%matrix_inversion, &
    !$acc        algorithms%matrix_inversion%b(:),algorithms%matrix_inversion%A(:,:), &
    !$acc        algorithms%matrix_inversion%Ainv(:,:),algorithms%matrix_inversion%Ax(:), &
    !$acc        algorithms%matrix_inversion%q0(:))    

    xyz_now = system%xyz_ions_step(1)

    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, system%coulomb_potential_atoms)

    if (system%field%field_type == FIELD_TYPE_ELECTRIC) then
       call MW_external_field_Efield_elec_potential(system%field, &
            system%electrodes, system%xyz_atoms, system%external_field_potential_atoms(:,FIELD_TYPE_ELECTRIC,1))
    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, &
            system%external_field_potential_atoms(:,FIELD_TYPE_DISPLACEMENT,1))
    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
          !$acc parallel
          !$acc loop
          do i = system%electrodes(ei)%offset+1, system%electrodes(ei)%offset + system%electrodes(ei)%count
             algorithms%maze%const_constr(i) = (-Vi + &
                  (system%coulomb_potential_atoms(i,1) + &
                   system%coulomb_potential_atoms(i,2) + &
                   system%coulomb_potential_atoms(i,3) + &
                   system%external_field_potential_atoms(i, FIELD_TYPE_ELECTRIC,1)+ &
                   system%external_field_potential_atoms(i, FIELD_TYPE_DISPLACEMENT,1)))
          end do
          !$acc end loop
          !$acc end parallel
       else if (algorithms%matrix_inversion%calc_elec) then
          !$acc parallel
          !$acc loop
          do i = system%electrodes(ei)%offset+1, system%electrodes(ei)%offset + system%electrodes(ei)%count
             algorithms%matrix_inversion%b(i) = Vi - &
                  (system%coulomb_potential_atoms(i,1) + &
                   system%coulomb_potential_atoms(i,2) + &
                   system%coulomb_potential_atoms(i,3) + &
                   system%external_field_potential_atoms(i, FIELD_TYPE_ELECTRIC,1)+ &
                   system%external_field_potential_atoms(i, FIELD_TYPE_DISPLACEMENT,1))
          end do
          !$acc end loop
          !$acc end parallel
       else
          !$acc parallel
          !$acc loop
          do i = system%electrodes(ei)%offset+1, system%electrodes(ei)%offset + system%electrodes(ei)%count
             algorithms%cg%b(i) = Vi - &
                  (system%coulomb_potential_atoms(i,1) + &
                   system%coulomb_potential_atoms(i,2) + &
                   system%coulomb_potential_atoms(i,3) + &
                   system%external_field_potential_atoms(i, FIELD_TYPE_ELECTRIC,1)+ &
                   system%external_field_potential_atoms(i, FIELD_TYPE_DISPLACEMENT,1))
          end do
          !$acc end loop
          !$acc end parallel
       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(system%num_atoms+1) = -system%total_ions_charge
          else
             algorithms%maze%const_constr(system%num_atoms+1) = 0.0_wp
          end if
       end if
    end do
    !$acc end data
  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_Qelec2Qelec_potential => Qelec2Qelec_potential
    use MW_external_field, only: &
         MW_external_field_Dfield_elec_potentialA => Dfield_elec_potentialA, &
         MW_external_field_elec_polarization => elec_polarization
    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 :: i


    !$acc data &
    !$acc present(y, system, system%coulomb_potential_atoms)

    call MW_coulomb_Qelec2Qelec_potential(system%num_PBC, system%localwork, system%ewald, system%box, &
         system%electrodes, system%xyz_atoms, x, system%coulomb_potential_atoms)

    ! 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, x, system%polarization_field, &
            system%external_field_potential_atoms(:,FIELD_TYPE_DISPLACEMENT,2))
    endif
    ! Compute the quantum corrections to A * x in the Thomas-Fermi framework
    call MW_thomas_fermi_kinetic_elec_potential(system%thomas_fermi, &
         system%electrodes, x, system%quantum_potential_atoms(:,1))

    !$acc parallel loop
    do i = 1, system%num_atoms
       y(i) = (system%coulomb_potential_atoms(i,1) + &
            system%coulomb_potential_atoms(i,2) + &
            system%coulomb_potential_atoms(i,3) + &
            system%coulomb_potential_atoms(i,4) + &
            system%external_field_potential_atoms(i, FIELD_TYPE_DISPLACEMENT,2) + &
            system%quantum_potential_atoms(i,1))
    end do
    !$acc end parallel loop
    !$acc end data
  end subroutine apply_A

  subroutine compute_hessian(system, hessian)
     use MW_coulomb, only: &
         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_atoms
     integer  :: i, j

     num_atoms = system%num_atoms

     call MW_coulomb_gradQelec_potential(system%num_PBC,system%localwork,system%ewald,system%box, &
     system%electrodes, system%xyz_atoms, hessian(:,:))

     ! 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(:,:))
     endif

     call MW_thomas_fermi_kinetic_gradQelec_potential(system%thomas_fermi, &
          system%electrodes, hessian(:,:))

     do i = 1, num_atoms
        do j = 1, i-1
           hessian(i,j) = hessian(j,i)
        end do
     end do


  end subroutine compute_hessian

  ! Parallelized routine for computing matrix vector products invloving electrodes
  ! The matrix has to be Symmetric
  subroutine SMatrix_vector_product_elec(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_elec_types
     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
     real(wp) :: matvec_product_i, vector_i

     call MW_timers_start(TIMER_MINQ_MATVEC)

     num_elec_types = size(system%electrodes,1)


     !$acc data present(matrix(:,:),vector(:),matvec_product(:))

     !$acc kernels
     matvec_product(:) = 0.0_wp
     !$acc end kernels
     ! Different kinds of electrodes - begin
     do itype = 1, num_elec_types
        count_n = system%electrodes(itype)%count
        offset_n = system%electrodes(itype)%offset
        do jtype = 1, itype-1
           count_m = system%electrodes(jtype)%count
           offset_m = system%electrodes(jtype)%offset

           ! Number of blocks with full interactions (below the diagonal)
           num_block_full = system%localwork%pair_atom2atom_num_block_full_local(itype,jtype)
           iblock_offset = system%localwork%pair_atom2atom_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)
              !$acc parallel loop private(matvec_product_i,vector_i)
              do i = istart, iend
                 matvec_product_i = 0.0_wp
                 vector_i = vector(i)
                 !$acc loop reduction(+:matvec_product_i)
                 do j = jstart, jend
                    matvec_product_i = matvec_product_i + matrix(j,i)*vector(j)
                    !$acc atomic update
                    matvec_product(j) = matvec_product(j) + matrix(j,i)*vector_i
                 end do
                 !$acc end loop
                 matvec_product(i) = matvec_product(i) + matvec_product_i
              end do
              !$acc end parallel loop
           end do
        end do
        ! Different kinds of electrodes - end

        ! Same kind of electrodes - begin
        ! Blocks on the diagonal compute only half of the pair interactions
        num_block_diag = system%localwork%pair_atom2atom_num_block_diag_local(itype)
        iblock_offset = system%localwork%pair_atom2atom_diag_iblock_offset(itype)
        do iblock = 1, num_block_diag
           call update_diag_block_boundaries(iblock_offset+iblock, offset_n, count_n, istart,  iend)
           !$acc parallel loop private(matvec_product_i)
           do i = istart, iend
              matvec_product_i = 0.0_wp
              !$acc loop reduction(+:matvec_product_i)
              do j = istart, iend
                 matvec_product_i = matvec_product_i + matrix(j,i)*vector(j)
              end do
              !$acc end loop
              matvec_product(i) = matvec_product(i) + matvec_product_i
           end do
           !$acc end parallel loop
        end do

        ! Number of blocks with full interactions (below the diagonal)
        num_block_full = system%localwork%pair_atom2atom_num_block_full_local(itype,itype)
        iblock_offset = system%localwork%pair_atom2atom_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)
           !$acc parallel loop private(matvec_product_i,vector_i)
           do i = istart, iend
              matvec_product_i = 0.0_wp
              vector_i = vector(i)
              !$acc loop reduction(+:matvec_product_i)
              do j = jstart, jend
                 matvec_product_i = matvec_product_i + matrix(j,i)*vector(j)
                 !$acc atomic update
                 matvec_product(j) = matvec_product(j) + matrix(j,i)*vector_i
              end do
              !$acc end loop
              matvec_product(i) = matvec_product(i) + matvec_product_i
           end do
           !$acc end parallel loop
        end do

     end do
     ! Same kind of electrodes - end

     !$acc update device(matvec_product(:))
     !$acc end data


#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
     call MW_timers_stop(TIMER_MINQ_MATVEC)

  end subroutine SMatrix_vector_product_elec

  subroutine vector_vector_inner_product_elec(system, vector_v, vector_w, inner_prod)
     implicit none

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

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

     !Locals
     !---------------
     integer  :: itype, i, iatom
     integer  :: num_elec_types
     integer  :: ierr

     inner_prod = 0.0_wp
     num_elec_types = size(system%electrodes, 1)

     do itype = 1, num_elec_types
        do i = 1, system%localwork%count_atoms(itype)
           iatom = system%localwork%offset_atoms(itype) + i
           inner_prod = inner_prod + vector_v(iatom)*vector_w(iatom)
        end do
     end do

#ifndef MW_SERIAL
      call MPI_Allreduce(MPI_IN_PLACE, inner_prod, 1, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)
#endif
  end subroutine vector_vector_inner_product_elec

  subroutine compute_jacobi_precond(system, algorithms)
     use MW_coulomb, only: &
     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 :: n_atoms

     n_atoms = system%num_atoms

     call MW_coulomb_diag_gradQelec_potential(system%num_PBC, system%localwork, system%ewald, system%box, &
         system%electrodes, algorithms%cg%jacobi_precond(:))
  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_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
    implicit none

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

    ! Local
    ! -----
    integer :: i, itype
    integer :: q_atoms_tmdt, q_atoms_t, q_atoms_tpdt
    real(wp) :: total_charges, average_charges

    !$acc data &
    !$acc present (system, system%electrodes(:), system%q_atoms(:,:), &
    !$acc          algorithms, algorithms%cg, algorithms%cg%b, algorithms%cg%Ad)

    q_atoms_tpdt = system%q_atoms_step(1)
    q_atoms_t = system%q_atoms_step(2)
    if (algorithms%maze%calc_elec) q_atoms_tmdt = system%q_atoms_step(3)

    if (algorithms%constant_charge) then
       !$acc parallel
       !$acc loop     
       do itype = 1, system%num_elec
          !$acc loop     
          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
          !$acc end loop
       end do
       !$acc end loop
       !$acc end parallel
    else

       if (algorithms%maze%calc_elec) then

          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, system, 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_elec) then

          call setup_b(system, algorithms)
          call MW_matrix_inversion_solve(algorithms%matrix_inversion, system, compute_hessian, SMatrix_vector_product_elec,&
               system%q_atoms(:, q_atoms_tpdt))

       else

          if (system%electrode_use_predictor) then
             call predictor(system%q_atoms_step, system%electrodes, system%q_atoms)
          else
             ! Initialize with previous step value
             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

          ! Enforce charge neutrality constraint on initial vector
          if (system%electrode_charge_neutrality) then
             !$acc kernels     
             total_charges = sum(system%q_atoms(:, q_atoms_tpdt))
             !$acc end kernels     
             average_charges = total_charges / real(system%num_atoms,wp)
             if (system%global_charge_neutrality) then
                average_charges = average_charges + system%total_ions_charge / real(system%num_atoms,wp)
             end if
             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_tpdt) - average_charges
                end do
             end do
          end if

          call setup_b(system, algorithms)
          ! call MW_cg_solve(algorithms%cg, system, apply_A, system%q_atoms(:, q_atoms_tpdt))
          call MW_cg_solve_preconditioned(algorithms%cg, system, apply_A, system%q_atoms(:, q_atoms_tpdt))

          ! compute \ki = 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
          ! ki(:) is stored in cg%Ad(:)
          if (system%electrode_charge_neutrality) then
             call apply_A(system, system%q_atoms(:, q_atoms_tpdt), algorithms%cg%Ad)
             system%chi(q_atoms_tpdt) = 0.0_wp
             do itype = 1, system%num_elec
                do i = system%electrodes(itype)%offset+1, system%electrodes(itype)%offset + system%electrodes(itype)%count
                   system%chi(q_atoms_tpdt) = system%chi(q_atoms_tpdt) + algorithms%cg%b(i) - algorithms%cg%Ad(i)
                end do
             end do
             system%chi(q_atoms_tpdt) = system%chi(q_atoms_tpdt) / real(system%num_atoms,wp)
          end if
       end if
    end if

    ! 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(:,:))

   !$acc update self(system%q_atoms(:, q_atoms_tpdt))
   !$acc end data

  end subroutine compute

  ! ================================================================================
  !> Predictor for wall atoms charges
  subroutine predictor(q_atoms_step, electrodes, q_atoms)
    implicit none
    ! Parameters
    ! ----------
    integer, intent(in) :: q_atoms_step(:)  !< time step index
    type(MW_electrode_t), intent(in) :: electrodes(:) !< electrode parameters
    real(wp), intent(inout) :: q_atoms(:,:) !< atom charges

    ! Local
    ! -----
    integer :: q_now, q_m1, q_m2, q_m3, q_m4, q_m5, q_m6
    integer :: num_atom_types
    integer :: itype, i

    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_atom_types = size(electrodes,1)

    q_now = q_atoms_step(1)
    q_m1 = q_atoms_step(2)
    q_m2 = q_atoms_step(3)
    q_m3 = q_atoms_step(4)
    q_m4 = q_atoms_step(5)
    q_m5 = q_atoms_step(6)
    q_m6 = q_atoms_step(7)

    do itype = 1, num_atom_types
       do i = electrodes(itype)%offset+1, electrodes(itype)%offset + electrodes(itype)%count
          q_atoms(i, q_now) = B1 * q_atoms(i, q_m1) + B2 * q_atoms(i, q_m2) + B3 * q_atoms(i, q_m3) &
               + B4 * q_atoms(i, q_m4) + B5 * q_atoms(i, q_m5) + B6 * q_atoms(i, q_m6)
       end do
    end do

  end subroutine predictor

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

end module MW_electrode_charge
