! Module defines matrix inversion type to store matrix inversion algorithm parameters
module MW_matrix_inversion
   use MW_timers
   use MW_kinds, only: wp, sp
   use MW_system, only: MW_system_t

   implicit none
   private

   ! Public type
   ! -----------
   public :: MW_matrix_inversion_t

   ! Public subroutines
   ! ------------------
   public :: set_tolerance
   public :: set_charge_neutrality
   public :: set_write_matrix
   public :: reset_statistics
   public :: allocate_arrays
   public :: void_type
   public :: print_type
   public :: print_statistics
   public :: solve

   type MW_matrix_inversion_t

      !Convergence parameters
      real(wp) :: tol           !< tolerance on residual norm

      !Logicals
      logical :: charge_neutrality = .false.
      logical :: constant_matrix = .false.
      logical :: firsttime = .true.
      logical :: write_hessian_matrix = .false.
      logical :: calc_elec = .false.
      logical :: calc_dip = .false.

      !Statistics
      real(wp) :: res_norm         !< residual norm after procedure

      !Work arrays
      integer :: n                  !< dimensions of the work arrays
      real(wp), allocatable :: b(:)    !< Constant vector b
      real(wp), allocatable :: A(:,:)           !< Hessian
      real(wp), allocatable :: Ainv(:,:)          !< Inverse of the Hessian
      real(wp), allocatable :: Ax(:)   !< vector to check the exactness of the solution
      real(wp), allocatable :: q0(:)    !< Constant vector q0

   end type MW_matrix_inversion_t

