#include "macros.h"

module property_vector_pchb_doubles_spinorb_fastweighted
    !! spin-orbital-resolved PropVecCI-PCHB using the "fast weighted" scheme
    use constants, only: dp, n_int, maxExcit, stdout, int64
    use util_mod, only: operator(.isclose.), near_zero, stop_all, operator(.implies.)
    use indexing_mod, only: fuse_symm_idx
    use sets_mod, only: operator(.in.), set
    use get_excit, only: make_double, exciteIlut
    use dSFMT_interface, only: genrand_real2_dSFMT
    use FciMCData, only: excit_gen_store_type
    use UMatCache, only: numBasisIndices
    use SystemData, only: nEl, nBasis
    use SymExcitDataMod, only: ScratchSize
    use sltcnd_mod, only: nI_invariant_sltcnd_excit
    use bit_rep_data, only: nIfTot
    use excitation_generators, only: DoubleExcitationGenerator_t
    use FciMCData, only: ProjEDet, excit_gen_store_type
    use MPI_wrapper, only: root
    use excitation_types, only: Excite_2_t, spin_allowed
    use aliasSampling, only: AliasSampler_2D_t
    use property_vector_index, only: PropertyIndexer_t, lookup_property_indexer
    use property_vector_pchb_doubles_select_particles, only: &
        ParticleSelector_t, PCHB_ParticleSelection_t, &
        allocate_and_init, get_PCHB_weight
    better_implicit_none

    private
    public :: PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t

    type, extends(DoubleExcitationGenerator_t) :: PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t
        !! PropVec PCHB excitation generator for doubles using spin orbitals, with
        !! "fast weighting"
        !! This means we choose holes via \(p(AB|IJ)\)
        private
        logical, public :: use_lookup = .false.
            !! use a lookup for the prop_vec index
        logical, public :: create_lookup = .false.
            !! create and manage the prop_vec index

        class(ParticleSelector_t), allocatable :: particle_selector
            !! particle selector for getting I, J
        type(AliasSampler_2D_t) :: AB_sampler
            !! sampler for holes, dubbed \(AB\) (fused index)
            !! It yields \( p(AB | IJ, i_{\text{sg}}) \)
            !! where IJ is a fused index I < J and `i_sg` is the prop_vec.

        class(PropertyIndexer_t), allocatable :: indexer
        integer, allocatable :: tgtOrbs(:, :)
    contains
        private
        procedure, public :: init => PropVec_doubles_PCHB_spinorb_init
        procedure, public :: finalize => PropVec_doubles_PCHB_spinorb_finalize
        procedure, public :: gen_exc => PropVec_doubles_PCHB_spinorb_gen_exc
        procedure, public :: get_pgen => PropVec_doubles_PCHB_spinorb_get_pgen
        procedure, public :: gen_all_excits => PropVec_doubles_PCHB_spinorb_gen_all_excits

        procedure :: compute_samplers => PropVec_doubles_PCHB_spinorb_compute_samplers
    end type PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t

