# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.7.2
#     language: julia
#     name: julia-1.7
# ---

pwd()

using Pkg; Pkg.activate("../../FermiCG/")
using FermiCG, NPZ
readdir()

# +
# load integrals from disk
h0 = npzread("cr_qubit_scf_integrals_h0.npz.npy")
h1 = npzread("cr_qubit_scf_integrals_h1.npz.npy")
h2 = npzread("cr_qubit_scf_integrals_h2.npz.npy")
ints = InCoreInts(h0, h1, h2)

ints_original = deepcopy(ints);

print(" Integrals have the following sizes: h0= ", size(h0), " h1= ", size(h1), " h2= ", size(h2))


# +
# Define clusters - this should probably be done in the python notebook, and then just read in here
clusters_in = [
    (1:10),   # metal
    (11:16), # Benzene 1
    (17:22), # Benzene 1
    (23:28), # Benzene 1
    (29:34), # Benzene 1
]

clusters = [Cluster(i,collect(clusters_in[i])) for i = 1:length(clusters_in)]

init_fspace = [
    (5,5),
    (3,3),
    (3,3),
    (3,3),
    (3,3)];
display(clusters)
display(init_fspace)

# + tags=[]
rdm1 = zeros(size(ints.h1))

e_cmf, U, da1, db1 = FermiCG.cmf_oo(ints, clusters, init_fspace, rdm1, rdm1, verbose=0, gconv=1e-6, method="cg", sequential=true);

# + tags=[]
using JLD2
# rotate integrals to new cmf optimized basis
ints = FermiCG.orbital_rotation(ints_original,U);

@save "cr_qubit_cmf_5clusters.jld2" ints clusters init_fspace da1 db1


Da = da1
Db = db1 
max_roots = 20

#
# form Cluster data
@time cluster_bases = FermiCG.compute_cluster_eigenbasis(ints, clusters, verbose=0,
                                                   max_roots=max_roots,
                                                   init_fspace=init_fspace,
                                                   rdm1a=Da, rdm1b=Db)


clustered_ham = FermiCG.extract_ClusteredTerms(ints, clusters)

cluster_ops = FermiCG.compute_cluster_ops(cluster_bases, ints);

FermiCG.add_cmf_operators!(cluster_ops, cluster_bases, ints, Da, Db);




# -

@save "cr_qubit_cmf_5clusters_bases_M20.jld2" ints clusters init_fspace da1 db1 cluster_bases


# +
# Create BST state

v = FermiCG.BSTstate(clusters, FockConfig(init_fspace), cluster_bases, R=3)

FermiCG.add_single_excitons!(v, FockConfig(init_fspace), cluster_bases)
FermiCG.randomize!(v)
FermiCG.orthonormalize!(v)
e_ci, v = FermiCG.tucker_ci_solve(v, cluster_ops, clustered_ham);


# + jupyter={"source_hidden": true} tags=[]
using Printf
display(v, root=2)
for ei in 1:length(e_ci)
    @printf(" %5i %12.8f %12.8f eV\n", ei, e_ci[ei]+ints.h0, (e_ci[ei]-e_ci[1])*27.21165)
end

# + tags=[]
 e_var, v_var = FermiCG.block_sparse_tucker(v, cluster_ops, clustered_ham,
                                        max_iter    = 20,
                                        max_iter_pt = 200,
                                        nbody       = 4,
                                        H0          = "Hcmf",
                                        thresh_var  = 1e-1,
                                        thresh_foi  = 1e-3,
                                        thresh_pt   = 1e-3,
                                        ci_conv     = 1e-5,
                                        do_pt       = true,                          
                                        resolve_ss  = false,
                                        tol_tucker  = 1e-4);

# -

@time e2 = FermiCG.compute_pt2_energy(v_var, cluster_ops, clustered_ham, thresh_foi=1e-5)
@time e2 = FermiCG.compute_pt2_energy(v_var, cluster_ops, clustered_ham, thresh_foi=1e-6)
@time e2 = FermiCG.compute_pt2_energy(v_var, cluster_ops, clustered_ham, thresh_foi=1e-7)

# + tags=[]
v1 = deepcopy(v_var);
display(v1,root=3)

# + tags=[]
 e2, v2 = FermiCG.block_sparse_tucker(v1, cluster_ops, clustered_ham,
                                        max_iter    = 20,
                                        max_iter_pt = 200,
                                        nbody       = 4,
                                        H0          = "Hcmf",
                                        thresh_var  = 1e-1,
                                        thresh_foi  = 1e-4,
                                        thresh_pt   = 1e-3,
                                        ci_conv     = 1e-5,
                                        do_pt       = true,                          
                                        resolve_ss  = false,
                                        tol_tucker  = 1e-4);

