#include "macros.h"

module property_vector_pchb_doubles_spinorb_fullyweighted
    !! precomputed heat bath implementation for PropVecCI using spin orbitals and full weighting
    use constants, only: n_int, dp, int64, maxExcit, stdout, bits_n_int
    use util_mod, only: getSpinIndex, near_zero, swap, &
        operator(.implies.), operator(.isclose.), isclose, stop_all, ByteSize_t
    use indexing_mod, only: fuse_symm_idx
    use get_excit, only: make_double, exciteIlut
    use bit_reps, only: decode_bit_det
    use bit_rep_data, only: nIfD
    use orb_idx_mod, only: calc_spin_raw, sum
    use SymExcitDataMod, only: pDoubNew, ScratchSize
    use excitation_types, only: Excite_2_t, excite, spin_allowed
    use aliasSampling, only: AliasSampler_2D_t, AliasSampler_3D_t, do_direct_calculation
    use FciMCData, only: excit_gen_store_type, projEDet
    use SystemData, only: nEl, nBasis, t_mol_3_body
    use sets_mod, only: set, operator(.cap.), operator(.in.)
    use bit_rep_data, only: NIfTot
    use MPI_wrapper, only: iProcIndex_intra, root
    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
    use excitation_generators, only: DoubleExcitationGenerator_t
    better_implicit_none

    private
    public :: PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t

    type, extends(DoubleExcitationGenerator_t) :: PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t
        !! The PropVec PCHB excitation generator for doubles using spin orbitals
        !! and doing full weighting.
        !! This means that first a hole is chosen via \( p( A | I J) |_{A \notin D_i} )\
        !! then a second hole is chosen via \( p( B | I J A) |_{B \notin D_i} )\.
        !! This means it is guaranteed that only unoccupied sites are sampled.
        private
        logical, public :: use_lookup = .false.
            !! Use a lookup for the prop_vec index in global_det_data
        logical, public :: create_lookup = .false.
            !! Create **and** manage! the prop_vec index lookup in global_det_data.

        class(ParticleSelector_t), allocatable :: particle_selector
            !! The particle selector for I, J
        type(AliasSampler_2D_t) :: A_sampler
            !! The sampler for the first hole `A`.
            !! It yields \( p(A | IJ, i_{\text{sg}}) \)
            !! where IJ is a fused index I < J and `i_sg` is the prop_vec.
        type(AliasSampler_3D_t) :: B_sampler
            !! The sampler for the second hole `B`.
            !! It yields \( p(B | A, IJ, i_{\text{sg}}) \)
            !! where IJ is a fused index I < J, A is the first hole and `i_sg` is the prop_vec.
        class(PropertyIndexer_t), allocatable :: indexer
        integer(n_int), private :: last_possible_occupied
            !! The last element of the ilut array has some elements
            !! which are not used, if the number of spinorbitals is not a multiple of
            !! bitsize_n_int.
            !! To correctly zero them this bitmask is 1 wherever a determinant
            !! could be occupied in the last element, and 0 otherwise.
    contains
        private
        procedure, public :: init => PropVec_doubles_PCHB_init
        procedure, public :: finalize => PropVec_doubles_PCHB_finalize
        procedure, public :: gen_exc => PropVec_doubles_PCHB_gen_exc
        procedure, public :: get_pgen => PropVec_doubles_PCHB_get_pgen
        procedure, public :: gen_all_excits => PropVec_doubles_PCHB_gen_all_excits
        procedure, public :: get_memory_demand

        procedure :: compute_samplers => PropVec_doubles_PCHB_compute_samplers
        procedure :: get_unoccupied
    end type PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t