contains

    subroutine PropVec_doubles_PCHB_spinorb_init(this, indexer, use_lookup, &
                            create_lookup, PCHB_particle_selection)
        !! initalises the spinorb-resolved PCHB doubles excitation generator
        !!
        !! more specifically, sets up a lookup table for ab -> (a,b) and
        !! sets up the alias table for picking ab given ij with prob ~ Hijab
        class(PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t), intent(inout) :: this
        class(PropertyIndexer_t), intent(in) :: indexer
        logical, intent(in) :: use_lookup, create_lookup
        type(PCHB_ParticleSelection_t), intent(in) :: PCHB_particle_selection
        character(*), parameter :: this_routine = 'PropVec_doubles_PCHB_spinorb_init'

        integer :: AB, A, B, abMax, nBI

        this%indexer = indexer
        this%create_lookup = create_lookup
        this%use_lookup = use_lookup

        if (this%create_lookup) then
            if (allocated(lookup_property_indexer)) then
                call stop_all(this_routine, 'Someone else is already managing the supergroup lookup.')
            else
                write(stdout, *) 'PropVec PCHB (spinorb) doubles is creating and managing the supergroup lookup'
                allocate(lookup_property_indexer, source=this%indexer)
            end if
        end if
        if (this%use_lookup) write(stdout, *) 'PropVec PCHB doubles is using the supergroup lookup'

        write(stdout, *) "Allocating PCHB excitation generator objects"
        ! number of *spin* orbs
        nBI = nBasis
        ! initialize the mapping ab -> (a, b)
        abMax = fuse_symm_idx(nBI, nBI)
        allocate(this%tgtOrbs(2, 0:abMax), source=0)
        do A = 1, nBI
            do B = 1, A
                AB = fuse_symm_idx(A, B)
                ! b comes first as this is an ordered list
                this%tgtOrbs(1, AB) = B
                this%tgtOrbs(2, AB) = A
            end do
        end do

        ! set up the alias table
        call this%compute_samplers(nBI, PCHB_particle_selection)

        write(stdout, *) "Finished excitation generator initialization"

    end subroutine PropVec_doubles_PCHB_spinorb_init

    subroutine PropVec_doubles_PCHB_spinorb_finalize(this)
        !! deallocates the sampler and mapper
        class(PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t), intent(inout) :: this

        if (allocated(this%particle_selector)) then
            call this%AB_sampler%finalize()
            call this%particle_selector%finalize()
            ! Yes, we assume, that either all or none are allocated
            deallocate(this%particle_selector, this%tgtOrbs, this%indexer)
            if (this%create_lookup) deallocate(lookup_property_indexer)
        end if
    end subroutine PropVec_doubles_PCHB_spinorb_finalize

    subroutine PropVec_doubles_PCHB_spinorb_gen_exc(&
                this, nI, ilutI, nJ, ilutJ, exFlag, ic, &
                ex, tParity, pGen, hel, store, part_type)
        !! given the initial determinant (both as nI and ilut), create a random
        !! doubles excitation using the Hamiltonian matrix elements as weights
        class(PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t), intent(inout) :: this
            !! the exctitation generator
        integer, intent(in) :: nI(nel), exFlag
            !! determinant to excite from
            !! unused in this generator
        integer(n_int), intent(in) :: ilutI(0:NIfTot)
            !! determint from which to excite, ilut format
        integer, intent(out) :: nJ(nel), ic, ex(2, maxExcit)
            !! the excited determinant upon return
            !! excitation order (for doubles generator, always == 2)
            !! excitation matrix nI -> nJ
        integer(n_int), intent(out) :: ilutJ(0:NifTot)
            !! excited determinant, ilut format
        real(dp), intent(out) :: pGen
            !! probability of generating the excitation
        logical, intent(out) :: tParity
            !! the parity of the excitation nI -> nJ
        HElement_t(dp), intent(out) :: hel
            !! matrix element Hijab
        type(excit_gen_store_type), intent(inout), target :: store
            !!
        integer, intent(in), optional :: part_type
            !! unused in this generator
        character(*), parameter :: this_routine = 'PropVec_doubles_PCHB_spinorb_gen_exc'

        integer :: i_sg ! prop_vec index
        integer :: src(2) ! particles (I, J)
        integer :: elecs(2) ! particle indices in nI
        integer :: tgt(2)

        integer :: ij
        logical :: invalid

        integer :: ab
        real(dp) :: pGenHoles


#ifdef WARNING_WORKAROUND_
        associate(exFlag => exFlag); end associate
        associate(part_type => part_type); end associate
#endif
        ic = 2
#ifdef WARNING_WORKAROUND_
        hel = h_cast(0._dp) ! macro
#endif
        if (this%use_lookup) then
            i_sg = this%indexer%lookup_prop_vec_idx(store%idx_curr_dets, nI)
        else
            i_sg = this%indexer%idx_nI(nI)
        end if

        ! pick two random electrons
        call this%particle_selector%draw(nI, ilutI, i_sg, elecs, src, pGen)
        if (src(1) == 0) then
            call invalidate()
            return
        end if

#ifdef DEBUG_
    block
        use util_mod, only: stop_all
        use constants, only: stderr
        if (.not. (pgen .isclose. this%particle_selector%get_pgen(nI, i_sg, src(1), src(2)))) then
            write(stderr, *) ""
            write(stderr, *) "Assertion pgen .isclose. this%particle_selector%get_pgen(nI, i_sg, src(1), src(2))"
            write(stderr, *) "failed in /home/cicdadmin/fkfest-gh-runners/altest-runner-7/_work/NECI_autoforward/NECI_autoforward/s&
                &rc/property_vector_constraint/property_vector_pchb_doubles_spinorb_fastweighted.fpp:188"
            call stop_all (this_routine, "Assert fail: pgen .isclose. this%particle_selector%get_pgen(nI, i_sg, src(1), src(2))")
        end if
    end block
#endif

        invalid = .false.
        ij = fuse_symm_idx(src(1), src(2))

        ! get a pair of orbitals using the precomputed weights
        call this%AB_sampler%sample(ij, i_sg, ab, pGenHoles)
#ifdef DEBUG_
    block
        use util_mod, only: stop_all
        use constants, only: stderr
        if (.not. (ab /= 0 .implies. (pGenHoles .isclose. this%AB_sampler%get_prob(ij, i_sg, ab)))) then
            write(stderr, *) ""
            write(stderr, *) "Assertion ab /= 0 .implies. (pGenHoles .isclose. this%AB_sampler%get_prob(ij, i_sg, ab))"
            write(stderr, *) "failed in /home/cicdadmin/fkfest-gh-runners/altest-runner-7/_work/NECI_autoforward/NECI_autoforward/s&
                &rc/property_vector_constraint/property_vector_pchb_doubles_spinorb_fastweighted.fpp:195"
            call stop_all (this_routine, "Assert fail: ab /= 0 .implies. (pGenHoles .isclose. this%AB_sampler%get_prob(ij, i_sg,&
                & ab))")
        end if
    end block
