using SparseArrays
using QuadGK
using LRUCache
include("helper.jl")
include("math.jl")

function V_numerical(V_of_r, l, n1, n2; μω_gen=1.0, atol=0, maxevals=10^7)
    integrand(r) = sqrt(μω_gen) * ho_basis(l, n1, sqrt(μω_gen) * r) * ho_basis(l, n2, sqrt(μω_gen) * r) * V_of_r(r)
    (integral, _) = quadgk(integrand, 0, Inf; atol=atol, maxevals=maxevals)
    return integral
end

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, Λ=-1)
    Es = Int[]
    n1s = Int[]
    l1s = Int[]
    n2s = Int[]
    l2s = Int[]

    # E = 2*n1 + l1 + 2*n2 + l2
    for E in E_max : -2 : 0 # same parity states only
        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
                    if Λ≥0 && !triangle_ineq(l1, l2, Λ); continue; end
                    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

function get_sp_T_matrix(ns, ls; mask=trues(length(ns),length(ns)), μω_gen=1.0, μ=1.0)
    mat = spzeros(length(ns), length(ns))
    for idx in CartesianIndices(mat)
        if !mask[idx]; continue; end
        (i, j) = Tuple(idx)
        if ls[i] == ls[j]
            if ns[i] == ns[j]
                mat[idx] = ns[j] + ls[i]/2 + 3/4
            elseif abs(ns[i]-ns[j]) == 1
                n_max = max(ns[i], ns[j])
                mat[idx] = -(1/2) * sqrt(n_max * (n_max + ls[i] + 1/2))
            end
        end
    end
    return (μω_gen / μ) .* mat
end

function get_sp_V_matrix(V_l, ns, ls; mask=trues(length(ns),length(ns)), dtype=Float64, cache=LRU{Tuple{UInt8, UInt8, UInt8}, dtype}(maxsize=(1+maximum(ns))^2))
    mat = zeros(dtype, length(ns), length(ns))
    Threads.@threads for idx in CartesianIndices(mat) 
        if !mask[idx]; continue; end
        (i, j) = Tuple(idx)
        if ls[i] == ls[j]
            l = UInt8(ls[i])
            n1, n2 = UInt8.(minmax(ns[i], ns[j])) # assuming transpose symmetry
            mat[idx] = (get!(cache, (l, n1, n2)) do; V_l(l, n1, n2); end)
        end
    end
    return sparse(mat)
end

function Moshinsky_transform(Es, n1s, l1s, n2s, l2s, Λ)
    NQMAX = maximum(Es)
    @assert all(mod.(Es, 2) .== mod(NQMAX, 2)) "Can only admit basis states with same parity"

    LMIN = Λ
    LMAX = Λ
    CO = 1/sqrt(2)
    SI = 1/sqrt(2)

    # dimensions BRAC(0:LMAX,0:(NQMAX-LMIN)/2,0:(NQMAX-LMIN)/2,0:(NQMAX-LMIN)/2,0:(NQMAX-LMIN)/2,0:LMAX,0:(NQMAX-LMIN)/2,LMIN:LMAX)
    BRAC = zeros(Float64, 1 + LMAX, 1 + (NQMAX - LMIN) ÷ 2, 1 + (NQMAX - LMIN) ÷ 2, 1 + (NQMAX - LMIN) ÷ 2, 1 + (NQMAX - LMIN) ÷ 2, 1 + LMAX, 1 + (NQMAX-LMIN) ÷ 2, 1 + LMAX-LMIN)
    @ccall "../OSBRACKETS/allosbrac.so".allosbrac_(NQMAX::Ref{Int32},LMIN::Ref{Int32},LMAX::Ref{Int32},CO::Ref{Float64},SI::Ref{Float64},BRAC::Ptr{Array{Float64}})::Cvoid
    
    mat = zeros(length(Es), length(Es))

    s = hcat(Es, n1s, l1s, n2s, l2s)
    Threads.@threads 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 && triangle_ineq(L, l, Λ) && triangle_ineq(l1, l2, Λ)
            mat[i, j] = (-1)^(n1 + n2 + N + n) * pick_Moshinsky_bracket(BRAC, n1, l1, n2, l2, N, L, n, l, Λ)
        end
    end

    return sparse(mat)
end

function pick_Moshinsky_bracket(BRAC, n1′, l1′, n2′, l2′, n1, l1, n2, l2, Λ) # Efros notation -- don't confuse
    ϵ = (l1 + l2 - Λ) % 2
    NP = (l1′ - l2′ + Λ - ϵ) ÷ 2
    MP = (l1′ + l2′ - Λ - ϵ) ÷ 2
    N = (l1 - l2 + Λ - ϵ) ÷ 2
    M = (l1 + l2 - Λ - ϵ) ÷ 2

    # BRAC(NP,N1P,MP,N1,N2,N,M,L)
    return BRAC[1 + NP, 1 + n1′, 1 + MP, 1 + n1, 1 + n2, 1 + N, 1 + M, 1]
end

function get_jacobi_V_matrix(V_of_r, E_max, Λ, μ1ω1, μω_global; atol=10^-6, maxevals=10^5)
    _, n1s, l1s, n2s, l2s = get_2p_basis(E_max, Λ)
    l_max = max(maximum(l1s), maximum(l2s))
    n_max = max(maximum(n1s), maximum(n2s))
    mask1 = (n2s .== n2s') .&& (l2s .== l2s')
    
    V1_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μ1ω1, atol=atol, maxevals=maxevals)
    V1_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
    V1 = get_sp_V_matrix(V1_elem, n1s, l1s; mask=mask1, dtype=ComplexF64, cache=V1_cache)

    V2 =  get_jacobi_V2_matrix(V_of_r, E_max, Λ, μω_global; atol=atol, maxevals=maxevals)

    return V1 + V2
end

function get_jacobi_V2_matrix(V_of_r, E_max, Λ, μω_global; atol=10^-6, maxevals=10^5)
    Es, n1s, l1s, n2s, l2s = get_2p_basis(E_max, Λ)
    l_max = max(maximum(l1s), maximum(l2s))
    n_max = max(maximum(n1s), maximum(n2s))
    mask1 = (n2s .== n2s') .&& (l2s .== l2s')
    mask2 = (n1s .== n1s') .&& (l1s .== l1s')
    
    V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals)
    V_relative_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)

    V_relative = get_sp_V_matrix(V_relative_elem, n1s, l1s; mask=mask1, dtype=ComplexF64, cache=V_relative_cache) + get_sp_V_matrix(V_relative_elem, n2s, l2s; mask=mask2, dtype=ComplexF64, cache=V_relative_cache)
    U = Moshinsky_transform(Es, n1s, l1s, n2s, l2s, Λ)
    V2 =  U' * V_relative * U

    return V2
end