! Module defines MaZe type to store MaZe algorithm parameters
module MW_maze
   use MW_timers
   use MW_kinds, only: wp, sp
   use MW_system, only: MW_system_t

   implicit none
   private

   ! Public type
   ! -----------
   public :: MW_maze_t

   ! Public subroutines
   ! ------------------
   public :: set_max_iterations
   public :: set_tolerance
   public :: set_blocks_number
   public :: set_charge_neutrality
   public :: set_write_matrix
   public :: set_sor_param
   public :: reset_statistics
   public :: set_dimensionality
   public :: select_algorithm
   public :: allocate_arrays
   public :: void_type
   public :: setup_history
   public :: update_history
   public :: print_type
   public :: print_statistics
   public :: compute_x0
   public :: solve_invert
   public :: solve_shake
   public :: solve_block_iterate
   public :: solve_CG_store
   public :: solve_CG_nostore

   integer, parameter :: CONSTR_GRAD_NUM_STEPS = 2    !< number of steps stored for gradient of the constraints

   type MW_maze_t

      !Convergence parameters
      integer :: max_iterations !< maximum number of iterations
      real(wp) :: tol           !< tolerance on constraints norm

      !Statistics
      integer  :: last_iteration_count    !< Number of iteration for the last MaZe
      real(wp) :: last_max_constr         !< Maximum constraint for last iteration
      real(sp) :: total_iteration_count   !< Sum of all iteration count

      !Logicals (common or shared)
      logical :: calc_elec = .false.
      logical :: calc_dip = .false.
      logical :: charge_neutrality = .false.
      logical :: constant_constr_grad = .false.
      logical :: firsttime = .true.
      logical :: write_shake_matrix = .false.

      !Type variables (common)
      integer :: n_constr, n_var                  !< dimensions of the work arrays
      real(wp)              :: sor_param          !< successive over relaxation parameter
      real(wp), allocatable :: const_constr(:)    !< Constant part of constraint
      real(wp), allocatable :: ggamma(:)          !< MaZe Lagrange Multipliers
      real(wp), allocatable :: constr_grad(:,:,:) !< Gradient of the constraints

      !History Pointers
      integer :: constr_grad_step(CONSTR_GRAD_NUM_STEPS)

      !Logicals (switches btwn algorithms)
      logical :: shake = .false.
      logical :: invert = .false.
      logical :: block_iterate = .false.
      logical :: CG_store_prec = .false.
      logical :: CG_store_noprec = .false.
      logical :: CG_nostore = .false.

      !Type variables (shake)
      real(wp), allocatable :: denom(:)           !< denominator (constraint gradient dependent)

      !Type variables (block invert and invert)
      integer :: n_blocks = 0, n_red, n_test          !< dimensions of the work arrays for block shake
      real(wp) :: uscale
      real(wp), allocatable :: block_ggamma(:)    !< MaZe Lagrange Multipliers for block shake
      real(wp), allocatable :: Aredinv(:,:,:)     !< Matrix to invert in block shake
      real(wp), allocatable :: constr(:)          !< Non-constant part of constraint
      real(wp), allocatable :: fconstr(:)         !< Constraint forces
      real(wp), allocatable :: Ainv(:,:)          !< Matrix to invert

      !Type variables (CG)
      real(wp), allocatable :: Aggamma(:)    !< MaZe Lagrange Multipliers for block shake
      real(wp), allocatable :: Ap(:)    !< MaZe Lagrange Multipliers for block shake
      real(wp), allocatable :: res(:)    !< MaZe Lagrange Multipliers for block shake
      real(wp), allocatable :: p(:)    !< MaZe Lagrange Multipliers for block shake
      real(wp), allocatable :: eta(:)    !< MaZe Lagrange Multipliers for block shake
      real(wp), allocatable :: jacobi_precond(:)

   end type MW_maze_t

