module tdhf_sf_gradient_mod

  use precision, only: dp
  use grd2, only: grd2_driver, grd2_compute_data_t
  use basis_tools, only: basis_set, bas_norm_matrix, build_cart_density
  use constants, only: HARMONIC_ACTIVE, NUM_CART_BF
  use types, only: information

  implicit none

  character(len=*), parameter :: module_name = "tdhf_sf_gradient_mod"

  public tdhf_sf_gradient

  type, extends(grd2_compute_data_t) :: grd2_sf_compute_data_t
    real(kind=dp), pointer :: d2(:,:,:) => null()
    real(kind=dp), pointer :: p2(:,:,:) => null()
    real(kind=dp), pointer :: v2(:,:) => null()
    ! Cartesian-effective (bfnrm-folded) copies + offsets for HARMONIC_ACTIVE.
    real(kind=dp), allocatable :: d2a_c(:,:), d2b_c(:,:), p2a_c(:,:), p2b_c(:,:), v2_c(:,:)
    integer, allocatable :: cart_off(:)
    integer :: nbf = 0
  contains
    procedure :: init => grd2_sf_compute_data_t_init
    procedure :: clean => grd2_sf_compute_data_t_clean
    procedure :: get_density => grd2_sf_compute_data_t_get_density
    procedure :: build_cart => grd2_sf_build_cart
  end type

contains

  subroutine sf_gradient_C(c_handle) bind(C, name="tdhf_sf_gradient")
    use c_interop, only: oqp_handle_t, oqp_handle_get_info
    use types, only: information
    type(oqp_handle_t) :: c_handle
    type(information), pointer :: inf
    inf => oqp_handle_get_info(c_handle)
    call tdhf_sf_gradient(inf)
  end subroutine sf_gradient_c

  subroutine tdhf_sf_gradient(infos)
    use io_constants, only: iw
    use oqp_tagarray_driver

    use types, only: information
    use strings, only: Cstring, fstring
    use basis_tools, only: basis_set
    use messages, only: show_message, with_abort

    use grd1, only: eijden, print_gradient
    use util, only: measure_time
    use tdhf_lib, only: &
      iatogen, mntoia
    use tdhf_sf_lib, only: &
      sfrorhs, &
      sfromcal, sfrogen, sfrolhs, &
      pcgb, sfropcal, sfrowcal
    use dft, only: dft_initialize, dftclean
    use mathlib, only: symmetrize_matrix
    use mod_dft_molgrid, only: dft_grid_t
    use mod_dft_gridint_tdxc_grad, only: utddft_xc_gradient
    use mathlib, only: unpack_matrix
    use printing, only: print_module_info

    implicit none

    character(len=*), parameter :: subroutine_name = "tdhf_sf_gradient"

    type(basis_set), pointer :: basis
    type(information), target, intent(inout) :: infos

    integer :: s_size

    integer :: nbf, nbf_tri
    logical :: roref = .false.

    type(dft_grid_t) :: molGrid

  ! General data
    logical :: dft
    integer :: scf_type, mol_mult

    real(kind=dp), allocatable :: p(:,:,:), v(:,:,:), d(:,:,:)

    ! tagarray
    real(kind=dp), contiguous, pointer :: &
      dmat_a(:), dmat_b(:), td_abxc(:,:), td_p(:,:)
    character(len=*), parameter :: tags_general(*) = (/ character(len=80) :: &
      OQP_DM_A, OQP_DM_B, OQP_td_abxc, OQP_td_p /)

    mol_mult = infos%mol_prop%mult
!    if (.not. (mol_mult == 3 .or. mol_mult == 4)) then
!      call show_message( &
!        'SF-TDDFT only supports mult=3 (triplet) or mult=4 (quartet) references', &
!        with_abort)
!    end if 

    scf_type = infos%control%scftype
    if (scf_type==3) roref = .true.

    dft = infos%control%hamilton == 20

  ! Files open
    open (unit=IW, file=infos%log_filename, position="append")
  !
    call print_module_info('SF_Grad','Computing Gradient of SF-TDDFT')
