import time
import wicked as w
import numpy as np
from examples_helpers import *

"""
Explanation: CCSDT theory for a closed-shell reference
This notebook extends the spinorbital-CCSD notebook to compute CCSDT
End of explanation
"""


molecule = "sr-h6-sto-3g"

with open(f"{molecule}.npy", "rb") as f:
    Eref = np.load(f)
    nocc, nvir = np.load(f)
    H = np.load(f, allow_pickle=True).item()

invD = compute_inverse_denominators(H, nocc, nvir, 3)

"""
Explanation: Read calculation information (integrals, number of orbitals)
We start by reading information about the reference state, integrals, and denominators from the file sr-h6-sto-3g.npy. The variable H is a dictionary that holds the blocks of the Hamiltonian normal-ordered with respect to the Hartree–Fock determinant. invD similarly is a dictionary that stores the denominators $(\epsilon_i + \epsilon_j + \ldots - \epsilon_a - \epsilon_b - \ldots)^{-1}$.
End of explanation
"""


w.reset_space()
w.add_space("o", "fermion", "occupied", ["i", "j", "k", "l", "m", "n"])
w.add_space("v", "fermion", "unoccupied", ["a", "b", "c", "d", "e", "f"])

Top = w.op("T", ["v+ o", "v+ v+ o o", "v+ v+ v+ o o o"])
Hop = w.utils.gen_op("H", 1, "ov", "ov") + w.utils.gen_op("H", 2, "ov", "ov")
# the similarity-transformed Hamiltonian truncated to the four-nested commutator term
Hbar = w.bch_series(Hop, Top, 4)

"""
Explanation: Define orbital spaces and the Hamiltonian and cluster operators
End of explanation
"""


wt = w.WickTheorem()
expr = wt.contract(w.rational(1), Hbar, 0, 6)
mbeq = expr.to_manybody_equation("R")

"""
Explanation: In the following lines, we apply Wick's theorem to simplify the similarity-transformed Hamiltonian $\bar{H}$ computing all contributions ranging from operator rank 0 to 6 (triple substitutions).
Then we convert all the terms into many-body equations accumulated into the residual R.
End of explanation
"""


energy_eq = generate_equation(mbeq, 0, 0)
t1_eq = generate_equation(mbeq, 1, 1)
t2_eq = generate_equation(mbeq, 2, 2)
t3_eq = generate_equation(mbeq, 3, 3)

exec(energy_eq)
exec(t1_eq)
exec(t2_eq)
exec(t3_eq)

# show what do these functions look like
print(energy_eq)

"""
Explanation: Here we generate the CCSDT equations.
End of explanation
"""


Ecorr_ref = -0.108354659115  # from forte sparse implementation

T = {
    "ov": np.zeros((nocc, nvir)),
    "oovv": np.zeros((nocc, nocc, nvir, nvir)),
    "ooovvv": np.zeros((nocc, nocc, nocc, nvir, nvir, nvir)),
}

header = "Iter.     Energy [Eh]    Corr. energy [Eh]       |R|       "
print("-" * len(header))
print(header)
print("-" * len(header))

start = time.perf_counter()

maxiter = 50

for i in range(maxiter):
    # 1. compute energy and residuals
    R = {}
    Ecorr_w = evaluate_residual_0_0(H, T)
    Etot_w = Eref + Ecorr_w
    R["ov"] = evaluate_residual_1_1(H, T)
    Roovv = evaluate_residual_2_2(H, T)
    R["oovv"] = antisymmetrize_residual_2_2(Roovv, nocc, nvir)
    Rooovvv = evaluate_residual_3_3(H, T)
    R["ooovvv"] = antisymmetrize_residual_3_3(Rooovvv, nocc, nvir)

    # 2. amplitude update
    update_cc_amplitudes(T, R, invD, 3)

    # 3. check for convergence
    norm_R = np.sqrt(np.linalg.norm(R["ov"]) ** 2 + np.linalg.norm(R["oovv"]) ** 2)
    print(f"{i:3d}    {Etot_w:+.12f}    {Ecorr_w:+.12f}    {norm_R:e}")
    if norm_R < 1.0e-8:
        break

end = time.perf_counter()
t = end - start

print("-" * len(header))
print(f"CCSDT total energy                   {Etot_w:+.12f} [Eh]")
print(f"CCSDT correlation energy             {Ecorr_w:+.12f} [Eh]")
print(f"Reference CCSDT correlation energy   {Ecorr_ref:+.12f} [Eh]")
print(f"Error                               {Ecorr_w - Ecorr_ref:+.12e} [Eh]")
print(f"Timing                              {t:+.12e} [s]")
assert np.isclose(Ecorr_w, Ecorr_ref)

"""
Explanation: CCSDT algorithm
Here we code a simple loop in which we evaluate the energy and residuals of the CCSD equations and update the amplitudes
End of explanation
"""