#endif

        if (ab == 0) then
            invalid = .true.
            tgt = 0
        else
            ! split ab -> a,b
            tgt = this%tgtOrbs(:, ab)
#ifdef DEBUG_
    block
        use util_mod, only: stop_all
        use constants, only: stderr
        if (.not. (all(tgt /= 0) .implies. tgt(1) /= tgt(2))) then
            write(stderr, *) ""
            write(stderr, *) "Assertion all(tgt /= 0) .implies. tgt(1) /= tgt(2)"
            write(stderr, *) "failed in /home/cicdadmin/fkfest-gh-runners/altest-runner-7/_work/NECI_autoforward/NECI_autoforward/s&
                &rc/property_vector_constraint/property_vector_pchb_doubles_spinorb_fastweighted.fpp:203"
            call stop_all (this_routine, "Assert fail: all(tgt /= 0) .implies. tgt(1) /= tgt(2)")
        end if
    end block
#endif
            invalid = any(tgt == 0) .or. any(tgt(1) == nI) .or. any(tgt(2) == nI)
        end if

        ! as in the spatially-resolved case, there is a very rare case where (due
        ! to floating point error) an excitation with pgen=0 is created. Invalidate.
        if (.not. invalid .and. near_zero(pGenHoles)) then
            invalid = .true.
            ! Yes, print. Those events are signficant enough to be always noted in the output
            write(stdout, *) "WARNING: Generated excitation with probability of 0"
        end if

        if (invalid) then
            ! if 0 returned, no excitations -> nulldet
            call invalidate()
        else
            ! construct the determinant given the selected orbitals and electrons
            call make_double(nI, nJ, elecs(1), elecs(2), tgt(1), tgt(2), ex, tParity)

            ilutJ = exciteIlut(ilutI, src, tgt)

            pgen = pgen * pgenholes

            block
                integer :: classCount2(ScratchSize), classCountUnocc2(ScratchSize)
#ifdef DEBUG_
    block
        use util_mod, only: stop_all
        use constants, only: stderr
        if (.not. (pGen .isclose. this%get_pgen(nI, ilutI, ex, ic, classCount2, classCountUnocc2))) then
            write(stderr, *) ""
            write(stderr, *) "Assertion pGen .isclose. this%get_pgen(nI, ilutI, ex, ic, classCount2, classCountUnocc2)"
            write(stderr, *) "failed in /home/cicdadmin/fkfest-gh-runners/altest-runner-7/_work/NECI_autoforward/NECI_autoforward/s&
                &rc/property_vector_constraint/property_vector_pchb_doubles_spinorb_fastweighted.fpp:228"
            call stop_all (this_routine, "Assert fail: pGen .isclose. this%get_pgen(nI, ilutI, ex, ic, classCount2,&
                & classCountUnocc2)")
        end if
    end block
#endif
            end block
        end if

    contains

        subroutine invalidate()
            nJ = 0
            ilutJ = 0_n_int
            ex(1, 1 : 2) = src
            ex(2, 1 : 2) = tgt
        end subroutine invalidate

    end subroutine PropVec_doubles_PCHB_spinorb_gen_exc

    real(dp) function PropVec_doubles_PCHB_spinorb_get_pgen(this, nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2) result(pgen)
        !! calculates the probability of drawing a given double excitation
        !! parametrised by the excitation matrix ex
        class(PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t), intent(inout) :: this
        integer, intent(in) :: nI(nel)
        integer(n_int), intent(in) :: ilutI(0:NIfTot)
        integer, intent(in) :: ex(2, maxExcit), ic
            !! excitation matrix
        integer, intent(in) :: ClassCount2(ScratchSize), ClassCountUnocc2(ScratchSize)
        integer :: i_sg, IJ, AB
        character(*), parameter :: this_routine = 'PropVec_doubles_PCHB_spinorb_get_pgen'

#ifdef WARNING_WORKAROUND_
        associate(ilutI => ilutI); end associate
        associate(ClassCount2 => ClassCount2); end associate
        associate(ClassCountUnocc2 => ClassCountUnocc2); end associate
#endif
#ifdef DEBUG_
    block
        use util_mod, only: stop_all
        use constants, only: stderr
        if (.not. (ic == 2)) then
            write(stderr, *) ""
            write(stderr, *) "Assertion ic == 2"
            write(stderr, *) "failed in /home/cicdadmin/fkfest-gh-runners/altest-runner-7/_work/NECI_autoforward/NECI_autoforward/s&
                &rc/property_vector_constraint/property_vector_pchb_doubles_spinorb_fastweighted.fpp:256"
            call stop_all (this_routine, "Assert fail: ic == 2")
        end if
    end block
