!> @file
!!   Basic routines for calling NTPoly's solver routines.
!! @author
!!   Copyright (C) 2016 CheSS developers
!!
!!   This file is part of CheSS.
!!   
!!   CheSS is free software: you can redistribute it and/or modify
!!   it under the terms of the GNU Lesser General Public License as published by
!!   the Free Software Foundation, either version 3 of the License, or
!!   (at your option) any later version.
!!   
!!   CheSS is distributed in the hope that it will be useful,
!!   but WITHOUT ANY WARRANTY; without even the implied warranty of
!!   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!!   GNU Lesser General Public License for more details.
!!   
!!   You should have received a copy of the GNU Lesser General Public License
!!   along with CheSS.  If not, see <http://www.gnu.org/licenses/>.


!> Basic routines for calling NTPoly's solver routines.
module bigpoly_solvers
  use futile
  use datatypesmodule, only : ntreal
  use densitymatrixsolversmodule, only : trs4
  use inversesolversmodule, only : invert
  use permutationmodule, only : permutation_t, constructrandompermutation, &
     destructpermutation
  use processgridmodule, only : constructprocessgrid, destructprocessgrid
  use psmatrixmodule, only : matrix_ps, destructmatrix, constructemptymatrix, &
     fillmatrixidentity, copymatrix
  use psmatrixalgebramodule, only : scalematrix
  use solverparametersmodule, only : solverparameters_t, &
     destructsolverparameters
  use squarerootsolversmodule, only : inversesquareroot, squareroot
  use sparsematrix_base, only : sparse_matrix, matrices
  use sparsematrix_types, only : mp
  use sparsematrix_timing
  use bigpoly, only : bigpoly_options
  implicit none
  private

  !> Public routines
  public :: ntpoly_solve
  public :: ntpoly_overlappowers

  contains

    !> Convert the required powers of the overlap matrix (-1/2, 1/2, -1)
    subroutine ntpoly_overlappowers(mat, val, out_mat, comm, &
        inv_val, sq_val, isq_val)
      !> The matrix to invert's sparsity pattern.
      type(sparse_matrix), intent(in):: mat
      !> The matrix values.
      type(matrices), intent(in) :: val
      !> Output matrix sparsity pattern.
      type(sparse_matrix), intent(inout):: out_mat
      !> The communicator to distribute on.
      integer, intent(inout) :: comm
      !> Output inverted values.
      type(matrices), intent(inout), optional :: inv_val
      !> Output square root values.
      type(matrices), intent(inout), optional :: sq_val
      !> Output inverse square root  values.
      type(matrices), intent(inout), optional :: isq_val
      ! Local variables
      type(matrix_ps) :: nt_mat
      type(matrix_ps) :: nt_inv
      type(matrix_ps) :: nt_sq
      type(matrix_ps) :: nt_isq
      type(solverparameters_t) :: param
      type(permutation_t) :: perm
      real(f_double) :: threshold, convergence

      call f_routine(id='ntpoly_ovlp_powers')

      ! Create the process grid
      call constructprocessgrid(comm)

      ! Conversion of input matrix.
      call chess_to_ntpoly(mat, val, nt_mat)

      ! Load Balancing.
      call constructrandompermutation(perm, nt_mat%logical_matrix_dimension)

      ! Conversion of parameters.
      param = solverparameters_t()
      threshold = bigpoly_options//'threshold_overlap'
      param%threshold = real(threshold, ntreal)
      convergence = bigpoly_options//'convergence_overlap'
      param%converge_diff = real(convergence, ntreal)
      param%BalancePermutation = perm

      call f_timing(TCAT_NTPOLY_INV, 'ON')
      if (present(inv_val)) then
         call invert(nt_mat, nt_inv, param)
         call ntpoly_to_chess(nt_inv, out_mat, inv_val)
      end if
      if (present(sq_val)) then
         call squareroot(nt_mat, nt_sq, param)
         call ntpoly_to_chess(nt_sq, out_mat, sq_val)
      end if
      if (present(isq_val)) then
         call inversesquareroot(nt_mat, nt_isq, param)
         call ntpoly_to_chess(nt_isq, out_mat, isq_val)
      end if
      call f_timing(TCAT_NTPOLY_INV, 'OFF')

      ! Cleanup
      call destructmatrix(nt_mat)
      call destructmatrix(nt_inv)
      call destructmatrix(nt_sq)
      call destructmatrix(nt_isq)
      call destructpermutation(perm)
      call destructsolverparameters(param)
      call destructprocessgrid()

      call f_release_routine()
    end subroutine ntpoly_overlappowers

    !> Solve for the density matrix using ntpoly.
    subroutine ntpoly_solve(hmat, hval, smat, sval, dmat, dval, &
        isq, nel, ebs, invert_overlap, comm)
      !> The Hamiltonian matrix.
      type(sparse_matrix), intent(in):: hmat
      !> The Hamiltonian values.
      type(matrices), intent(in) :: hval
      !> The overlap matrix.
      type(sparse_matrix), intent(inout):: smat
      !> The overlap values.
      type(matrices), intent(in) :: sval
      !> Output density matrix.
      type(sparse_matrix), intent(inout):: dmat
      !> Output density values.
      type(matrices), intent(inout) :: dval
      !> The inverse square root of the the overlap matrix (cache).
      type(matrices), intent(inout) :: isq
      !> The number of electrons.
      integer, intent(in) :: nel
      !> The energy value
      real(kind=mp),intent(out) :: ebs
      !> Whether we need to invert the overlap matrix.
      logical, intent(in) :: invert_overlap
      !> The communicator to distribute on.
      integer, intent(inout) :: comm
      ! Local variables
      type(matrix_ps) :: nt_hmat
      type(matrix_ps) :: nt_smat
      type(matrix_ps) :: nt_isq
      type(matrix_ps) :: nt_dmat
      type(solverparameters_t) :: param
      type(permutation_t) :: perm
      real(f_double) :: threshold, convergence

      call f_routine(id='ntpoly_solve')

      ! Create the process grid
      call constructprocessgrid(comm)

      ! Conversion of input matrix.
      call chess_to_ntpoly(hmat, hval, nt_hmat)

      ! Load Balancing.
      call constructrandompermutation(perm, nt_hmat%logical_matrix_dimension)

      ! Conversion of parameters for the inverse square root.
      param = solverparameters_t()
      threshold = bigpoly_options//'threshold_overlap'
      param%threshold = real(threshold, ntreal)
      convergence = bigpoly_options//'convergence_overlap'
      param%converge_diff = real(convergence, ntreal)
      param%BalancePermutation = perm

      ! Calculation of Inverse Square Root.
      if (invert_overlap) then
         call chess_to_ntpoly(smat, sval, nt_smat)
         call f_timing(TCAT_NTPOLY_INV, 'ON')
         call inversesquareroot(nt_smat, nt_isq, param)
         call f_timing(TCAT_NTPOLY_INV, 'OFF')
         call ntpoly_to_chess(nt_isq, dmat, isq)
      else
         call chess_to_ntpoly(dmat, isq, nt_isq)
      end if

      ! Conversion of parameters for purification.
      call destructsolverparameters(param)
      param = solverparameters_t()
      threshold = bigpoly_options//'threshold_density'
      param%threshold = real(threshold, ntreal)
      convergence = bigpoly_options//'convergence_density'
      param%converge_diff = real(convergence, ntreal)
      param%BalancePermutation = perm

      ! Actual purification
      call f_timing(TCAT_NTPOLY_PURIFY, 'ON')
      call trs4(nt_hmat, nt_isq, nel, nt_dmat, energy_value_out = ebs, &
                solver_parameters_in = param)
      call f_timing(TCAT_NTPOLY_PURIFY, 'OFF')
      call scalematrix(nt_dmat, 2.0_ntreal)

      ! Conversion of output matrix
      call ntpoly_to_chess(nt_dmat, dmat, dval)

      ! Cleanup
      call destructmatrix(nt_hmat)
      call destructmatrix(nt_smat)
      call destructmatrix(nt_isq)
      call destructmatrix(nt_dmat)
      call destructpermutation(perm)
      call destructsolverparameters(param)
      call destructprocessgrid()

      call f_release_routine()
    end subroutine ntpoly_solve

end module bigpoly_solvers
