!> @file
!! Constrained DFT (based on linear version)
!! @author
!!    Copyright (C) 2007-2013 BigDFT group
!!    This file is distributed under the terms of the
!!    GNU General Public License, see ~/COPYING file
!!    or http://www.gnu.org/copyleft/gpl.txt .
!!    For the list of contributors, see ~/AUTHORS


!> Module to perform constrained DFT calculations
module constrained_dft
  use module_precisions
  use sparsematrix_base, only: sparse_matrix, matrices
  use module_bigdft_arrays
  implicit none

  private

  ! these two types are only needed to build the weight matrix
  type, public :: fragment_constraint_components
    integer :: frag ! fragment which this information is associated to
    integer :: spin ! should be 1 for optical case
    ! variables which are relevant for electronic and optical constraints
    integer :: ncomp ! number of orbitals which contribute
    integer, dimension(:), pointer :: orbital_indices
    real(gp), dimension(:), pointer :: orbital_weights
  end type fragment_constraint_components

  type, public :: input_constraint_data
     character(len=16) :: constraint_type  ! can be spatial, charge_transfer, optical or electronic

     ! electronic
     integer :: nfrags_orbital
     type(fragment_constraint_components), dimension(:), pointer :: orbital_frag_constraints
     real(gp), dimension(:), pointer :: orbital ! orbital in the SF basis

     ! optical
     character(len=7) :: excitation_type ! singlet or triplet
     integer :: nfrags_electron, nfrags_hole ! this should be 1 in the hole case, unless you want to somehow merge fragments??
     type(fragment_constraint_components), dimension(:), pointer :: electron_frag_constraints, hole_frag_constraints
     real(gp), dimension(:), pointer :: hole, electron ! hole and electron orbitals in the SF basis
 
     ! spatial
     integer :: nfrags_spatial
     type(fragment_constraint_components), dimension(:), pointer :: spatial_frag_constraints

     ! CT
     integer :: nfrags_source, nfrags_dest
     type(fragment_constraint_components), dimension(:), pointer :: source_frag_constraints, dest_frag_constraints
  end type input_constraint_data

  ! the following two types are the only ones used externally
  type, public :: constraint_data
     real(gp) :: lag_mult ! the Lagrange multiplier used to enforce the constraint
     type(sparse_matrix) :: weight_matrix ! matrix elements of the weight function
     type(matrices) :: weight_matrix_ ! matrix elements of the weight function
  end type constraint_data

  type, public :: cdft_data
     real(kind=8),dimension(:),pointer :: psit_c, psit_f ! ground state SFs - needed for defining orbital constraint
     integer :: num_constraints ! the number of constraint objects, which should correspond to the length of constraints
     type(constraint_data),dimension(:),pointer :: constraints ! the list of constraints to be imposed
     type(input_constraint_data),dimension(:),pointer :: input_constraints ! the input data associated with the list of constraints to be imposed
  end type cdft_data

  ! the remaining routines are only used internally
  public :: cdft_from_dict, input_constraint_data_free
  public :: nullify_cdft_data, cdft_data_free, cdft_data_allocate