contains


    !>  @brief
    !>  Initialize the pchb excitation generator
    !>
    !>  @details
    !>  This does two things:
    !>  1. setup the lookup table for the mapping ab -> (a,b)
    !>  2. setup the alias table for picking ab given ij with probability ~<ij|H|ab>
    subroutine PropVec_doubles_PCHB_init(this, indexer, &
            use_lookup, create_lookup, PCHB_particle_selection)
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_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_init'

        allocate(this%indexer, source=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 property vector lookup.')
            else
                write(stdout, *) 'PropVec PCHB (RHF) doubles is creating and managing the property vector lookup'
                allocate(lookup_property_indexer, source=this%indexer)
            end if
        end if
        if (this%use_lookup) write(stdout, *) 'PropVec PCHB doubles is using the property vector lookup'

        this%last_possible_occupied = 0_n_int
        block
            integer :: i
            do i = 0, ilut_off(nBasis)
                this%last_possible_occupied = ibset(this%last_possible_occupied, i)
            end do
        end block

        call this%compute_samplers(PCHB_particle_selection)

        write(stdout, *) "Finished excitation generator initialization"
    end subroutine PropVec_doubles_PCHB_init

    elemental function get_memory_demand(this) result(res)
        !! Get the memory demand
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t), intent(in) :: this
        type(ByteSize_t) :: res
        ! Strictly speaking this is not the whole memory demand.
        ! e.g. the gas specification is missing
        res = this%A_sampler%get_memory_demand() &
              + this%B_sampler%get_memory_demand() &
              + this%particle_selector%get_memory_demand()
    end function


    subroutine PropVec_doubles_PCHB_finalize(this)
        !! Finalize everything
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t), intent(inout) :: this

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

    end subroutine PropVec_doubles_PCHB_finalize


    subroutine PropVec_doubles_PCHB_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
        !! double excitation using the hamiltonian matrix elements as weights
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t), intent(inout) :: this
        integer, intent(in) :: nI(nel), exFlag
        integer(n_int), intent(in) :: ilutI(0:NIfTot)
        integer, intent(out) :: nJ(nel), ic, ex(2, maxExcit)
        integer(n_int), intent(out) :: ilutJ(0:NifTot)
        real(dp), intent(out) :: pGen
        logical, intent(out) :: tParity
        HElement_t(dp), intent(out) :: hel
        type(excit_gen_store_type), intent(inout), target :: store
        integer, intent(in), optional :: part_type
        character(*), parameter :: this_routine = 'PropVec_doubles_PCHB_gen_exc'

        integer :: elecs(2), src(2), tgt(2), IJ
        real(dp) :: p_IJ, p_first(2), p_second(2), renorm_first, renorm_second(2)
        integer :: unoccupied(nBasis - nEl)
        integer(n_int) :: ilut_unoccupied(0 : nIfD)
        integer :: i_sg

#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.0_dp)
#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

        ! first, pick two random elecs
        call this%particle_selector%draw(nI, ilutI, i_sg, elecs, src, p_IJ)
        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. (p_IJ .isclose. this%particle_selector%get_pgen(nI, i_sg, src(1), src(2)))) then
            write(stderr, *) ""
            write(stderr, *) "Assertion p_IJ .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_fullyweighted.fpp:188"
            call stop_all (this_routine, "Assert fail: p_IJ .isclose. this%particle_selector%get_pgen(nI, i_sg, src(1), src(2))")
        end if
    end block
#endif

        call this%get_unoccupied(ilutI(0 : nIfD), ilut_unoccupied, unoccupied)
        IJ = fuse_symm_idx(src(1), src(2))

        renorm_first = 1._dp - sum(this%A_sampler%get_prob(IJ, i_sg, nI))
        if (do_direct_calculation(renorm_first)) then
            renorm_first = sum(this%A_sampler%get_prob(IJ, i_sg, unoccupied))
        end if
        block
            integer :: dummy
            call this%A_sampler%constrained_sample(&
                IJ, i_sg, unoccupied, ilut_unoccupied, renorm_first, dummy, tgt(1), p_first(1))
        end block
        if (tgt(1) == 0) then
            call invalidate()
            return
        end if

        renorm_second(1) = 1._dp - sum(this%B_sampler%get_prob(tgt(1), IJ, i_sg, nI))
        if (do_direct_calculation(renorm_second(1))) then
            renorm_second(1) = sum(this%B_sampler%get_prob(tgt(1), IJ, i_sg, unoccupied))
        end if
        if (near_zero(renorm_second(1))) then
            call invalidate()
            return
        end if

        block
            integer :: dummy
            call this%B_sampler%constrained_sample(&
                tgt(1), IJ, i_sg, unoccupied, ilut_unoccupied, &
                renorm_second(1), dummy, tgt(2), p_second(1))
        end block
        if (tgt(2) == 0) then
            call invalidate()
            return
        end if

        ! We could have picked them the other way round.
        ! Account for that.
        ! The renormalization for the first hole is the same
        p_first(2) = this%A_sampler%constrained_getProb(&
            IJ, i_sg, unoccupied, renorm_first, tgt(2))
        renorm_second(2) = 1._dp - sum(this%B_sampler%get_prob(tgt(2), IJ, i_sg, nI))
        if (do_direct_calculation(renorm_second(2))) then
            renorm_second(2) = sum(this%B_sampler%get_prob(tgt(2), IJ, i_sg, unoccupied))
        end if
        p_second(2) = this%B_sampler%constrained_getProb(&
            tgt(2), IJ, i_sg, unoccupied, renorm_second(2), tgt(1))

        pGen = p_IJ * sum(p_first * p_second)

        call make_double(nI, nJ, elecs(1), elecs(2), tgt(1), tgt(2), ex, tParity)
        ilutJ = exciteIlut(ilutI, src, tgt)
        block
            integer :: ClassCount2(ScratchSize), ClassCountUnocc2(ScratchSize)
