include("../p_space.jl")
include("../EC.jl")

Λ = 0
m = 1.0

# Distinguishable particles: V12 = bound and V13 & V23 = resonant
Vsubsystem_of_r(r) = -2 * exp(-r^2/4)
Vdecay_of_r(r) = -exp(-r^2 / 3) + exp(-r^2 / 10)

ϕ = 0.1
vertices = [0, 6 * exp(-1im * ϕ)]
subdivisions = [50]
jmax = 4

E_max = 40
μω_global = 0.4

# Jacobi coordinates
μ1ω1 = μω_global * 1/2
μ2ω2 = μω_global * 2
μ1 = m * 1/2
μ2 = m * 2/3

atol=10^-5; maxevals=10^5; R_cutoff=16; verbose=true;

###########

verbose && println("No of threads = ", Threads.nthreads())

V_l(j, k, kp) = Vl_mat_elem(Vsubsystem_of_r, j, k, kp; atol=atol, maxevals=maxevals, R_cutoff=R_cutoff)

ks, ws = get_mesh(vertices, subdivisions)
weights = repeat(kron(ws, ws), jmax + 1)
block_size = length(ks)
tri((j1, j2)) = triangle_ineq(j1, j2, Λ)
js = collect(Iterators.filter(tri, iter_prod(0:jmax, 0:jmax)))
basis = iter_prod(js, zip(ks, ws), zip(ks, ws)) # basis = ((j1, j2), (k1, w1), (k2, w2))
basis_size = length(js) * length(ks)^2
@assert length(basis) == basis_size "Something wrong with the basis"
verbose && println("Basis size = $basis_size")

@time "Block diagonal part" begin
    blocks = [kron_sum(get_H_matrix((k, kp) -> V_l(j1, k, kp), ks, ws, μ1), get_T_matrix(ks, μ2)) for (j1, _) in js]
    Ha = blockdiag(sparse.(blocks)...)
end

basis_ho = ho_basis_2B(E_max, Λ)
verbose && println("HO basis size = ", basis_ho.dim)

@time "V2_HO" V2_HO =  get_jacobi_V2_matrix(Vdecay_of_r, basis_ho, μω_global)

@time "W_right" W_right = get_W_matrix(basis, basis_ho, μ1ω1, μ2ω2; weights=true)
@time "W_left" W_left = get_W_matrix(basis, basis_ho, μ1ω1, μ2ω2; weights=false)

@time "V2" Vb = W_left * V2_HO * transpose(W_right)

###########

training_c = [-0.55, -0.7, -0.85, -1, -1.2]
extrapolating_c = [0.2, 0.1, 0.0, -0.1, -0.2, -0.3]

ref_E = -0.5173809356244544

exact_ref = [-0.31360746615280954-0.07689284936870341im
    -0.3233372403877718-0.06011323914565665im
    -0.339615582074795-0.04239442037174759im
    -0.36333816534241997-0.02648721825958402im
    -0.39376650561322885-0.014382935339817332im	
    -0.4299479825535172-0.006510710745123606im]

EC = affine_EC(Ha, Vb)
train!(EC, training_c; ref_eval=ref_E, CAEC=true)
extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_ref)

exportCSV(EC, "temp/dis_CSM_B2R.csv")
plot(EC, "temp/dis_CSM_B2R.pdf")