contains

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

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

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

   ! ================================================================================
   ! reset statistic counters
   subroutine reset_statistics(this)
      implicit none
      type(MW_matrix_inversion_t), intent(inout) :: this
   end subroutine reset_statistics

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

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

      ! Parameters Input
      ! -------------
      integer, intent(in) :: n !< size of the data arrays

      ! Local
      ! -----
      integer :: i, j
      integer :: ierr

      this%n = n
      allocate(this%b(n), this%q0(n), this%A(n,n), this%Ainv(n,n), this%Ax(n), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("allocate_arrays", "matrix_inversion.F90", ierr)
      end if

      do i = 1, this%n
         this%b(i) = 0.0_wp
         this%q0(i) = 0.0_wp
         do j = 1, this%n
            this%A(j,i) = 0.0_wp
            this%Ainv(j,i) = 0.0_wp
         end do
         this%Ax(i) = 0.0_wp
      end do

   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_matrix_inversion_t), intent(inout) :: this

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

      this%tol = 0.0_wp
      this%res_norm = 0.0_wp
      this%n = 0

      if (allocated(this%b)) then
         deallocate(this%b, this%q0, this%A, this%Ainv, this%Ax, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "matrix_inversion.F90", ierr)
         end if
      end if
   end subroutine void_type

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

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

      write(ounit, '("|matrix_inversion| Tolerance:               ",es12.5)') this%tol
      if (this%charge_neutrality) then
         write(ounit, '("|matrix_inversion| Force charge neutrality:      ",a3)') "yes"
      else
         write(ounit, '("|matrix_inversion| 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_matrix_inversion_t), intent(in) :: this
      integer,       intent(in) :: ounit

      write(ounit, '("|matrix_inversion| Norm of the residual after procedure:        ",es12.5)') this%res_norm
   end subroutine print_statistics

   subroutine get_hessian_inverse(this, system, compute_hessian)
      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_matrix_inversion_t), intent(inout) :: this
      type(MW_system_t),  intent(inout) :: system

      integer  :: xyz_now
      integer  :: n
      character(64) :: fname
      logical  :: matrix_data_present

      interface
         subroutine compute_hessian(system, hessian)
            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) :: hessian(:,:)
         end subroutine compute_hessian
      end interface

      n = this%n
      xyz_now = system%xyz_ions_step(1)

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

         call compute_hessian(system, this%A)

         !Inverting matrix (The matrix is symmetric)
         if (.not.(system%elecrun.and.system%pimrun)) then
            call MW_timers_start(TIMER_MINQ_INVMAT)
            call invert_matrix(this%n, this%Ainv, this%A)
            call MW_timers_stop(TIMER_MINQ_INVMAT)
         end if

         if (this%constant_matrix) call write_hessian_matrix(this, system, fname)

         if (system%elecrun.and.system%pimrun) then
            call MW_errors_runtime_error("get_hessian_inverse", "matrix_inversion_f90", &
            "Stopping execution after matrix writing: full matrix inversion method not implemented for dipoles + electrodes")
         end if

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

         if (system%elecrun.and.system%pimrun) then
            call MW_errors_runtime_error("get_hessian_inverse", "matrix_inversion_f90", &
            "Stopping execution before matrix reading: full matrix inversion method not implemented for dipoles + electrodes")
         end if

         call read_hessian_matrix(this, system, fname)

      end if

      if (this%write_hessian_matrix) then
         fname = "hessian_matrix.out"
         call write_hessian_matrix_ascii(this, system, fname)
      end if

   end subroutine get_hessian_inverse

   subroutine write_hessian_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_matrix_inversion_t), intent(inout) :: this
      type(MW_system_t),           intent(inout) :: system
      character(64),               intent(in)    :: fname

      integer,                     parameter     :: n_atoms_check = 10

      integer                                    :: funit, ierr

      ! 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("MW_matrix_inversion_get_A", "matrix_inversion_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("MW_matrix_inversion_get_A", "matrix_inversion_f90",&
         funit, ierr)
      end if

   end subroutine write_hessian_matrix

   subroutine read_hessian_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_matrix_inversion_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_atoms_check_read
      integer                                    :: funit, ierr


      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("MW_matrix_inversion_get_A", "matrix_inversion_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("MW_matrix_inversion_get_A", "matrix_inversion_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("MW_matrix_inversion_get_A", "matrix_inversion_f90", &
         funit, ierr)
      end if
   end subroutine read_hessian_matrix

   subroutine write_hessian_matrix_ascii(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_matrix_inversion_t), intent(inout) :: this
      type(MW_system_t),           intent(inout) :: system
      character(64),               intent(in)    :: fname

      integer,                     parameter     :: n_atoms_check = 10

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

      n = this%n

      ! 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_matrix_inversion_get_A", "matrix_inversion_f90", fname, ierr)
      end if

      ! Defining the format string based on the number of constraints
      write (matrix_fmt, '("(", I6, "(es23.15E3,:,1x))")' ) n
      ! writing file
      write(funit, '("# Inverted Matrix for electrode charges used with matrix_inversion_algorithm")')
      write(funit, '("# electrode_position_check")')
      write(funit, '("n_atoms_check",(I3))') n_atoms_check
      do j = 1, n_atoms_check
         write(funit,'(3(es22.15, 1x))') system%xyz_atoms(j,1), system%xyz_atoms(j,2), system%xyz_atoms(j,3)
      end do
      write(funit, '("# matrix")')
      write(funit, '("n ",(I6))') n
      do j = 1, n
         write(funit, matrix_fmt)  this%A(:,j)
      end do
      if (.not.(system%elecrun.and.system%pimrun)) then
         do j = 1, n
            write(funit, matrix_fmt)  this%Ainv(:,j)
         end do
      end if

      ! closing file
      close(funit, iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_close_error("MW_matrix_inversion_get_A", "matrix_inversion_f90",&
         funit, ierr)
      end if

   end subroutine write_hessian_matrix_ascii

   subroutine read_hessian_matrix_ascii(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_matrix_inversion_t), intent(inout) :: this
      type(MW_system_t),           intent(inout) :: system
      character(64),               intent(in)    :: fname

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

      integer                                    :: j
      integer                                    :: n
      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 = this%n

      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_matrix_inversion_get_A", "matrix_inversion_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_matrix_inversion_get_A", "matrix_inversion_f90", &
         "Expecting # electrode_position_check in matrix_inversion_matrix input file")
      end if
      read(funit, *) keyword, n_atoms_check_read
      if ((keyword /= "n_atoms_check")) then
         call MW_errors_runtime_error("MW_matrix_inversion_get_A", "matrix_inversion_f90", &
         "Expecting n_atoms_check in matrix_inversion_matrix input file")
      end if
      do j = 1, n_atoms_check
         read(funit,'(3(es22.15, 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_matrix_inversion_get_A", "matrix_inversion_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 /= "matrix")) then
         call MW_errors_runtime_error("MW_matrix_inversion_get_A", "matrix_inversion_f90", &
         "Expecting # matrix in matrix_inversion_matrix input file")
      end if
      read(funit, *) keyword, n
      if ((keyword /= "n")) then
         call MW_errors_runtime_error("MW_matrix_inversion_get_A", "matrix_inversion_f90", &
         "Expecting # n in matrix_inversion_matrix input file")
      end if
      ! Defining the format string based on the number of constraints
      write (matrix_fmt, '("(", I6, "(es23.15E3,:,1x))")' ) n
      do j = 1, n
         read(funit, matrix_fmt)  this%A(:,j)
      end do
      do j = 1, n
         read(funit, matrix_fmt)  this%Ainv(:,j)
      end do

      close(funit, iostat=ierr)
      if (ierr /= 0) then
         call MW_errors_close_error("MW_matrix_inversion_get_A", "matrix_inversion_f90", &
         funit, ierr)
      end if
   end subroutine read_hessian_matrix_ascii

   ! ================================================================================
   subroutine solve(this, system, compute_hessian, matvec_product, x)
      use MW_errors, only: MW_errors_runtime_error => runtime_error, &
                           MW_errors_allocate_error => allocate_error, &
                           MW_errors_deallocate_error => deallocate_error
      implicit none

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

      interface
         subroutine compute_hessian(system, hessian)
            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) :: hessian(:,:)
         end subroutine compute_hessian
      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 (common)
      !---------------
      integer :: n, ierr, i, j
      real(wp) :: EAinvE, jAinvE, target_charge
      real(wp), allocatable :: AinvE(:), E(:)

      n = this%n

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

         call get_hessian_inverse(this, system, compute_hessian)
         this%firsttime = .false.

         if (this%charge_neutrality) then
            target_charge = - system%total_ions_charge

            allocate(E(n), AinvE(n), stat=ierr)
            if (ierr /= 0) then
               call MW_errors_allocate_error("solve", "matrix_inversion.F90", ierr)
            end if

            E(:) = 1.0_wp

            ! compute S = Ainv - (Ainv * E * Et * Ainv) / (Et * Ainv * E)
            ! and replace it in Ainv
            call matvec_product(system, this%Ainv(:,:), E(:), AinvE(:))

            EAinvE = sum(AinvE(:))

            do j = 1, n
               jAinvE = AinvE(j)
               do i = 1, n
                  this%Ainv(i, j) = this%Ainv(i, j) - AinvE(i) * jAinvE / EAinvE
                  this%q0(i) = AinvE(i) * target_charge / EAinvE
               end do
            end do

            deallocate(E, AinvE, stat=ierr)
            if (ierr /= 0) then
               call MW_errors_deallocate_error("solve", "matrix_inversion.F90", ierr)
            end if

         end if

      end if

      call matvec_product(system, this%Ainv(:,:), this%b(:), x(:))

      if (system%global_charge_neutrality) then
         do i = 1, n
            x(i) = x(i) + this%q0(i)
         end do
      end if

      !call matvec_product(system, this%A(:,:), x(:), this%Ax(:))

      !Checking the residual on the constraint is less than tolerance
      !discr = maxval(dabs(this%Ax - this%b))
      ! write(*,*) discr

      !if (discr.gt.this%tol) then
      !   call MW_errors_runtime_error("solve", "matrix_inversion.F90", &
      !   "matrix inversion failed to satisfy constraints up the given tolerance")
      !end if

   end subroutine solve

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

end module MW_matrix_inversion