contains

  subroutine nullify_cdft_data(cdft)
    implicit none
    type(cdft_data), intent(out) :: cdft

    nullify(cdft%constraints)
    nullify(cdft%input_constraints)
    nullify(cdft%psit_c)
    nullify(cdft%psit_f)
    cdft%num_constraints = 0

  end subroutine nullify_cdft_data

  subroutine cdft_data_free(cdft)
    implicit none
    type(cdft_data), intent(inout) :: cdft

    character(len=200), parameter :: subname='cdft_data_free'
    integer :: i

    call f_free_ptr(cdft%psit_c)
    call f_free_ptr(cdft%psit_f)
    do i=1,cdft%num_constraints
      call constraint_data_free(cdft%constraints(i))
      ! to be written:
      !call input_constraint_data_free(cdft%input_constraints(i))
    end do
    do i=1,cdft%num_constraints
      call nullify_constraint_data(cdft%constraints(i))
      ! to be written:
      !call nullify_input_constraint_data(cdft%input_constraints(i))
    end do
    if (associated(cdft%constraints)) then
       deallocate(cdft%constraints)
    end if
    if (associated(cdft%input_constraints)) then
       deallocate(cdft%input_constraints)
    end if
    call nullify_cdft_data(cdft)

  end subroutine cdft_data_free

  subroutine cdft_data_allocate(cdft,ham,collcom)
    use sparsematrix_base, only: sparse_matrix, sparsematrix_malloc_ptr, &
         SPARSE_FULL, assignment(=),copy_sparse_matrix
    use communications_base, only: comms_linear
    use module_bigdft_profiling
    implicit none
    type(cdft_data), intent(inout) :: cdft
    type(sparse_matrix), intent(in) :: ham
    type(comms_linear), intent(in) :: collcom

    character(len=200), parameter :: subname='cdft_data_allocate'
    integer :: i

    call f_routine(id='cdft_data_allocate')
    do i=1,cdft%num_constraints
      call constraint_data_allocate(cdft%constraints(i),ham)
    end do
    cdft%psit_c = f_malloc_ptr(collcom%ndimind_c,id='cdft%psit_c')
    cdft%psit_f = f_malloc_ptr(7*collcom%ndimind_f,id='cdft%psit_f')

    call f_release_routine()

  end subroutine cdft_data_allocate

  subroutine nullify_constraint_data(constraint)
    use sparsematrix_base, only: sparse_matrix_null, matrices_null
    use module_bigdft_errors
    implicit none
    type(constraint_data), intent(out) :: constraint
    constraint%lag_mult = 0.0_gp
    constraint%weight_matrix = sparse_matrix_null()
    constraint%weight_matrix_ = matrices_null()
  end subroutine nullify_constraint_data

  subroutine constraint_data_free(constraint)
    use sparsematrix_base, only: deallocate_sparse_matrix, deallocate_matrices
    implicit none
    type(constraint_data), intent(inout) :: constraint

    character(len=200), parameter :: subname='constraint_data_free'

    call deallocate_sparse_matrix(constraint%weight_matrix)
    call deallocate_matrices(constraint%weight_matrix_)
    call nullify_constraint_data(constraint)
  end subroutine constraint_data_free

  subroutine constraint_data_allocate(constraint,ham)
    use sparsematrix_base, only: sparse_matrix, sparsematrix_malloc_ptr, &
         SPARSE_FULL, assignment(=),copy_sparse_matrix
    use module_bigdft_profiling
    implicit none
    type(constraint_data), intent(inout) :: constraint
    type(sparse_matrix), intent(in) :: ham

    character(len=200), parameter :: subname='constraint_data_allocate'
    !!integer :: istat

    call f_routine(id='constraint_data_allocate')
    !call sparse_copy_pattern(ham, cdft%weight_matrix, bigdft_mpi%iproc, subname)
    call copy_sparse_matrix(ham, constraint%weight_matrix)
    !cdft%weight_matrix_%matrix_compr=f_malloc_ptr(cdft%weight_matrix%nvctr,id='cdft%weight_matrix%matrix_compr')
    constraint%weight_matrix_%matrix_compr=sparsematrix_malloc_ptr(constraint%weight_matrix,iaction=SPARSE_FULL, &
                                                                   id='constraint%weight_matrix%matrix_compr')

    call f_release_routine()

  end subroutine constraint_data_allocate

  subroutine nullify_fragment_constraint_components(frag_constraint)
    implicit none
    type(fragment_constraint_components), intent(out) :: frag_constraint

    ! set spin and frag to 1, so if they are not specified it should still behave sensibly
    frag_constraint%spin = 1
    frag_constraint%frag = 1
    frag_constraint%ncomp = 0
    nullify(frag_constraint%orbital_indices)
    nullify(frag_constraint%orbital_weights)

  end subroutine nullify_fragment_constraint_components
 
  subroutine fragment_constraint_components_free(frag_constraint)
    implicit none
    type(fragment_constraint_components), intent(inout) :: frag_constraint

    if (associated(frag_constraint%orbital_weights)) then
       call f_free_ptr(frag_constraint%orbital_weights)
    end if
    if (associated(frag_constraint%orbital_indices)) then
       call f_free_ptr(frag_constraint%orbital_indices)
    end if
    call nullify_fragment_constraint_components(frag_constraint)
 
  end subroutine fragment_constraint_components_free

  subroutine fragment_constraint_components_allocate(frag_constraint)
    implicit none
    type(fragment_constraint_components), intent(inout) :: frag_constraint

    frag_constraint%orbital_indices = f_malloc_ptr(frag_constraint%ncomp,id='orbital_indices')
    frag_constraint%orbital_weights = f_malloc_ptr(frag_constraint%ncomp,id='orbital_weights')

  end subroutine fragment_constraint_components_allocate

  subroutine nullify_input_constraint_data(input_constraint)
    implicit none
    type(input_constraint_data), intent(out) :: input_constraint

    input_constraint%constraint_type = 'NONE'
    input_constraint%excitation_type = 'NONE'
    input_constraint%nfrags_orbital = 0
    input_constraint%nfrags_electron = 0
    input_constraint%nfrags_hole = 0
    input_constraint%nfrags_spatial = 0
    input_constraint%nfrags_source = 0
    input_constraint%nfrags_dest = 0
    nullify(input_constraint%orbital_frag_constraints)
    nullify(input_constraint%orbital)
    nullify(input_constraint%electron_frag_constraints)
    nullify(input_constraint%electron)
    nullify(input_constraint%hole_frag_constraints)
    nullify(input_constraint%hole)
    nullify(input_constraint%spatial_frag_constraints)
    nullify(input_constraint%source_frag_constraints)
    nullify(input_constraint%dest_frag_constraints)

  end subroutine nullify_input_constraint_data
 
  subroutine input_constraint_data_free(input_constraint)
    implicit none
    type(input_constraint_data), intent(inout) :: input_constraint
    integer :: ifrag

    do ifrag=1,input_constraint%nfrags_orbital
       call fragment_constraint_components_free(input_constraint%orbital_frag_constraints(ifrag))
    end do
    if (associated(input_constraint%orbital_frag_constraints)) then
       deallocate(input_constraint%orbital_frag_constraints)
    end if
    if (associated(input_constraint%orbital)) then
       call f_free_ptr(input_constraint%orbital)
    end if

    do ifrag=1,input_constraint%nfrags_electron
       call fragment_constraint_components_free(input_constraint%electron_frag_constraints(ifrag))
    end do
    if (associated(input_constraint%electron_frag_constraints)) then
       deallocate(input_constraint%electron_frag_constraints)
    end if
    if (associated(input_constraint%electron)) then
       call f_free_ptr(input_constraint%electron)
    end if

    do ifrag=1,input_constraint%nfrags_hole
       call fragment_constraint_components_free(input_constraint%hole_frag_constraints(ifrag))
    end do
    if (associated(input_constraint%hole_frag_constraints)) then
       deallocate(input_constraint%hole_frag_constraints)
    end if
    if (associated(input_constraint%hole)) then
       call f_free_ptr(input_constraint%hole)
    end if

    do ifrag=1,input_constraint%nfrags_spatial
       call fragment_constraint_components_free(input_constraint%spatial_frag_constraints(ifrag))
    end do
    if (associated(input_constraint%spatial_frag_constraints)) then
       deallocate(input_constraint%spatial_frag_constraints)
    end if

    do ifrag=1,input_constraint%nfrags_source
       call fragment_constraint_components_free(input_constraint%source_frag_constraints(ifrag))
    end do
    if (associated(input_constraint%source_frag_constraints)) then
       deallocate(input_constraint%source_frag_constraints)
    end if

    do ifrag=1,input_constraint%nfrags_dest
       call fragment_constraint_components_free(input_constraint%dest_frag_constraints(ifrag))
    end do
    if (associated(input_constraint%dest_frag_constraints)) then
       deallocate(input_constraint%dest_frag_constraints)
    end if

    call nullify_input_constraint_data(input_constraint)
 
  end subroutine input_constraint_data_free

  subroutine cdft_from_dict(dict,cdft)
    use dictionaries
    implicit none
    type(dictionary), pointer :: dict
    type(cdft_data), intent(out) :: cdft
    !local variables
    type(dictionary), pointer :: dict_tmp
    integer :: iconstraint

    call nullify_cdft_data(cdft)

    cdft%num_constraints = dict_len(dict)
    allocate(cdft%input_constraints(cdft%num_constraints))
    allocate(cdft%constraints(cdft%num_constraints))

    ! loop over the list of constraints which are present
    iconstraint = 1
    dict_tmp=>dict_iter(dict)
    do while(associated(dict_tmp))
       ! should we also have a nullify input_constraints routine here? or is it enough if we are setting defaults when we read the dict?
       call input_constraint_from_dict(dict_tmp,cdft%input_constraints(iconstraint),cdft%constraints(iconstraint))
       dict_tmp=>dict_next(dict_tmp)
       iconstraint = iconstraint + 1
    end do

  end subroutine cdft_from_dict

  subroutine input_constraint_from_dict(dict,input_constraint,constraint)
    use dictionaries
    implicit none
    type(dictionary), pointer :: dict
    type(input_constraint_data), intent(out) :: input_constraint
    type(constraint_data), intent(out) :: constraint
    !local variables
    type(dictionary), pointer :: dict_tmp, dict_tmp2
    integer :: ifrag

    call nullify_constraint_data(constraint)
    call nullify_input_constraint_data(input_constraint)

    dict_tmp=>dict_iter(dict)
    do while(associated(dict_tmp))
       select case(trim(dict_key(dict_tmp)))

       case('Vc')
           constraint%lag_mult = dict_tmp
       case('constraint_type')
           input_constraint%constraint_type = dict_tmp
       case('excitation_type')
          input_constraint%excitation_type = dict_tmp
       case('electron')
           input_constraint%nfrags_electron = dict_len(dict_tmp)
           allocate(input_constraint%electron_frag_constraints(input_constraint%nfrags_electron))
           ifrag = 1
           dict_tmp2=>dict_iter(dict_tmp)
           do while(associated(dict_tmp2))
              call fragment_constraint_components_from_dict(dict_tmp2, input_constraint%electron_frag_constraints(ifrag))
              ifrag = ifrag + 1
              dict_tmp2=>dict_next(dict_tmp2)
           end do

       case('hole')
           input_constraint%nfrags_hole = dict_len(dict_tmp)
           allocate(input_constraint%hole_frag_constraints(input_constraint%nfrags_hole))
           ifrag = 1
           dict_tmp2=>dict_iter(dict_tmp)
           do while(associated(dict_tmp2))
              call fragment_constraint_components_from_dict(dict_tmp2, input_constraint%hole_frag_constraints(ifrag))
              ifrag = ifrag + 1
              dict_tmp2=>dict_next(dict_tmp2)
           end do

       case('orbital')
           ! should this be 1 or can an electronic excitation have more than one orbital component?
           input_constraint%nfrags_orbital = dict_len(dict_tmp)
           allocate(input_constraint%orbital_frag_constraints(input_constraint%nfrags_orbital))
           ifrag = 1
           dict_tmp2=>dict_iter(dict_tmp)
           do while(associated(dict_tmp2))
              call fragment_constraint_components_from_dict(dict_tmp2, input_constraint%orbital_frag_constraints(ifrag))
              ifrag = ifrag + 1
              dict_tmp2=>dict_next(dict_tmp2)
           end do

       case('source')
           input_constraint%nfrags_source = dict_len(dict_tmp)
           allocate(input_constraint%source_frag_constraints(input_constraint%nfrags_source))
           ifrag = 1
           dict_tmp2=>dict_iter(dict_tmp)
           do while(associated(dict_tmp2))
              call fragment_constraint_components_from_dict(dict_tmp2, input_constraint%source_frag_constraints(ifrag))
              ifrag = ifrag + 1
              dict_tmp2=>dict_next(dict_tmp2)
           end do

       case('destination')
           input_constraint%nfrags_dest = dict_len(dict_tmp)
           allocate(input_constraint%dest_frag_constraints(input_constraint%nfrags_dest))
           ifrag = 1
           dict_tmp2=>dict_iter(dict_tmp)
           do while(associated(dict_tmp2))
              call fragment_constraint_components_from_dict(dict_tmp2, input_constraint%dest_frag_constraints(ifrag))
              ifrag = ifrag + 1
              dict_tmp2=>dict_next(dict_tmp2)
           end do

       case('spatial')
           input_constraint%nfrags_spatial = dict_len(dict_tmp)
           allocate(input_constraint%spatial_frag_constraints(input_constraint%nfrags_spatial))
           ifrag = 1
           dict_tmp2=>dict_iter(dict_tmp)
           do while(associated(dict_tmp2))
              call fragment_constraint_components_from_dict(dict_tmp2, input_constraint%spatial_frag_constraints(ifrag))
              ifrag = ifrag + 1
              dict_tmp2=>dict_next(dict_tmp2)
           end do

       !case default
           ! do we put an error message here (and in the other routines) in case of invalid input?

       end select

       dict_tmp=>dict_next(dict_tmp)
    end do

  end subroutine input_constraint_from_dict

  subroutine fragment_constraint_components_from_dict(dict,frag_constraint)
     use dictionaries
     implicit none
     type(dictionary), pointer :: dict
     type(fragment_constraint_components), intent(out) :: frag_constraint
     !local variables
     type(dictionary), pointer :: iter, orb_iter
     integer :: iorb
     character(len=24) :: orb_str

     call nullify_fragment_constraint_components(frag_constraint)  
     nullify(iter)
     do while(iterating(iter,on=dict))
         select case(dict_key(iter))
         case('fragment')
            frag_constraint%frag = iter
         case('spin')
            frag_constraint%spin = iter
         case('orbitals')
            frag_constraint%ncomp = dict_size(iter)  ! we could eventually remove this variable and other size-related ones, but keeping for now
            call fragment_constraint_components_allocate(frag_constraint)
            nullify(orb_iter)
            iorb = 1
            do while(iterating(orb_iter,on=iter))
               orb_str = dict_key(orb_iter)
               read(orb_str,*) frag_constraint%orbital_indices(iorb)
               frag_constraint%orbital_weights(iorb) = orb_iter
               iorb = iorb + 1
            end do
         end select
     end do

  end subroutine fragment_constraint_components_from_dict

end module constrained_dft