# + jupyter={"outputs_hidden": true} tags=[]
display(v2, root=1, thresh=.1)


# + tags=[]
 e3, v3 = FermiCG.block_sparse_tucker(v2, cluster_ops, clustered_ham,
                                        max_iter    = 20,
                                        max_iter_pt = 200,
                                        nbody       = 4,
                                        H0          = "Hcmf",
                                        thresh_var  = 1e-1,
                                        thresh_foi  = 1e-5,
                                        thresh_pt   = 1e-4,
                                        ci_conv     = 1e-5,
                                        do_pt       = true,                          
                                        resolve_ss  = false,
                                        tol_tucker  = 1e-4);

# + jupyter={"outputs_hidden": true} tags=[]
# rotate integrals to new cmf optimized basis
ints = FermiCG.orbital_rotation(ints_original,U);

Da = da1
Db = db1 
max_roots = 50

#
# form Cluster data
cluster_bases = FermiCG.compute_cluster_eigenbasis(ints, clusters, verbose=1,
                                                   max_roots=max_roots,
                                                   init_fspace=init_fspace,
                                                   rdm1a=Da, rdm1b=Db,
                                                    delta_elec=2)

clustered_ham = FermiCG.extract_ClusteredTerms(ints, clusters)

cluster_ops = FermiCG.compute_cluster_ops(cluster_bases, ints);

FermiCG.add_cmf_operators!(cluster_ops, cluster_bases, ints, Da, Db);


# + jupyter={"outputs_hidden": true} tags=[]
# Create BST state

v = FermiCG.BSTstate(clusters, FockConfig(init_fspace), cluster_bases, R=3)

FermiCG.add_single_excitons!(v, FockConfig(init_fspace), cluster_bases)
FermiCG.randomize!(v)
FermiCG.orthonormalize!(v)
e_ci, v = FermiCG.tucker_ci_solve(v, cluster_ops, clustered_ham);


# + tags=[]
 e3_50, v3_50 = FermiCG.block_sparse_tucker(v, cluster_ops, clustered_ham,
                                        max_iter    = 20,
                                        max_iter_pt = 200,
                                        nbody       = 4,
                                        H0          = "Hcmf",
                                        thresh_var  = 1e-1,
                                        thresh_foi  = 1e-5,
                                        thresh_pt   = 1e-4,
                                        ci_conv     = 1e-5,
                                        do_pt       = true,                          
                                        resolve_ss  = false,
                                        tol_tucker  = 1e-4);

# + jupyter={"source_hidden": true} tags=[]
# rotate integrals to new cmf optimized basis
ints = FermiCG.orbital_rotation(ints_original,U);

Da = da1
Db = db1 
max_roots = 100

#
# form Cluster data
@time cluster_bases = FermiCG.compute_cluster_eigenbasis(ints, clusters, verbose=0,
                                                   max_roots=max_roots,
                                                   init_fspace=init_fspace,
                                                   rdm1a=Da, rdm1b=Db)

clustered_ham = FermiCG.extract_ClusteredTerms(ints, clusters)

@time cluster_ops = FermiCG.compute_cluster_ops(cluster_bases, ints);

FermiCG.add_cmf_operators!(cluster_ops, cluster_bases, ints, Da, Db);


# + jupyter={"outputs_hidden": true, "source_hidden": true} tags=[]
# Create BST state

v = FermiCG.BSTstate(clusters, FockConfig(init_fspace), cluster_bases, R=3)

FermiCG.add_single_excitons!(v, FockConfig(init_fspace), cluster_bases)
FermiCG.randomize!(v)
FermiCG.orthonormalize!(v)
e_ci, v = FermiCG.tucker_ci_solve(v, cluster_ops, clustered_ham);


# + jupyter={"source_hidden": true} tags=[]
 e3_50, v3_50 = FermiCG.block_sparse_tucker(v, cluster_ops, clustered_ham,
                                        max_iter    = 20,
                                        max_iter_pt = 200,
                                        nbody       = 4,
                                        H0          = "Hcmf",
                                        thresh_var  = 1e-1,
                                        thresh_foi  = 1e-5,
                                        thresh_pt   = 1e-4,
                                        ci_conv     = 1e-5,
                                        do_pt       = true,                          
                                        resolve_ss  = false,
                                        tol_tucker  = 1e-4);
# -



StarcoderdataJupyterScriptsDedupFiltered