contains

   ! ================================================================================
   ! Set maximum number of iterations
   subroutine set_max_iterations(this, max_iterations)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      integer, intent(in) :: max_iterations
      this%max_iterations = max_iterations
   end subroutine set_max_iterations

   ! ================================================================================
   ! Set tolerance parameter
   subroutine set_tolerance(this, tol)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      real(wp), intent(in) :: tol
      this%tol = tol
   end subroutine set_tolerance

   ! ================================================================================
   ! Set number of blocks
   subroutine set_blocks_number(this, n_blocks)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      integer, intent(in) :: n_blocks
      this%n_blocks = n_blocks
   end subroutine set_blocks_number

   ! ================================================================================
   ! Set force neutral
   subroutine set_charge_neutrality(this, charge_neutrality, uscale)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      logical, intent(in) :: charge_neutrality
      real(wp), intent(in) :: uscale
      this%charge_neutrality = charge_neutrality
      this%uscale = uscale
   end subroutine set_charge_neutrality

   ! ================================================================================
   ! Set write_matrix
   subroutine set_write_matrix(this, write_matrix)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      logical, intent(in) :: write_matrix
      this%write_shake_matrix = write_matrix
   end subroutine set_write_matrix

   ! ================================================================================
   ! Set successive over relaxation parameter
   subroutine set_sor_param(this, sor_param)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      real(wp), intent(in) :: sor_param
      this%sor_param = sor_param
   end subroutine set_sor_param

   ! ================================================================================
   ! reset statistic counters
   subroutine reset_statistics(this)
      implicit none
      type(MW_maze_t), intent(inout) :: this
      this%last_iteration_count = 0
      this%last_max_constr = 0.0_wp
      this%total_iteration_count = 0
   end subroutine reset_statistics

   ! ================================================================================
   ! Set dimensionality
   subroutine set_dimensionality(this, n)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      implicit none
      type(MW_maze_t), intent(inout) :: this
      integer, intent(in) :: n

      this%n_var = n
      if (this%charge_neutrality.and.this%calc_elec) then
         this%n_constr = this%n_var+1
      else
         this%n_constr = this%n_var
      end if

      if ((this%n_blocks .gt. this%n_constr) .or. (this%n_blocks .lt. -3)) then
         call MW_errors_runtime_error("set_dimensionality", "maze.f90", "Unacceptable number of blocks.")
      end if

      if (this%n_blocks .gt. 0) then
         this%n_red = int(this%n_constr/this%n_blocks)
         this%n_test = mod(this%n_constr, this%n_blocks)
      else
         this%n_red = this%n_constr
         this%n_test = 0
      end if

      if (this%n_test /= 0) then
         call MW_errors_runtime_error("set_dimensionality", "maze.f90", "Number of blocks does not divide number of "// &
               "constraints exactly. Function not yet implemented.")
      end if

   end subroutine set_dimensionality

   ! ================================================================================
   ! Select algorithm for constraint enforcement
   subroutine select_algorithm(this, maze_algorithm, maze_cg_store, maze_cg_prec, n_blocks)
      use MW_configuration_line, only: max_word_length
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      implicit none

      integer                   , intent(in) :: n_blocks
      character(max_word_length), intent(in) :: maze_algorithm
      character(max_word_length), intent(in) :: maze_cg_store
      character(max_word_length), intent(in) :: maze_cg_prec

      type(MW_maze_t), intent(inout) :: this

      select case (maze_algorithm)
      case ("inversion")
         this%invert = .true.
      case ("iterative_shake")
         if (n_blocks > 0) then
            call set_blocks_number(this, n_blocks)
            this%block_iterate = .true.
         else
            this%shake = .true.
         end if
      case ("iterative_cg")
         if (maze_cg_store.eq."store") then
            if (maze_cg_prec.eq."preconditioned") then
               this%CG_store_prec = .true.
            else
               this%CG_store_noprec = .true.
            end if
         else
            this%CG_nostore = .true.
         end if
      case default
         call MW_errors_runtime_error("select_algorithm", "mlshake.f90", &
               "Invalid algorithm for mlshake method.")
      end select

   end subroutine select_algorithm

   !================================================================================
   ! Allocate data arrays
   subroutine allocate_arrays(this)
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_runtime_error => runtime_error
      implicit none

      ! Parameters Input/Output
      ! ----------------
      type(MW_maze_t), intent(inout) :: this

      ! Local
      ! -----
      integer :: n_var, n_constr, n_steps, n_blocks, n_red
      integer :: ierr

      n_var = this%n_var
      n_constr = this%n_constr
      n_steps = CONSTR_GRAD_NUM_STEPS
      n_blocks = this%n_blocks
      n_red = this%n_red

      if (this%CG_nostore) then
         allocate(this%const_constr(n_constr), this%ggamma(n_constr), stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
      else
         allocate(this%const_constr(n_constr), this%ggamma(n_constr), this%constr_grad(n_constr, n_constr, n_steps), stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
         this%constr_grad(:,:,:) = 0.0_wp
      end if

      this%const_constr(:) = 0.0_wp
      this%ggamma(:) = 0.0_wp

      if (this%shake) then

         allocate(this%denom(n_constr), stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
         this%denom(:) = 0.0_wp

      else if (this%block_iterate) then

         allocate(this%constr(n_red), this%block_ggamma(n_red), this%Aredinv(n_red, n_red, n_blocks), stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
         this%constr(:) = 0.0_wp
         this%block_ggamma = 0.0_wp
         this%Aredinv(:,:,:) = 0.0_wp

      else if (this%invert) then

         allocate(this%constr(n_constr), this%fconstr(n_constr), this%Ainv(n_constr, n_constr), stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
         this%constr(:) = 0.0_wp
         this%fconstr(:) = 0.0_wp
         this%Ainv(:,:) = 0.0_wp

      else if (this%CG_store_prec.or.this%CG_store_noprec) then

         allocate(this%constr(n_constr), this%fconstr(n_constr), this%Ainv(n_constr,n_constr), this%Aggamma(n_constr), &
               this%Ap(n_constr), this%res(n_constr), this%p(n_constr), this%eta(n_constr), this%jacobi_precond(n_constr), &
               stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
         this%constr(:) = 0.0_wp
         this%fconstr(:) = 0.0_wp
         this%Ainv(:,:) = 0.0_wp
         this%Aggamma(:) = 0.0_wp
         this%Ap(:) = 0.0_wp
         this%res(:) = 0.0_wp
         this%p(:) = 0.0_wp
         this%eta(:) = 0.0_wp
         this%jacobi_precond(:) = 1.0_wp

      else if (this%CG_nostore) then

         allocate(this%constr(n_constr), this%fconstr(n_constr), this%Aggamma(n_constr), this%Ap(n_constr), this%res(n_constr), &
               this%p(n_constr), stat=ierr)
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "maze.f90", ierr)
         end if
         this%constr(:) = 0.0_wp
         this%fconstr(:) = 0.0_wp
         this%Aggamma(:) = 0.0_wp
         this%Ap(:) = 0.0_wp
         this%res(:) = 0.0_wp
         this%p(:) = 0.0_wp

      end if

   end subroutine allocate_arrays

   !================================================================================
   ! Void the data structure
   subroutine void_type(this)
      use MW_errors, only: MW_errors_deallocate_error => deallocate_error
      implicit none

      ! Parameters Input/Output
      ! ----------------
      type(MW_maze_t), intent(inout) :: this

      ! Local
      ! -----
      integer :: ierr

      this%max_iterations = 0
      this%tol = 0.0_wp
      this%last_iteration_count = 0
      this%total_iteration_count = 0.0_sp
      this%last_max_constr = 0.0_wp
      this%n_constr = 0
      this%n_var = 0

      if (allocated(this%const_constr)) then

         if (this%CG_nostore) then
            deallocate(this%const_constr, this%ggamma, stat=ierr)
         else
            deallocate(this%const_constr, this%ggamma, this%constr_grad, stat=ierr)
         end if

         if (this%shake) then

            if (allocated(this%denom)) deallocate(this%denom, stat=ierr)

         else if (this%block_iterate) then

            if (allocated(this%constr)) deallocate(this%constr, this%Aredinv, this%block_ggamma, stat=ierr)

         else if (this%invert) then

            if (allocated(this%constr)) deallocate(this%constr, this%fconstr, this%Ainv, stat=ierr)

         else if (this%CG_store_prec.or.this%CG_store_noprec) then

            if (allocated(this%constr)) deallocate(this%constr, this%fconstr, this%Ainv, this%Aggamma, this%Ap, this%res, this%p, &
                  this%jacobi_precond, stat=ierr)

         else if (this%CG_nostore) then

            if (allocated(this%constr)) deallocate(this%constr, this%fconstr, this%Aggamma, this%Ap, this%res, this%p, stat=ierr)

         end if

         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "maze.f90", ierr)
         end if

      end if

   end subroutine void_type

   subroutine setup_history(this)

      implicit none

      ! Parameters Input/Output
      ! ----------------
      type(MW_maze_t), intent(inout) :: this

      !Locals
      !------
      integer:: i

      do i = 1, CONSTR_GRAD_NUM_STEPS
         this%constr_grad_step(i) = i
      end do

   end subroutine setup_history

   subroutine update_history(this)

      implicit none

      ! Parameters Input/Output
      ! ----------------
      type(MW_maze_t), intent(inout) :: this

      !Locals
      !------
      integer:: step, step_now

      step_now = this%constr_grad_step(CONSTR_GRAD_NUM_STEPS)
      do step = CONSTR_GRAD_NUM_STEPS, 2, -1
         this%constr_grad_step(step) = this%constr_grad_step(step-1)
      end do
      this%constr_grad_step(1) = step_now

   end subroutine update_history

   !================================================================================
   ! Print the data structure parameters
   subroutine print_type(this, ounit)
      implicit none

      ! Parameters Input
      ! ----------------
      type(MW_maze_t), intent(in) :: this
      integer,            intent(in) :: ounit

      write(ounit, '("|maze| maximum number of iterations: ",i12)') this%max_iterations
      write(ounit, '("|maze| Tolerance:               ",es12.5)') this%tol
      if (this%charge_neutrality) then
         write(ounit, '("|maze| Force charge neutrality:      ",a3)') "yes"
      else
         write(ounit, '("|maze| Force charge neutrality:      ",a3)') " no"
      end if
   end subroutine print_type

   !================================================================================
   ! Print the data structure parameters
   subroutine print_statistics(this, ounit)
      implicit none

      ! Parameters Input
      ! ----------------
      type(MW_maze_t), intent(in) :: this
      integer,       intent(in) :: ounit

      write(ounit, '("|maze| number of iterations: ",i12)') this%last_iteration_count
      write(ounit, '("|maze| Maximum of the constraints:        ",es12.5)') this%last_max_constr
   end subroutine print_statistics

   ! ================================================================================
   !Compute starting vector
   subroutine compute_x0(this, system, alpha, x_tmdt, x_t, x_0)
      implicit none

      !Parameters Input
      !----------------
      real(wp), intent(in) :: alpha
      real(wp), intent(in) :: x_t(:), x_tmdt(:)
      type(MW_system_t),  intent(inout) :: system

      ! Parameters Input/Output
      ! -----------------------
      type(MW_maze_t), intent(inout) :: this

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

      !Locals
      integer :: i, k
      integer :: n_var, n_constr
      integer :: prev, chi_tpdt=0, chi_t=0, chi_tmdt=0
      real(wp) :: init_fconstr_i = 0.0_wp

      n_constr = this%n_constr
      n_var = this%n_var
      prev = this%constr_grad_step(2)

      if (this%charge_neutrality) then
         chi_tpdt = system%potential_shift_step(1)
         chi_t    = system%potential_shift_step(2)
         chi_tmdt = system%potential_shift_step(3)
      end if

      if(this%invert.or.this%CG_nostore) then
         do i = 1, n_var
            x_0(i) = x_t(i) + (x_t(i) - x_tmdt(i))*alpha
         end do
         if (this%charge_neutrality) then
            system%chi(chi_tpdt) = 2*system%chi(chi_t) - system%chi(chi_tmdt)
         end if
      else
         do i = 1, n_var
            init_fconstr_i = 0.0_wp
            do k = 1, n_constr
               init_fconstr_i = init_fconstr_i + this%ggamma(k)*this%constr_grad(k,i,prev)
            end do
            x_0(i) = x_t(i) + (x_t(i) - x_tmdt(i))*alpha + init_fconstr_i
         end do
      end if

   end subroutine compute_x0

   subroutine get_shake_matrix(this, system, compute_constr_gradient, mode)
      use MW_fileunit, only: MW_fileunit_get_new_unit => get_new_unit
      use MW_errors, only: MW_errors_open_error => open_error, &
      MW_errors_close_error => close_error, &
      MW_errors_runtime_error => runtime_error
      implicit none

      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t),  intent(inout) :: system

      interface
         subroutine compute_constr_gradient(system, constr_grad)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(out) :: constr_grad(:,:)
         end subroutine compute_constr_gradient
      end interface

      character(len=*), intent(in) :: mode

      integer       :: n_constr, n_var
      integer       :: constr_grad_now, constr_grad_prev
      character(64) :: fname
      logical       :: matrix_data_present

      n_var = this%n_var
      n_constr = this%n_constr
      constr_grad_now = this%constr_grad_step(1)
      constr_grad_prev = this%constr_grad_step(2)

      call  compute_constr_gradient(system, this%constr_grad(1:n_var, 1:n_var, constr_grad_now))
      if (this%charge_neutrality) then
         this%constr_grad(1:n_var , n_constr, constr_grad_now) = 1.0_wp
         this%constr_grad(n_constr,  1:n_var, constr_grad_now) = this%uscale
         this%constr_grad(n_constr, n_constr, constr_grad_now) = 0.0_wp
      end if

      !Gradient is constant during simulation. Setting same values for t and tpdt.
      if (this%firsttime) then
         this%constr_grad(:,:,constr_grad_prev) = this%constr_grad(:,:,constr_grad_now)
         this%firsttime=.false.
      end if

      fname = "maze_matrix.inpt"
      INQUIRE(FILE=fname, EXIST=matrix_data_present)
      if(.not. matrix_data_present) then ! The full matrix has to be computed

         select case (mode)
            case ("symmetric")
               call compute_shake_matrix(this, constr_grad_prev, constr_grad_prev)
               if (this%CG_store_prec) call compute_jacobi_preconditioner(this)
            case ("invert")
               call compute_shake_matrix(this, constr_grad_now, constr_grad_prev)

               !Inverting matrix (The matrix is symmetric)
               call MW_timers_start(TIMER_MINQ_INVMAT)
               call invert_matrix(this%n_constr, this%Ainv(1:n_constr,1:n_constr))
               call MW_timers_stop(TIMER_MINQ_INVMAT)
            case default
               call MW_errors_runtime_error("get_shake_matrix", "maze.f90", &
                     "Invalid mode for getting shake matrix")
         end select

         if (this%constant_constr_grad) then
            if (mode .eq. "symmetric") then
               call write_shake_matrix_ascii(this, system, fname, mode)
            else if (mode .eq. "invert") then
               call write_shake_matrix(this, system, fname)
            end if
         end if

      else !then we read the file and fill this%Ainv array

         if (mode .eq. "symmetric") then
            call read_shake_matrix_ascii(this, system, fname, mode)
            if (this%CG_store_prec) call compute_jacobi_preconditioner(this)
         else if (mode .eq. "invert") then
            call read_shake_matrix(this, system, fname)
         end if

      end if

      if (this%write_shake_matrix) then
         fname = "maze_matrix.out"
         call write_shake_matrix_ascii(this, system, fname, mode)
      end if

   end subroutine get_shake_matrix

   subroutine compute_shake_matrix(this, constr_grad_t_jk, constr_grad_t_jh)
      implicit none

      type(MW_maze_t), intent(inout) :: this
      integer, intent(in) :: constr_grad_t_jk, constr_grad_t_jh

      integer  :: j, h, k, n_constr, n_var
      real(wp) :: A_kh

      n_constr = this%n_constr
      n_var = this%n_var

      !$acc parallel loop private(A_kh)
      do h = 1, n_constr
         do k = 1, n_constr
            A_kh = 0.0_wp
            !$acc loop reduction(+:A_kh)
            do j = 1, n_constr
               A_kh = A_kh + this%constr_grad(j,k,constr_grad_t_jk)*this%constr_grad(j,h,constr_grad_t_jh)
            end do
            !$acc end loop
            this%Ainv(k,h) = A_kh
         end do
      end do
      !$acc end parallel loop

   end subroutine compute_shake_matrix

   subroutine compute_jacobi_preconditioner(this)
      implicit none

      type(MW_maze_t), intent(inout) :: this

      integer :: h

      do h = 1, this%n_constr
         this%jacobi_precond(h) = 1.0_wp/this%Ainv(h,h)
      end do

   end subroutine compute_jacobi_preconditioner

   subroutine write_shake_matrix(this, system, fname)
      use MW_fileunit, only: MW_fileunit_get_new_unit => get_new_unit
      use MW_errors, only: MW_errors_open_error => open_error, &
      MW_errors_close_error => close_error
      implicit none

      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t),  intent(inout) :: system
      character(64),      intent(in)    :: fname

      integer,            parameter     :: n_atoms_check = 10

      integer                           :: n_constr
      integer                           :: funit, ierr

      n_constr = this%n_constr

      ! opening file
      call MW_fileunit_get_new_unit(funit)
      open(unit=funit, file=fname, access="SEQUENTIAL", action="WRITE", &
      position="REWIND", form="UNFORMATTED", status="REPLACE", &
      iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_open_error("write_shake_matrix", "maze.f90", fname, ierr)
      end if

      write(funit) n_atoms_check
      write(funit) system%xyz_atoms(1:n_atoms_check,:)
      write(funit) this%Ainv(:,:)

      ! closing file
      close(funit, iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_close_error("write_shake_matrix", "maze.f90",&
         funit, ierr)
      end if

   end subroutine write_shake_matrix

   subroutine read_shake_matrix(this, system, fname)
      use MW_fileunit, only: MW_fileunit_get_new_unit => get_new_unit
      use MW_errors, only: MW_errors_open_error => open_error, &
      MW_errors_close_error => close_error, &
      MW_errors_runtime_error => runtime_error
      implicit none

      type(MW_maze_t),         intent(inout) :: this
      type(MW_system_t),       intent(inout) :: system
      character(64),           intent(in)    :: fname

      integer,                 parameter     :: n_atoms_check = 10
      real(wp),   dimension(n_atoms_check,3) :: xyz
      real(wp),                parameter     :: tol = 1e-14

      integer                                :: j
      integer                                :: n_constr
      integer                                :: n_atoms_check_read
      integer                                :: funit, ierr

      n_constr = this%n_constr

      call MW_fileunit_get_new_unit(funit)
      open(unit=funit, file=fname, access="SEQUENTIAL", action="READ", &
      position="REWIND", form="UNFORMATTED", status="OLD", iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_open_error("read_shake_matrix", "maze.f90",&
         fname, ierr)
      end if
      read(funit) n_atoms_check_read

      read(funit) xyz
      do j = 1, n_atoms_check

         if (abs(xyz(j,1) - system%xyz_atoms(j,1)) > tol .or. &
               abs(xyz(j,2) - system%xyz_atoms(j,2)) > tol .or. &
               abs(xyz(j,3) - system%xyz_atoms(j,3)) > tol ) then
            call MW_errors_runtime_error("read_shake_matrix", "maze.f90", &
                  "Electrode atoms not at the same positions as the ones used to compute the"// &
                  "inverted matrix that should be loaded => not consistent.")
         end if
      end do
      read(funit) this%Ainv
      close(funit, iostat=ierr)

      if (ierr /= 0) then
         call MW_errors_close_error("read_shake_matrix", "maze.f90", &
         funit, ierr)
      end if
   end subroutine read_shake_matrix

   subroutine write_shake_matrix_ascii(this, system, fname, mode)
      use MW_fileunit, only: MW_fileunit_get_new_unit => get_new_unit
      use MW_errors, only: MW_errors_open_error => open_error, &
      MW_errors_close_error => close_error
      implicit none

      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t),  intent(inout) :: system
      character(64),      intent(in)    :: fname
      character(len=*),   intent(in)    :: mode

      integer,            parameter     :: n_atoms_check = 10

      integer                           :: j
      integer                           :: n_constr
      integer                           :: funit, ierr
      character(64)                     :: matrix_fmt

      n_constr = this%n_constr

      ! opening file
      call MW_fileunit_get_new_unit(funit)
      open(unit=funit, file=fname, access="SEQUENTIAL", action="WRITE", &
      position="REWIND", form="FORMATTED", status="REPLACE", &
      iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_open_error("MW_maze_compute_gradient", "maze.f90", fname, ierr)
      end if

      ! Defining the format string based on the number of constraints
      write (matrix_fmt, '("(", I6, "(es30.20E3,:,1x))")' ) n_constr
      ! writing file
      write(funit, '("# Matrix for electrode charges used with maze algorithm")')
      write(funit, '("# electrode_position_check")')
      write(funit, '("n_atoms_check",(I3))') n_atoms_check
      do j = 1, n_atoms_check
         write(funit,'(3(es29.20, 1x))') system%xyz_atoms(j,1), system%xyz_atoms(j,2), system%xyz_atoms(j,3)
      end do
      write(funit, '("# maze_matrix")')
      write(funit, '("# ",(a10))') mode
      write(funit, '("n_constr ",(I6))') n_constr
      do j = 1, n_constr
         write(funit, matrix_fmt) this%Ainv(:,j)
      end do

      ! closing file
      close(funit, iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_close_error("MW_maze_compute_gradient", "maze.f90",&
         funit, ierr)
      end if

   end subroutine write_shake_matrix_ascii

   subroutine read_shake_matrix_ascii(this, system, fname, mode)
      use MW_fileunit, only: MW_fileunit_get_new_unit => get_new_unit
      use MW_errors, only: MW_errors_open_error => open_error, &
      MW_errors_close_error => close_error, &
      MW_errors_runtime_error => runtime_error
      implicit none

      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t),  intent(inout) :: system
      character(64),      intent(in)    :: fname
      character(len=*),   intent(in)    :: mode

      integer,            parameter     :: n_atoms_check = 10
      real(wp),           parameter     :: tol = 1e-14

      integer                           :: j
      integer                           :: n_constr
      integer                           :: n_atoms_check_read
      integer                           :: funit, ierr
      character(64)                     :: keyword, matrix_fmt
      character(1)                      :: hash_char
      real(wp)                          :: x_read, y_read, z_read

      n_constr = this%n_constr

      call MW_fileunit_get_new_unit(funit)
      open(unit=funit, file=fname, access="SEQUENTIAL", action="READ", &
            position="REWIND", form="FORMATTED", status="OLD", iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_open_error("MW_maze_compute_gradient", "maze.f90",&
               fname, ierr)
      end if
      read(funit, *) hash_char, keyword
      read(funit, *) hash_char, keyword
      if ((keyword /= "electrode_position_check")) then
         call MW_errors_runtime_error("MW_maze_compute_gradient", "maze.f90", &
               "Expecting # electrode_position_check in maze_matrix input file")
      end if
      read(funit, *) keyword, n_atoms_check_read
      if ((keyword /= "n_atoms_check")) then
         call MW_errors_runtime_error("MW_maze_compute_gradient", "maze.f90", &
               "Expecting n_atoms_check in maze_matrix input file")
      end if
      do j = 1, n_atoms_check
         read(funit,'(3(es29.20, 1x))') x_read, y_read, z_read

         if (abs(x_read - system%xyz_atoms(j,1)) > tol .or. &
               abs(y_read - system%xyz_atoms(j,2)) > tol .or. &
               abs(z_read - system%xyz_atoms(j,3)) > tol ) then
            call MW_errors_runtime_error("MW_maze_compute_gradient", "maze.f90", &
                  "Electrode atoms not at the same positions as the ones used to compute the"// &
                  "inverted matrix that should be loaded => not consistent.")
         end if
      end do

      read(funit, *) hash_char, keyword
      if ((keyword /= "maze_matrix")) then
         call MW_errors_runtime_error("MW_maze_compute_gradient", "maze.f90", &
               "Expecting # maze_matrix in maze_matrix input file")
      end if
      read(funit, *) hash_char, keyword
      if ((keyword /= mode)) then
         call MW_errors_runtime_error("MW_maze_compute_gradient", "maze.f90", &
               "Expecting # "//mode//" in maze_matrix input file")
      end if
      read(funit, *) keyword, n_constr
      if ((keyword /= "n_constr")) then
         call MW_errors_runtime_error("MW_maze_compute_gradient", "maze.f90", &
               "Expecting # n_constr in maze_matrix input file")
      end if
      ! Defining the format string based on the number of constraints
      write (matrix_fmt, '("(", I6, "(es30.20E3,:,1x))")' ) n_constr
      do j= 1, n_constr
         read(funit, matrix_fmt)  this%Ainv(:,j)
      end do

      close(funit, iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_close_error("MW_maze_compute_gradient", "maze.f90", &
         funit, ierr)
      end if
   end subroutine read_shake_matrix_ascii

   subroutine solve_invert(this, system, compute_constr_gradient, matvec_product, x)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_stdio, only: MW_stderr
      implicit none

      ! Parameters Input/Output
      ! -----------------------
      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t), intent(inout) :: system  !< electrode parameters

      interface
         subroutine compute_constr_gradient(system, constr_grad)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(out) :: constr_grad(:,:)
         end subroutine compute_constr_gradient
      end interface

      interface
         subroutine matvec_product(system, mat, vec, matvec_prod)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(in) :: system
            real(wp), intent(in) :: mat(:,:)
            real(wp), intent(in) :: vec(:)
            real(wp), intent(out) :: matvec_prod(:)
         end subroutine matvec_product
      end interface

      real(wp), intent(inout) :: x(:)

      !Locals
      !------
      integer  :: n_var, n_constr
      integer  :: prev, now, chi_new=0
      real(wp) :: discr

      n_var = this%n_var
      n_constr = this%n_constr
      now = this%constr_grad_step(1)
      prev = this%constr_grad_step(2)
      if (this%charge_neutrality) chi_new = system%potential_shift_step(1)

      if ((.not.this%constant_constr_grad).or.(this%firsttime)) then
         call get_shake_matrix(this, system, compute_constr_gradient, "invert")
         this%firsttime = .false.
      end if

      !Computing constraints sigma_old
      call matvec_product(system, this%constr_grad(1:n_var,1:n_var,now), x(1:n_var), this%constr(1:n_var))
      this%constr(1:n_var) = this%const_constr(1:n_var) + this%constr(1:n_var)
      if (this%charge_neutrality) then
         this%constr(1:n_var) = this%constr(1:n_var) + system%chi(chi_new)
         this%constr(n_constr) = sum(x(1:n_var)) - this%const_constr(n_constr)
      end if

      !Computing Lagrange multipliers for system without additional constraint
      call matvec_product(system, this%Ainv(1:n_var,1:n_var), -this%constr(1:n_var), this%ggamma(1:n_var))
      if (this%charge_neutrality) then
         this%ggamma(1:n_var) = this%ggamma(1:n_var) -this%constr(n_constr)*this%Ainv(n_constr,1:n_var)
         this%ggamma(n_constr) = -sum(this%Ainv(n_constr,1:n_constr)*this%constr(1:n_constr))
      end if

      !Computing constraint forces
      call matvec_product(system, this%constr_grad(1:n_var,1:n_var,prev), this%ggamma(1:n_var), this%fconstr(1:n_var))
      if (this%charge_neutrality) then
         this%fconstr(1:n_var) = this%fconstr(1:n_var) + this%ggamma(n_constr)
         this%fconstr(n_constr) = this%uscale*this%uscale*sum(this%ggamma(1:n_var))
      end if

      !Updating MaZe variables
      x(1:n_var) = x(1:n_var) + this%fconstr(1:n_var)
      if (this%charge_neutrality) system%chi(chi_new) = system%chi(chi_new) + this%fconstr(n_constr)

      !Computing constraints with updated variables
      call matvec_product(system, this%constr_grad(1:n_var,1:n_var,now), x(1:n_var), this%constr(1:n_var))
      this%constr(1:n_var) = this%const_constr(1:n_var) + this%constr(1:n_var)

      if (this%charge_neutrality) then
         this%constr(1:n_var) = this%constr(1:n_var) + system%chi(chi_new)
         this%constr(n_constr) = sum(x(1:n_var)) - this%const_constr(n_constr)
      end if

      !Checking the residual on the constraint is less than tolerance
      discr = maxval(dabs(this%constr(1:n_constr)))

      !write(*,*)this%constr(1),this%constr(n_constr-1),this%constr(n_constr),discr,(discr.eq.dabs(this%constr(n_constr))),"#M-shake"

      if (discr.lt.this%tol) then
         this%last_max_constr = discr
      else
         this%last_max_constr = discr
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "maze.f90", &
         "M-SHAKE failed to satisfy constraints up the given tolerance")
      end if

   end subroutine solve_invert

   subroutine solve_shake(this, system, compute_constr_gradient, x)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_stdio, only: MW_stderr
      implicit none

      ! Parameters Input/Output
      ! -----------------------
      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t), intent(inout) :: system  !< electrode parameters

      interface
         subroutine compute_constr_gradient(system, constr_grad)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(out) :: constr_grad(:,:)
         end subroutine compute_constr_gradient
      end interface

      real(wp), intent(inout) :: x(:)

      !Locals (common)
      !---------------
      logical  :: maze_converged
      integer  :: j, k
      integer  :: n_var, n_constr
      integer  :: prev, now
      integer  :: discr_k
      real(wp) :: discr

      integer  :: iter
      real(wp) :: gam_k, denom_k, constr_k

      n_var = this%n_var
      n_constr = this%n_constr
      now = this%constr_grad_step(1)
      prev = this%constr_grad_step(2)

      !If dipoles need to be computed, denominator is calculated at each timestep
      if ((.not.this%constant_constr_grad).or.(this%firsttime)) then

         call  compute_constr_gradient(system, this%constr_grad(1:n_var, 1:n_var, now))
         if (this%charge_neutrality) this%constr_grad(:, n_constr, now) = 1.0_wp

         !Gradient is constant during simulation. Setting same values for t and tpdt.
         if (this%firsttime) then
            this%constr_grad(:,:,prev) = this%constr_grad(:,:,now)
            this%firsttime=.false.
         end if

         this%denom(:) = 0.0_wp
         do k = 1, n_constr
            denom_k = 0.0_wp
            do j = 1, k-1
               denom_k       = denom_k       + this%constr_grad(j,k,prev)*this%constr_grad(j,k,now)
               this%denom(j) = this%denom(j) + this%constr_grad(j,k,prev)*this%constr_grad(j,k,now)
            end do
            this%denom(k) = denom_k + this%constr_grad(k,k,prev)*this%constr_grad(k,k,now)
         end do
      end if

      this%ggamma = 0.0_wp

      !Starting iterative procedure
      maze_converged = .false.
      constr_k = 0.0_wp
      do iter = 1, this%max_iterations

         discr   = 0.0_wp
         discr_k = -1

         do k = 1, n_constr !Looping over constraints

            !Computing sigma_k
            if (k.le.n_var) then
               !Computing constraints for maze
               constr_k = this%const_constr(k)
               do j = 1, n_var
                  constr_k = constr_k + this%constr_grad(j,k,now)*x(j)
               end do
            else
               !Computing additional constraint
               constr_k = 0.0_wp
               do j = 1, n_var
                  constr_k = constr_k + x(j)
               end do
            end if

            !Checking if greater than tolerance
            if (dabs(constr_k).gt.this%tol) then
               if (dabs(constr_k).gt.discr) then

                  discr = dabs(constr_k)
                  discr_k = k
               endif

               !Computing part of Lagrange multiplier
               gam_k = -this%sor_param*constr_k/this%denom(k)
               this%ggamma(k) = this%ggamma(k) + gam_k

               !Updating MaZe variables
               do j = 1, n_var
                  x(j) = x(j) + gam_k*this%constr_grad(j,k,prev)
               end do

            end if

         end do !End looping over constraints

         ! write(*,*)iter,discr,discr_k,"#shake"

         !Checking the residual on the constraint is less than tolerance
         if (discr.eq.0.0_wp) then
            this%last_iteration_count = iter
            this%total_iteration_count = this%total_iteration_count + iter
            do k = 1, n_constr
               if (k.le.n_var) then
                  constr_k = this%const_constr(k)
                  do j = 1, n_var
                     constr_k = constr_k + this%constr_grad(j, k, now)*x(j)
                  end do
               else
                  constr_k = 0.0_wp
                  do j = 1, n_var
                     constr_k = constr_k + x(j)
                  end do
               end if
               if (dabs(constr_k).gt.discr) discr = dabs(constr_k)
            end do
            this%last_max_constr = discr
            maze_converged = .true.
            exit
         end if

      end do !End of iterative procedure

      if (.not.maze_converged) then
         this%last_iteration_count = iter
         this%total_iteration_count = this%total_iteration_count + iter
         discr = 0.0_wp
         do k = 1, n_constr
            if (k.le.n_var) then
               constr_k = this%const_constr(k)
               do j = 1, n_var
                  constr_k = constr_k + this%constr_grad(j, k, now)*x(j)
               end do
            else
               constr_k = 0.0_wp
               do j = 1, n_var
                  constr_k = constr_k + x(j)
               end do
            end if
            if (dabs(constr_k).gt.discr) discr = dabs(constr_k)
         end do
         this%last_max_constr = discr
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "maze.f90", "MaZe failed to converge")
      end if

   end subroutine solve_shake

   subroutine solve_block_iterate(this, system, compute_constr_gradient, x)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_stdio, only: MW_stderr
      implicit none

      ! Parameters Input/Output
      ! -----------------------
      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t), intent(inout) :: system  !< electrode parameters

      interface
         subroutine compute_constr_gradient(system, constr_grad)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(out) :: constr_grad(:,:)
         end subroutine compute_constr_gradient
      end interface

      real(wp), intent(inout) :: x(:)

      !Locals (common)
      !---------------
      logical  :: mlbshake_converged
      integer  :: j, k, kk, h, hh, iblock, iter
      integer  :: n_var, n_red, n_blocks
      integer  :: prev, now
      integer  :: discr_k
      real(wp) :: discr
      real(wp) :: gam_kk, constr_kk, x_j

      n_var = this%n_var
      n_red = this%n_red
      n_blocks = this%n_blocks
      now = this%constr_grad_step(1)
      prev = this%constr_grad_step(2)

      !If dipoles need to be computed, denominator is calculated at each timestep
      if ((.not.this%constant_constr_grad).or.(this%firsttime)) then

         call  compute_constr_gradient(system, this%constr_grad(:,:, now))

         !Gradient is constant during simulation. Setting same values for t and tpdt.
         if (this%firsttime) then
            this%constr_grad(:, :, prev) = this%constr_grad(:, :, now)
            this%firsttime=.false.
         end if

         do iblock = 1, n_blocks
            hh = 1
            do h = (iblock-1)*n_red+1, (iblock)*n_red
               kk = 1
               do k = (iblock-1)*n_red+1, (iblock)*n_red
                  this%Aredinv(kk, hh, iblock) = 0.0_wp
                  do j = 1, n_var
                     this%Aredinv(kk, hh, iblock) = this%Aredinv(kk, hh, iblock) + &
                           this%constr_grad(j, h, now)*this%constr_grad(j, k, prev) !Inverted order of k and h for efficiency reasons.
                  end do
                  kk = kk + 1
               end do
               hh = hh + 1
            end do

            call invert_matrix(this%n_red, this%Aredinv(:,:,iblock))

         end do
      end if

      this%ggamma(:) = 0.0_wp

      mlbshake_converged = .false.
      do iter = 1, this%max_iterations !Loop over iterations

         discr = 0.0_wp
         discr_k = -1

         do iblock = 1, n_blocks !Loop over blocks

            kk = 1
            do k = (iblock-1)*n_red+1, (iblock)*n_red !Loop over constraints inside block

               constr_kk = this%const_constr(k)
               do j = 1, n_var
                  constr_kk = constr_kk + this%constr_grad(j, k, now)*x(j)
               end do

               if (discr .lt. dabs(constr_kk)) then
                  discr = dabs(constr_kk)
                  discr_k = k
               end if

               this%constr(kk) = constr_kk
               kk = kk + 1
            end do

            if (discr .gt. this%tol) then !If maximum constraint inside block is greater than tolerance than solve reduced system

               kk = 1
               do k = (iblock-1)*n_red+1, (iblock)*n_red !Computing lagrange multipliers for reduced system

                  gam_kk = 0.0_wp

                  hh = 1
                  do h = (iblock-1)*n_red+1, (iblock)*n_red
                     gam_kk = gam_kk - this%Aredinv(hh, kk, iblock)*this%constr(hh) !A_hk is in fact A_kh
                     hh = hh + 1
                  end do
                  gam_kk = this%sor_param*gam_kk
                  this%ggamma(k) = this%ggamma(k) + gam_kk
                  this%block_ggamma(kk) = gam_kk

                  kk = kk + 1
               end do

               do j = 1, n_var !Computing new "coordinates" for reduced system

                  x_j = x(j)

                  kk = 1
                  do k = (iblock-1)*n_red+1, (iblock)*n_red

                     x_j = x_j + this%block_ggamma(kk)*this%constr_grad(k, j, prev)

                     kk = kk + 1
                  end do
                  x(j) = x_j

               end do

            end if

         end do !End loop over blocks

         ! write(*,*) iter, discr, "#B-shake"

         if (discr.lt.this%tol) then
            this%last_iteration_count = iter
            this%total_iteration_count = this%total_iteration_count + iter
            discr = 0.0_wp
            do k = 1, n_var
               constr_kk = this%const_constr(k)
               do j = 1, n_var
                  constr_kk = constr_kk + this%constr_grad(j, k, now)*x(j)
               end do
               if ((constr_kk).gt.discr) discr = dabs(constr_kk)
            end do
            this%last_max_constr = discr
            mlbshake_converged = .true.
            exit
         end if

      end do !End of iterative loop

      if (.not.mlbshake_converged) then
         this%last_iteration_count = iter
         this%total_iteration_count = this%total_iteration_count + iter
         discr = 0.0_wp
         do k = 1, n_var
            constr_kk = this%const_constr(k)
            do j = 1, n_var
               constr_kk = constr_kk + this%constr_grad(j, k, now)*x(j)
            end do
            if (dabs(constr_kk).gt.discr) discr = dabs(constr_kk)
         end do
         this%last_max_constr = discr
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "maze.f90", "MaZe failed to converge")
      end if

   end subroutine solve_block_iterate

   subroutine solve_CG_store(this, system, compute_constr_gradient, matvec_product, inner_product, x)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_stdio, only: MW_stderr
      implicit none

      ! Parameters Input/Output
      ! -----------------------
      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t), intent(inout) :: system  !< electrode parameters

      interface
         subroutine compute_constr_gradient(system, constr_grad)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(out) :: constr_grad(:,:)
         end subroutine compute_constr_gradient
      end interface

      interface
         subroutine matvec_product(system, mat, vec, matvec_prod)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(in) :: system
            real(wp), intent(in) :: mat(:,:)
            real(wp), intent(in) :: vec(:)
            real(wp), intent(out) :: matvec_prod(:)
         end subroutine matvec_product
      end interface

      interface
         subroutine inner_product(system, vec_v, vec_w, inner_prod)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(in) :: system
            real(wp), intent(in) :: vec_v(:)
            real(wp), intent(in) :: vec_w(:)
            real(wp), intent(out) :: inner_prod
         end subroutine inner_product
      end interface

      real(wp), intent(inout) :: x(:)

      !Locals (common)
      !---------------
      integer  :: n_constr, n_var
      integer  :: iter = 0
      integer  :: prev, now
      real(wp) :: discr, res_norm
      real(wp) :: res_dot_eta_new, res_dot_eta_old, pAp
      real(wp) :: alpha, beta

      n_var = this%n_var
      n_constr = this%n_constr
      now = this%constr_grad_step(1)
      prev = this%constr_grad_step(2)

      if ((.not.this%constant_constr_grad).or.(this%firsttime)) then

         call get_shake_matrix(this, system, compute_constr_gradient, "symmetric")
         this%firsttime = .false.
      end if

      call matvec_product(system, this%constr_grad(1:n_var,1:n_var,now), x(1:n_var), this%constr(1:n_var))
      this%constr(1:n_var) = this%const_constr(1:n_var) + this%constr(1:n_var)
      if (this%charge_neutrality) this%constr(n_constr) = sum(x)

      discr = maxval(abs(this%constr))

      call matvec_product(system, this%Ainv(1:n_var,1:n_var), this%ggamma(1:n_var), this%Aggamma(1:n_var))
      if (this%charge_neutrality) this%Aggamma(n_constr) = 1.0_wp !TO BE IMPLEMENTED

      this%res(:) = this%Aggamma(:) + this%constr(:)
      this%eta(:) = this%jacobi_precond(:)*this%res(:)
      this%p(:)   = -this%eta(:)

      call inner_product(system, this%res(1:n_var), this%eta(1:n_var), res_dot_eta_old)
      if (this%charge_neutrality) res_dot_eta_old = res_dot_eta_old + 1.0_wp !TO BE IMPLEMENTED
      res_norm = maxval(abs(this%res))

      if (res_norm.gt.this%tol) then

         do iter = 1, this%max_iterations

            call matvec_product(system, this%Ainv(1:n_var,1:n_var), this%p(1:n_var), this%Ap(1:n_var))
            if (this%charge_neutrality) this%Ap(n_constr) = 1.0_wp !TO BE IMPLEMENTED
            call inner_product(system, this%p(1:n_var), this%Ap(1:n_var), pAp)
            if (this%charge_neutrality) pAp = pAp + 1.0_wp !TO BE IMPLEMENTED
            alpha = res_dot_eta_old/pAp

            this%ggamma(:) = this%ggamma(:) + alpha*this%p(:)

            if (mod(iter, 100) .eq. 0) then
               call matvec_product(system, this%Ainv(1:n_var,1:n_var), this%ggamma(1:n_var), this%Aggamma(1:n_var))
               if (this%charge_neutrality) this%Aggamma(n_constr) = 1.0_wp !TO BE IMPLEMENTED
               this%res(:) = this%Aggamma(:) + this%constr(:)
            else
               this%res(:) = this%res(:) + alpha*this%Ap(:)
            end if

            this%eta(:) = this%jacobi_precond(:)*this%res(:)

            call inner_product(system, this%res(1:n_var), this%eta(1:n_var), res_dot_eta_new)
            if (this%charge_neutrality) res_dot_eta_new = res_dot_eta_new + 1.0_wp !TO BE IMPLEMENTED
            res_norm = maxval(abs(this%res))

            ! write(*,*) iter, res_norm, "#CG-shake-store"

            if (res_norm.lt.this%tol) exit

            beta = res_dot_eta_new/res_dot_eta_old
            res_dot_eta_old = res_dot_eta_new

            this%p(:) = -this%eta(:) + beta*this%p(:)

         end do

      end if

      call matvec_product(system, this%constr_grad(1:n_var,1:n_var,prev), this%ggamma(1:n_var), this%fconstr(1:n_var))
      if (this%charge_neutrality) this%fconstr(:) = this%fconstr(:) + this%constr_grad(:,n_constr,prev)*this%ggamma(n_constr)

      x(:) = x(:) + this%fconstr(:)

      call matvec_product(system, this%constr_grad(1:n_var,1:n_var,now), x(1:n_var), this%constr(1:n_var))
      this%constr(1:n_var) = this%const_constr(1:n_var) + this%constr(1:n_var)
      if (this%charge_neutrality) this%constr(n_constr) = sum(x)

      discr = maxval(abs(this%constr))

      ! write(*,*) !discr, "#CG-shake-store discr"

      if (discr.lt.this%tol) then
         this%last_iteration_count = iter
         this%total_iteration_count = this%total_iteration_count + iter
         this%last_max_constr = discr
      else
         this%last_iteration_count = iter
         this%total_iteration_count = this%total_iteration_count + iter
         this%last_max_constr = discr
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "maze.f90", "MaZe failed to converge")
      end if

   end subroutine solve_CG_store

   subroutine solve_CG_nostore(this, system, apply_constr_gradient, inner_product, x)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_stdio, only: MW_stderr
      implicit none

      ! Parameters Input/Output
      ! -----------------------
      type(MW_maze_t), intent(inout) :: this
      type(MW_system_t), intent(inout) :: system  !< electrode parameters

      interface
         subroutine apply_constr_gradient(system, x, y)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(in) :: x(:)
            real(wp), intent(out) :: y(:)
         end subroutine apply_constr_gradient
      end interface

      interface
         subroutine inner_product(system, vec_v, vec_w, inner_prod)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(in) :: system
            real(wp), intent(in) :: vec_v(:)
            real(wp), intent(in) :: vec_w(:)
            real(wp), intent(out) :: inner_prod
         end subroutine inner_product
      end interface

      real(wp), intent(inout) :: x(:)

      !Locals (common)
      !---------------
      integer  :: n_constr, n_var
      integer  :: iter = 0
      integer  :: prev, now
      real(wp) :: discr, res_norm
      real(wp) :: res_dot_res_new, res_dot_res_old, pAp
      real(wp) :: alpha, beta

      n_var = this%n_var
      n_constr = this%n_constr
      now = this%constr_grad_step(1)
      prev = this%constr_grad_step(2)

      call apply_constr_gradient(system, x, this%constr(1:n_var))
      ! call matvec_product(system, this%constr_grad(1:n_var,1:n_var,now), x(1:n_var), this%constr(1:n_var))
      this%constr(1:n_var) = this%const_constr(1:n_var) + this%constr(1:n_var)
      if (this%charge_neutrality) this%constr(n_constr) = sum(x)

      discr = maxval(abs(this%constr))

      call apply_constr_gradient(system, this%ggamma(1:n_var), this%Aggamma(1:n_var))
      call apply_constr_gradient(system, this%Aggamma(1:n_var), this%Aggamma(1:n_var))
      ! call matvec_product(system, this%Ainv(1:n_var,1:n_var), this%ggamma(1:n_var), this%Aggamma(1:n_var))
      if (this%charge_neutrality) this%Aggamma(n_constr) = 1.0_wp !TO BE IMPLEMENTED

      this%res(:) = this%Aggamma(:) + this%constr(:)
      this%p(:)   = -this%res(:)

      call inner_product(system, this%res(1:n_var), this%res(1:n_var), res_dot_res_old)
      if (this%charge_neutrality) res_dot_res_old = res_dot_res_old + 1.0_wp !TO BE IMPLEMENTED
      res_norm = maxval(abs(this%res))

      if (res_norm.gt.this%tol) then

         do iter = 1, this%max_iterations

            call apply_constr_gradient(system, this%p(1:n_var), this%Ap(1:n_var))
            call apply_constr_gradient(system, this%Ap(1:n_var), this%Ap(1:n_var))
            ! call matvec_product(system, this%Ainv(1:n_var,1:n_var), this%p(1:n_var), this%Ap(1:n_var))
            if (this%charge_neutrality) this%Ap(n_constr) = 1.0_wp !TO BE IMPLEMENTED
            call inner_product(system, this%p(1:n_var), this%Ap(1:n_var), pAp)
            if (this%charge_neutrality) pAp = pAp + 1.0_wp !TO BE IMPLEMENTED
            alpha = res_dot_res_old/pAp

            this%ggamma(:) = this%ggamma(:) + alpha*this%p(:)

            if (mod(iter, 100) .eq. 0) then
               call apply_constr_gradient(system, this%ggamma(1:n_var), this%Aggamma(1:n_var))
               call apply_constr_gradient(system, this%Aggamma(1:n_var), this%Aggamma(1:n_var))
               ! call matvec_product(system, this%Ainv(1:n_var,1:n_var), this%ggamma(1:n_var), this%Aggamma(1:n_var))
               if (this%charge_neutrality) this%Aggamma(n_constr) = 1.0_wp !TO BE IMPLEMENTED
               this%res(:) = this%Aggamma(:) + this%constr(:)
            else
               this%res(:) = this%res(:) + alpha*this%Ap(:)
            end if

            call inner_product(system, this%res(1:n_var), this%res(1:n_var), res_dot_res_new)
            if (this%charge_neutrality) res_dot_res_new = res_dot_res_new + 1.0_wp !TO BE IMPLEMENTED
            res_norm = maxval(abs(this%res))

            ! write(*,*) iter, res_norm, "#CG-shake-nostore"

            if (res_norm.lt.this%tol) exit

            beta = res_dot_res_new/res_dot_res_old
            res_dot_res_old = res_dot_res_new

            this%p(:) = -this%res(:) + beta*this%p(:)

         end do

      end if

      call apply_constr_gradient(system, this%ggamma(1:n_var), this%fconstr(1:n_var))
      ! call matvec_product(system, this%constr_grad(1:n_var,1:n_var,prev), this%ggamma(1:n_var), this%fconstr(1:n_var))
      if (this%charge_neutrality) this%fconstr(:) = this%fconstr(:) + this%constr_grad(:,n_constr,prev)*this%ggamma(n_constr)

      x(:) = x(:) + this%fconstr(:)

      call apply_constr_gradient(system, x(1:n_var), this%constr(1:n_var))
      ! call matvec_product(system, this%constr_grad(1:n_var,1:n_var,now), x(1:n_var), this%constr(1:n_var))
      this%constr(1:n_var) = this%const_constr(1:n_var) + this%constr(1:n_var)
      if (this%charge_neutrality) this%constr(n_constr) = sum(x)

      discr = maxval(abs(this%constr))

      ! write(*,*) !discr, "#CG-shake-nostore discr"

      if (discr.lt.this%tol) then
         this%last_iteration_count = iter
         this%total_iteration_count = this%total_iteration_count + iter
         this%last_max_constr = discr
      else
         this%last_iteration_count = iter
         this%total_iteration_count = this%total_iteration_count + iter
         this%last_max_constr = discr
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "maze.f90", "MaZe failed to converge")
      end if

   end subroutine solve_CG_nostore

   ! ================================================================================
   include 'matrix_inversion.inc'

end module MW_maze