#ifdef DEBUG_
    block
        use util_mod, only: stop_all
        use constants, only: stderr
        if (.not. (isclose(pgen, this%get_pgen(nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2), atol=1e-10_dp))) then
            write(stderr, *) ""
            write(stderr, *) "Assertion isclose(pgen, this%get_pgen(nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2),&
                & atol=1e-10_dp)"
            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_fullyweighted.fpp:245"
            write(stderr, *) "Values leading to this error:"
            write(stderr, *) "pgen = ", pgen
            write(stderr, *) "this%get_pgen(nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2) = ", this%get_pgen(nI, ilutI, ex,&
                & ic, ClassCount2, ClassCountUnocc2)
            write(stderr, *) "abs(pgen - this%get_pgen(nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2)) = ", abs(pgen -&
                & this%get_pgen(nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2))
            call stop_all (this_routine, "Assert fail: isclose(pgen, this%get_pgen(nI, ilutI, ex, ic, ClassCount2,&
                & ClassCountUnocc2), atol=1e-10_dp)")
        end if
    end block
#endif
        end block

    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_gen_exc


    !>  @brief
    !>  Calculate the probability of drawing a given double excitation ex
    !>
    !>  @param[in] ex  2x2 excitation matrix
    !>
    !>  @return pgen  probability of generating this double with the pchb double excitgen
    function PropVec_doubles_PCHB_get_pgen(this, nI, ilutI, ex, ic, ClassCount2, ClassCountUnocc2) result(pgen)
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t), intent(inout) :: this
        integer, intent(in) :: nI(nel)
        integer(n_int), intent(in) :: ilutI(0:NIfTot)
        integer, intent(in) :: ex(2, maxExcit), ic
        integer, intent(in) :: ClassCount2(ScratchSize), ClassCountUnocc2(ScratchSize)
        real(dp) :: pgen
        real(dp) :: pgen_particle
        character(*), parameter :: this_routine = 'PropVec_doubles_PCHB_get_pgen'

        integer :: IJ, i_sg
        integer :: unoccupied(nBasis - nEl)
        real(dp) :: renorm_first, p_first(2), renorm_second(2), p_second(2)

#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_fullyweighted.fpp:284"
            call stop_all (this_routine, "Assert fail: ic == 2")
        end if
    end block