!
    write(iw,'(/5X,"Gradient options"/&
                &5X,18("-")/&
                &5X,"Target State: ",I8/&
                &,5X,"*Note that the ground state of SF-TDDFT is 1.*"/)')&
                & infos%tddft%target_state

  ! Load basis set
    basis => infos%basis
    basis%atoms => infos%atoms

    call data_has_tags(infos%dat, tags_general, module_name, subroutine_name, WITH_ABORT)
    call tagarray_get_data(infos%dat, OQP_DM_A, dmat_a)
    call tagarray_get_data(infos%dat, OQP_DM_B, dmat_b)
    call tagarray_get_data(infos%dat, OQP_td_abxc, td_abxc)
    call tagarray_get_data(infos%dat, OQP_td_p, td_p)

  ! Allocate H, S ,T and D matrices
    nbf = basis%nbf
    nbf_tri = nbf*(nbf+1)/2
    s_size = (basis%nshell**2+basis%nshell)/2

!   Compute 1e gradient
    call flush(iw)

    call sf_1e_grad(infos, basis)

    write(iw,"(' ..... End Of 1-Eelectron Gradient ......')")
    call measure_time(print_total=1, log_unit=iw)
    call flush(iw)

    allocate(d(nbf,nbf,2), source=0.0d0)
    allocate(p(nbf,nbf,2), source=0.0d0)

    call unpack_matrix(td_p(:,1), p(:,:,1))
    call unpack_matrix(td_p(:,2), p(:,:,2))

    call unpack_matrix(dmat_a, d(:,:,1))
    call unpack_matrix(dmat_b, d(:,:,2))

!   Compute xc gradient
    if (dft) then
      call dft_initialize(infos, basis, molGrid, verbose=.true.)

      call utddft_xc_gradient(basis=basis, &
           molGrid=molGrid, &
           dedft=infos%atoms%grad, &
           da=d(:,:,1), &
           db=d(:,:,2), &
           pa=p(:,:,1:1), &
           pb=p(:,:,2:2), &
           nmtx=1, &
           !threshold=1.0d-15, &
           threshold=0.0d0, &
           infos=infos)

      call dftclean(infos)
      call measure_time(print_total=1, log_unit=iw)
      call flush(iw)
    end if

!   Compute 2e gradient
    allocate(v(nbf,nbf,2), source=0.0d0)
    v(:,:,1) = td_abxc
    call sf_2e_grad(basis, infos, d, p, v(:,:,1))

    call print_gradient(infos)

!   Print timings
    call measure_time(print_total=1, log_unit=iw)

    close(iw)

  end subroutine tdhf_sf_gradient

!###############################################################################

  subroutine sf_1e_grad(infos, basis)

    use oqp_tagarray_driver
    use types, only: information
    use basis_tools, only: basis_set
    use util, only: measure_time
    use messages, only: show_message, WITH_ABORT
    use precision, only: dp
    use constants, only: tol_int
    use grd1, only: eijden, print_gradient, &
            grad_nn, grad_ee_overlap, &
            grad_ee_kinetic, grad_en_hellman_feynman, grad_en_pulay, grad_1e_ecp

    use mathlib, only: symmetrize_matrix

    implicit none

    character(len=*), parameter :: subroutine_name = "sf_1e_grad"

    type(information), intent(inout) :: infos
    type(basis_set), intent(inout) :: basis

    real(kind=dp), allocatable :: dens(:)
    real(kind=dp) :: tol
    integer :: nbf, nbf_tri, ok

    ! tagarray
    real(kind=dp), pointer :: dmat_a(:), dmat_b(:), wao(:), td_p(:,:)
    character(len=*), parameter :: tags_general(4) = (/ character(len=80) :: &
      OQP_DM_A, OQP_DM_B, OQP_WAO, OQP_td_p /)

    nbf = basis%nbf
    nbf_tri = nbf*(nbf+1)/2

    tol = tol_int*log(10.0_dp)