#endif

        IJ = fuse_symm_idx(ex(1, 1), ex(1, 2))
        AB = fuse_symm_idx(ex(2, 1), ex(2, 2))
        i_sg = this%indexer%idx_nI(nI)

        pgen = this%particle_selector%get_pgen(nI, i_sg, ex(1, 1), ex(1, 2))
        pgen = pgen * this%AB_sampler%get_prob(IJ, i_sg, AB)

    end function PropVec_doubles_PCHB_spinorb_get_pgen

    subroutine PropVec_doubles_PCHB_spinorb_gen_all_excits(this, nI, n_excits, det_list)
        class(PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t), intent(in) :: this
        integer, intent(in) :: nI(nEl)
        integer, intent(out) :: n_excits
        integer(n_int), allocatable, intent(out) :: det_list(:,:)

        call this%indexer%gen_all_excits(nI, n_excits, det_list, ic=2)

    end subroutine

    subroutine PropVec_doubles_PCHB_spinorb_compute_samplers(this, nBI, PCHB_particle_selection)
        !! computes and stores values for the alias (spin-independent) sampling table
        class(PropVec_PCHB_DoublesSpinOrbFastWeightExcGen_t), intent(inout) :: this
        integer, intent(in) :: nBI
        type(PCHB_ParticleSelection_t), intent(in) :: PCHB_particle_selection
        integer :: I, J, IJ, IJMax
        integer :: A, B, AB, ABMax
        integer(int64) :: memCost
            !! n_prop_vec * num_fused_indices * bytes_per_sampler
        real(dp), allocatable :: w(:), IJ_weights(:, :, :)
        integer, allocatable :: prop_vecs(:, :)
        integer :: i_sg
        ! possible prop_vecs
        prop_vecs = this%indexer%get_prop_vecs()

        ! number of possible spin orbital pairs
        ijMax = fuse_symm_idx(nBI, nBI)
        abMax = ijMax

        memCost = size(prop_vecs, 2, kind=int64) &
                  * int(ijMax, int64) &
                  * int(abMax, int64) * 8_int64

        write(stdout, *) "Excitation generator requires", real(memCost, dp) / 2.0_dp**30, "GB of memory"
        write(stdout, *) "The number of supergroups is", size(prop_vecs, 2)
        write(stdout, *) "Generating samplers for PCHB excitation generator"
        write(stdout, *) "Depending on the number of supergroups this can take up to 10min."

        call this%AB_sampler%shared_alloc([ijMax, size(prop_vecs, 2)], abMax, 'AB_PCHB_spinorb_sampler')

        ! One could allocate only on the intra-node-root here, if memory
        ! at initialization ever becomes an issue.
        ! Look at `gasci_pchb_doubles_spin_fulllyweighted.fpp` for inspiration.
        allocate(w(abMax), IJ_weights(nBI, nBI, size(prop_vecs, 2)), source=0._dp)

        prop_vec: do i_sg = 1, size(prop_vecs, 2)
            if (mod(i_sg, 100) == 0) write(stdout, *) 'Still generating the samplers'
            first_particle: do I = 1, nBI
                ! A,B are ordered, so we can assume I < J
                second_particle: do J = I + 1, nBI
                    w(:) = 0._dp
                    ! for each (I,J), get all matrix elements <IJ|H|AB> and use them as
                    ! weights to prepare the sampler
                    first_hole: do A = 1, nBI
                        if (any(A == [I, J])) cycle
                        second_hole: do B = 1, nBI
                            if (any(B == [I, J, A])) cycle
                            AB = fuse_symm_idx(A, B)
                            associate(exc => merge(Excite_2_t(I, A, J, B), Excite_2_t(I, B, J, A), A < B))
                                if (this%indexer%is_allowed(exc, prop_vecs(:, i_sg)) .and. spin_allowed(exc)) then
                                    w(AB) = get_PCHB_weight(exc)
                                end if
                            end associate
                        end do second_hole
                    end do first_hole
                    IJ = fuse_symm_idx(I, J)
                    call this%AB_sampler%setup_entry(IJ, i_sg, root, w)

                    IJ_weights(I, J, i_sg) = sum(w)
                    IJ_weights(J, I, i_sg) = IJ_weights(I, J, i_sg)
                end do second_particle
            end do first_particle
        end do prop_vec

        call allocate_and_init(PCHB_particle_selection, this%indexer, &
            IJ_weights, root, this%use_lookup, this%particle_selector)

    end subroutine PropVec_doubles_PCHB_spinorb_compute_samplers

end module property_vector_pchb_doubles_spinorb_fastweighted