#endif

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

        pgen_particle = this%particle_selector%get_pgen(nI, i_sg, ex(1, 1), ex(1, 2))

        block
            integer(n_int) :: ilut_unoccupied(0 : nIfD)
            call this%get_unoccupied(ilutI(0 : nIfD), ilut_unoccupied, unoccupied)
        end block
        ! The renormalization for the first hole is the same, regardless of order.
        renorm_first = 1._dp - sum(this%A_sampler%get_prob(IJ, i_sg, nI))
        if (do_direct_calculation(renorm_first)) then
            renorm_first = sum(this%A_sampler%get_prob(IJ, i_sg, unoccupied))
        end if

        associate(A => ex(2, 1), B => ex(2, 2))
            p_first(1) = this%A_sampler%constrained_getProb(IJ, i_sg, unoccupied, renorm_first, A)
            p_first(2) = this%A_sampler%constrained_getProb(IJ, i_sg, unoccupied, renorm_first, B)

            renorm_second(1) = 1._dp - sum(this%B_sampler%get_prob(A, IJ, i_sg, nI))
            if (do_direct_calculation(renorm_second(1))) then
                renorm_second(1) = sum(this%B_sampler%get_prob(A, IJ, i_sg, unoccupied))
            end if
            p_second(1) = this%B_sampler%constrained_getProb(&
                A, IJ, i_sg, unoccupied, renorm_second(1), B)

            renorm_second(2) = 1._dp - sum(this%B_sampler%get_prob(B, IJ, i_sg, nI))
            if (do_direct_calculation(renorm_second(2))) then
                renorm_second(2) = sum(this%B_sampler%get_prob(B, IJ, i_sg, unoccupied))
            end if
            p_second(2) = this%B_sampler%constrained_getProb(&
                B, IJ, i_sg, unoccupied, renorm_second(2), A)
        end associate

        pgen = pgen_particle * sum(p_first * p_second)

    end function PropVec_doubles_PCHB_get_pgen


    subroutine PropVec_doubles_PCHB_compute_samplers(this, PCHB_particle_selection)
        !! computes and stores values for the alias sampling table
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t), intent(inout) :: this
        type(PCHB_ParticleSelection_t), intent(in) :: PCHB_particle_selection
        integer :: I, J, IJ, IJ_max, A, B ! Uppercase because they are indexing spin orbitals
        real(dp), allocatable :: w_A(:), w_B(:), IJ_weights(:, :, :)
        integer, allocatable :: prop_vecs(:, :)
        integer :: i_sg
        ! possible prop_vecs
        prop_vecs = this%indexer%get_prop_vecs()

        ! number of possible source orbital pairs
        IJ_max = fuse_symm_idx(nBasis, nBasis)

        write(stdout, *) "Generating samplers for PCHB excitation generator"
        write(stdout, *) "Depending on the number of property vectors this can take up to 10min."
        call this%A_sampler%shared_alloc([IJ_max, size(prop_vecs, 2)], nBasis, 'A_sampler_PCHB_spinorb_FullyWeighted')
        call this%B_sampler%shared_alloc([nBasis, IJ_max, size(prop_vecs, 2)], nBasis, 'A_sampler_PCHB_spinorb_FullyWeighted')

        if (iProcIndex_intra == root) then
            allocate(w_A(nBasis), w_B(nBasis), source=0._dp)
            allocate(IJ_weights(nBasis, nBasis, size(prop_vecs, 2)), source=0._dp)
        else
            allocate(w_A(0), w_B(0), IJ_weights(0, 0, 0))
        end if

        ! initialize the three samplers
        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, nBasis
                second_particle: do J = I + 1, nBasis
                    IJ = fuse_symm_idx(I, J)
                    w_A(:) = 0.0_dp
                    first_hole: do A = 1, nBasis
                        if (any(A == [I, J])) cycle
                        w_B(:) = 0.0_dp
                        second_hole: do B = 1, nBasis
                            if (any(B == [I, J, A])) cycle
                            associate(exc => merge(Excite_2_t(I, A, J, B), Excite_2_t(I, B, J, A), A < B))
                            if (iProcIndex_intra == root) then
                                if (this%indexer%is_allowed(exc, prop_vecs(:, i_sg)) &
                                        .and. spin_allowed(exc)) then
                                    w_B(B) = get_PCHB_weight(exc)
                                end if
                            end if
                            end associate
                        end do second_hole
                        call this%B_sampler%setup_entry(A, IJ, i_sg, root, w_B(:))
                        if (iProcIndex_intra == root) then
                            w_A(A) = sum(w_B)
                        end if
                    end do first_hole
                    call this%A_sampler%setup_entry(IJ, i_sg, root, w_A(:))

                    if (iProcIndex_intra == root) then
                        IJ_weights(I, J, i_sg) = sum(w_A)
                        IJ_weights(J, I, i_sg) = IJ_weights(I, J, i_sg)
                    end if
                end do second_particle
            end do first_particle
        end do

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

        block
            type(ByteSize_t) :: memory_demand
            memory_demand = this%get_memory_demand()
            write(stdout, *) "The number of supergroups is", size(prop_vecs, 2)
            write(stdout, *) "The excitation generator requires roughly: ", &
                memory_demand%to_human_str(binary=.true.), ' or ', memory_demand%to_human_str()
        end block

    end subroutine PropVec_doubles_PCHB_compute_samplers


    subroutine PropVec_doubles_PCHB_gen_all_excits(this, nI, n_excits, det_list)
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_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 PropVec_doubles_PCHB_gen_all_excits


    pure subroutine get_unoccupied(this, ilutI, ilut_unoccupied, unoccupied)
        !! Return a bitmask and enumeration of the unoccupied spin orbitals.
        class(PropVec_PCHB_DoublesSpinorbFullyWeightExcGen_t), intent(in) :: this
        integer(n_int), intent(in) :: ilutI(0 : nIfD)
        integer(n_int), intent(out) :: ilut_unoccupied(0 : nIfD)
        integer, intent(out) :: unoccupied(nBasis - nEl)
        ilut_unoccupied = not(ilutI)
        ilut_unoccupied(nIfd) = iand(ilut_unoccupied(nIfd), this%last_possible_occupied)
        call decode_bit_det(unoccupied, ilut_unoccupied)
    end subroutine

end module property_vector_pchb_doubles_spinorb_fullyweighted