!   initial memory allocation
    allocate(dens(nbf_tri), source=0.0_dp, stat=ok)
    if(ok/=0) call show_message('Cannot allocate memory', WITH_ABORT)

    call data_has_tags(infos%dat, tags_general, module_name, subroutine_name, WITH_ABORT)
    call tagarray_get_data(infos%dat, OQP_DM_A, dmat_a)
    call tagarray_get_data(infos%dat, OQP_DM_B, dmat_b)
    call tagarray_get_data(infos%dat, OQP_WAO, wao)
    call tagarray_get_data(infos%dat, OQP_td_p, td_p)

    associate( grad   => infos%atoms%grad &
             , xyz    => infos%atoms%xyz &
             , zn     => infos%atoms%zn &
             , p      => td_p &
             , w      => wao &
             , urohf  => infos%control%scftype>=2 &
        )

!     Zero out gradient
      grad = 0.0d0
!     Nuclear repulsion force
      call grad_nn(infos%atoms, infos%basis%ecp_zn_num)

!     Obtain Lagrangian matrix (`dens`)
      call eijden(dens, nbf, infos)

!     Add W matrix:
      dens = dens + 2*w

!     Overlap gradient
      call grad_ee_overlap(basis, dens, grad, logtol=tol)

!     Compute total density matrix, discard Lagrangian
      dens = dmat_a + p(:,1)
      if (infos%control%scftype>=2) then
        dens = dens + dmat_b + p(:,2)
      end if

!     Hellmann-Feynman force
      call grad_en_hellman_feynman(basis, xyz, zn, dens, grad, logtol=tol)

!     KE gradient
      call grad_ee_kinetic(basis, dens, grad, logtol=tol)

!     Pulay force
      call grad_en_pulay(basis, xyz, zn, dens, grad, logtol=tol)

!     Effective core potential gradient
      call grad_1e_ecp(infos, basis, xyz, dens, grad, logtol=tol)

    end associate

   end subroutine

!###############################################################################

!> @brief The driver for the two electron gradient
  subroutine sf_2e_grad(basis, infos, d, p, v)

    use basis_tools, only: basis_set
    use precision, only: dp
    use messages, only: show_message, WITH_ABORT
    use types, only: information

    implicit none

    type(information), target, intent(inout) :: infos
    type(basis_set) :: basis
    real(kind=dp), contiguous, target :: p(:,:,:), d(:,:,:), v(:,:)

    logical :: urohf, dft
    real(kind=dp) :: scale_exch  !> HF scale in Reference
    real(kind=dp) :: scale_exch2 !> HF scale in Response

    integer :: ok
    real(kind=dp), allocatable :: de(:,:)
    class(grd2_compute_data_t), allocatable :: gcomp

    dft = infos%control%hamilton == 20 ! dft or hf
    urohf = infos%control%scftype >= 2

    scale_exch = 1.0_dp
    scale_exch2 = 1.0_dp
    if (dft) then
      scale_exch = infos%dft%HFscale
      scale_exch2 = infos%tddft%HFscale
    end if

    allocate(de(3,ubound(infos%atoms%zn,1)), &
            source=0.0d0, &
            stat=ok)

    if(ok/=0) call show_message('cannot allocate memory', WITH_ABORT)

    write(*, '(/7x,"Fitting parameters")')
    if (.not.infos%dft%cam_flag) then
      write(*, '(10x,"Exact HF exchange:")')
      write(*, '(5x,"Reference: |", t20, f6.3, t29, "|")') scale_exch
      write(*, '(5x,"Response:  |", t20, f6.3, t29, "|")') scale_exch2
    else
      write(*, '(10x,"CAM parametres:")')
      write(*, '(16x,"|   alpha   |    beta   |     mu    |")')
      write(*, '(5x,"Reference: |", t20, f6.3, t29, "|", t32, f6.3, t41, "|", t44, f6.3, t53, "|")') &
         infos%dft%cam_alpha, infos%dft%cam_beta, infos%dft%cam_mu
      write(*, '(5x,"Response:  |", t20, f6.3, t29, "|", t32, f6.3, t41, "|", t44, f6.3, t53, "|")') &
         infos%tddft%cam_alpha, infos%tddft%cam_beta, infos%tddft%cam_mu
    end if
    write(*, '(10x,"Spin-pair coupling parametres:")')
    write(*, '(16x,"|   CO-CO   |   OV-OV   |   CO-OV   |")')
    write(*, '(16x,"|", t20, f6.3, t29, "|", t32, f6.3, t41, "|", t44, f6.3, t53, "|")') &
       infos%tddft%HFscale, infos%tddft%HFscale, infos%tddft%HFscale

    gcomp =  grd2_sf_compute_data_t( d2 = d &
                                   , p2 = p &
                                   , v2 = v &
                                   , nbf = basis%nbf )

    call gcomp%init()

    select type (gcomp)
    class is (grd2_sf_compute_data_t)
      call gcomp%build_cart(basis)
    end select

    call grd2_driver(infos, basis, de, gcomp, &
                     cam = dft.and.infos%dft%cam_flag, &
                     alpha = infos%tddft%cam_alpha, &
                     beta = infos%tddft%cam_beta, &
                     mu = infos%tddft%cam_mu)

    infos%atoms%grad = infos%atoms%grad + de

    call gcomp%clean()

  end subroutine

