using NuclearToolkit
using SpecialFunctions
include("helper.jl")

# Gaussian potentials in HO space
inv_factorial(n) = Iterators.prod(inv.(1:n))
sqrt_factorial(n) = Iterators.prod(sqrt.(n:-1:1))
N_lnk(l, n, k) = 1/sqrt_factorial(n) * binomial(n, k) * sqrt(gamma(n + l + 3/2)) / gamma(k + l + 3/2)
Talmi(l, R, k1, k2; ω=1.0) = (-1)^(k1 + k2) * (1 + 1/(ω * R^2))^-(3/2 + l + k1 + k2) * gamma(3/2 + l + k1 + k2)
V_Gaussian(R, l, n1, n2; ω=1.0) = (-1)^(n1 + n2) * better_sum([N_lnk(l, n1, k1) * N_lnk(l, n2, k2) * Talmi(l, R, k1, k2; ω=ω) for (k1, k2) in Iterators.product(0:n1, 0:n2)])

function get_sp_basis(E_max)
    Es = Int[]
    ns = Int[]
    ls = Int[]

    # E = 2*n + l
    for E in 0 : E_max
        for n in 0 : E ÷ 2
            l = E - 2*n
            push!(Es, E)
            push!(ns, n)
            push!(ls, l)
        end
    end

    return (Es, ns, ls)
end

function get_2p_basis(E_max)
    Es = Int[]
    n1s = Int[]
    l1s = Int[]
    n2s = Int[]
    l2s = Int[]

    # E = 2*n1 + l1 + 2*n2 + l2
    for E in 0 : 2*E_max
        for n1 in 0 : E ÷ 2
            for n2 in 0 : (E - 2*n1) ÷ 2
                for l1 in 0 : (E - 2*n1 - 2*n2)
                    l2 = E - 2*n1 - 2*n2 - l1
                    push!(Es, E)
                    push!(n1s, n1)
                    push!(l1s, l1)
                    push!(n2s, n2)
                    push!(l2s, l2)
                end
            end
        end
    end

    return (Es, n1s, l1s, n2s, l2s)
end

get_V_matrix(V_l, ls, ns) = throw("unimplemented")

function sp_T_matrix(ns, ls; ω=1.0, μ=1.0)
    mat = zeros(length(ns), length(ns))
    for idx in CartesianIndices(mat)
        (i, j) = Tuple(idx)
        if ls[i] == ls[j]
            mat[idx] = ls[i]
            if ns[i] == ns[j]
                mat[idx] += 1/4 * (4*ns[j] + 3)
            elseif ns[i] == ns[j] + 1
                mat[idx] += -1/4 * sqrt((2*ns[j] + 2) * (2*ns[j] + 3))
            elseif ns[i] == ns[j] - 1
                mat[idx] += -1/4 * sqrt((2*ns[j] + 1) * 2*ns[j])
            end
        end
    end
    return (ω / μ) .* mat
end

get_H_matrix(V_l, ns, ls) = get_T_matrix(ns, ls) + get_V_matrix(V_l, ns, ls)

function Moshinsky_transform(Es, n1s, l1s, n2s, l2s, Λ)
    l_max = 2*max(maximum(l1s), maximum(l2s)) # OPTIMIZE
    E_max = maximum(Es) # OPTIMIZE
    j_max = l_max # OPTIMIZE
    to = 0 # unused
    
    dtri = NuclearToolkit.prep_dtri(l_max) ;     println("dtri prepared")
    dcgm0 = NuclearToolkit.prep_dcgm0(l_max) ;     println("dcgm0 prepared")
    d6j = NuclearToolkit.prep_d6j_int(E_max, j_max, to) ;     println("d6j prepared")

    mat = zeros(length(Es), length(Es))
    s = hcat(Es, n1s, l1s, n2s, l2s)
    for idx in CartesianIndices(mat)
        (i, j) = Tuple(idx)
        (Elhs, N, L, n, l) = s[i, :]
        (Erhs, n1, l1, n2, l2) = s[j, :]
        if Elhs == Erhs
            mat[i, j] = NuclearToolkit.HObracket_d6j(N, L, n, l, n1, l1, n2, l2, Λ, 1.0, dtri, dcgm0, d6j, to)
        end
    end

    return mat
end