!###############################################################################

  subroutine grd2_sf_compute_data_t_init(this)
    implicit none
    class(grd2_sf_compute_data_t), target, intent(inout) :: this

    call this%clean()

    this%d2(:,:,1) = this%d2(:,:,1) +   this%d2(:,:,2)
    this%d2(:,:,2) = this%d2(:,:,1) - 2*this%d2(:,:,2)

    this%p2(:,:,1) = this%p2(:,:,1) +   this%p2(:,:,2)
    this%p2(:,:,2) = this%p2(:,:,1) - 2*this%p2(:,:,2)

  end subroutine

!###############################################################################

!> @brief Cartesian-effective copies of the SF gradient densities (alpha/beta
!>   d and p, transition v) for HARMONIC_ACTIVE. Call AFTER init (which
!>   combines the spin densities).
  subroutine grd2_sf_build_cart(this, basis)
    class(grd2_sf_compute_data_t), intent(inout) :: this
    type(basis_set), intent(in) :: basis
    integer, allocatable :: od(:)
    integer :: nc
    if (.not. HARMONIC_ACTIVE) return
    call sf_cart_one(basis, this%d2(:,:,1), this%d2a_c, this%cart_off, nc)
    call sf_cart_one(basis, this%d2(:,:,2), this%d2b_c, od, nc)
    call sf_cart_one(basis, this%p2(:,:,1), this%p2a_c, od, nc)
    call sf_cart_one(basis, this%p2(:,:,2), this%p2b_c, od, nc)
    call sf_cart_one(basis, this%v2,        this%v2_c,  od, nc)
  end subroutine grd2_sf_build_cart

  subroutine sf_cart_one(basis, m, m_cart, off, nc)
    type(basis_set), intent(in) :: basis
    real(kind=dp), intent(in) :: m(:,:)
    real(kind=dp), allocatable, intent(out) :: m_cart(:,:)
    integer, allocatable, intent(out) :: off(:)
    integer, intent(out) :: nc
    real(kind=dp), allocatable :: tmp(:,:)
    tmp = m
    call bas_norm_matrix(tmp, basis%bfnrm, basis%nbf)
    call build_cart_density(basis, tmp, m_cart, off, nc)
  end subroutine sf_cart_one

!###############################################################################

  subroutine grd2_sf_compute_data_t_clean(this)
    implicit none
    class(grd2_sf_compute_data_t), target, intent(inout) :: this
  end subroutine

!###############################################################################

!> @brief This routine forms the product of density
!>        matrices for use in forming the two electron
!>        gradient. Valid for closed and open shell SCF.
  subroutine grd2_sf_compute_data_t_get_density(this, basis, id, dab, dabmax)

    implicit none

    class(grd2_sf_compute_data_t), target, intent(inout) :: this
    type(basis_set), intent(in) :: basis
    integer, intent(in) :: id(4)
    real(kind=dp), target, intent(out) :: dab(*)
    real(kind=dp), intent(out) :: dabmax

    real(kind=dp) :: df1, dq1, dt2, bfn
    real(kind=dp) :: coulfact, xcfact, xcfact2
    integer :: i, j, k, l
    integer :: loc(4)
    integer :: nbf(4)
    real(kind=dp), pointer :: ab(:,:,:,:)
    real(kind=dp), pointer :: d2a(:,:), d2b(:,:), p2a(:,:), p2b(:,:), v2(:,:)
    logical :: usecart
    integer :: i1, j1, k1, l1

    coulfact = 4*this%coulscale
    xcfact = this%hfscale
    xcfact2 = this%hfscale2
    dabmax = 0

    usecart = HARMONIC_ACTIVE
    if (usecart) then
      d2a => this%d2a_c;  d2b => this%d2b_c
      p2a => this%p2a_c;  p2b => this%p2b_c;  v2 => this%v2_c
      loc = this%cart_off(id) - 1
      nbf = NUM_CART_BF(basis%am(id))
    else
      d2a => this%d2(:,:,1);  d2b => this%d2(:,:,2)
      p2a => this%p2(:,:,1);  p2b => this%p2(:,:,2);  v2 => this%v2
      loc = basis%ao_offset(id) - 1
      nbf = basis%naos(id)
    end if

    ab(1:nbf(4),1:nbf(3),1:nbf(2),1:nbf(1)) => dab(1:product(nbf))

    do i = 1, nbf(1)
      i1 = loc(1) + i

      do j = 1, nbf(2)
        j1 = loc(2) + j

        do k = 1, nbf(3)
          k1 = loc(3) + k

          do l = 1, nbf(4)
            l1 = loc(4) + l
            df1 = (d2a(i1,j1)+p2a(i1,j1))*d2a(k1,l1) &
                +  d2a(i1,j1)                  *p2a(k1,l1)
            df1 = df1 * coulfact

            if (xcfact /= 0.0_dp .or. xcfact2 /= 0.0_dp) then
              dq1 = (d2a(i1,k1)+p2a(i1,k1))*d2a(j1,l1) &
                  +  d2a(i1,k1)                  *p2a(j1,l1) &
                  + (d2a(i1,l1)+p2a(i1,l1))*d2a(j1,k1) &
                  +  d2a(i1,l1)                  *p2a(j1,k1) &
                  + (d2b(i1,k1)+p2b(i1,k1))*d2b(j1,l1) &
                  +  d2b(i1,k1)                  *p2b(j1,l1) &
                  + (d2b(i1,l1)+p2b(i1,l1))*d2b(j1,k1) &
                  +  d2b(i1,l1)                  *p2b(j1,k1)
              dt2 = v2(i1,k1)*v2(j1,l1) &
                  + v2(k1,i1)*v2(l1,j1) &
                  + v2(i1,l1)*v2(j1,k1) &
                  + v2(l1,i1)*v2(k1,j1)

              df1 = df1-xcfact*dq1-xcfact2*2.0_dp*dt2
            end if
            dabmax = max(dabmax, abs(df1))
            bfn = 1.0_dp
            if (.not. usecart) bfn = product(basis%bfnrm([i1,j1,k1,l1]))
            ab(l,k,j,i) = df1*bfn
          end do
        end do
      end do
    end do
  end subroutine grd2_sf_compute_data_t_get_density

!###############################################################################

end module tdhf_sf_gradient_mod
