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

"""
Explanation: MP2 theory for a closed-shell reference
In this notebook we will use wicked to generate equations for the MP2 method using an orbital-invariant formalism
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, 2)

import wicked as w

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"])

"""
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
"""


T2op = w.op("T2", ["v+ v+ o o"])
F0op = w.op("F0", ["o+ o", "v+ v"])
F1op = w.op("F1", ["o+ v", "v+ o"])
V1op = w.utils.gen_op("V", 2, "ov", "ov")

wt = w.WickTheorem()

expr = wt.contract(w.commutator(F0op, T2op), 0, 4)
expr += wt.contract(V1op, 0, 4)

mbeq = expr.to_manybody_equation("R")

# code generation
code = [
    "def evaluate_residual(F0,T2,V):",
    "    # contributions to the residual",
    "    Roovv = np.zeros((nocc,nocc,nvir,nvir))",
]
for eq in mbeq["oo|vv"]:
    contraction = eq.compile("einsum")
    code.append(f"    {contraction}")
code.append(f"    return Roovv")

funct = "\n".join(code)

print(f"Defined the function:\n\n{funct}\n")
exec(funct)

def antisymmetrize_residual(Roovv):
    # antisymmetrize the residual
    Roovv_anti = np.zeros((nocc, nocc, nvir, nvir))
    Roovv_anti += np.einsum("ijab->ijab", Roovv)
    Roovv_anti -= np.einsum("ijab->jiab", Roovv)
    Roovv_anti -= np.einsum("ijab->ijba", Roovv)
    Roovv_anti += np.einsum("ijab->jiba", Roovv)
    return Roovv_anti


def update_amplitudes_mp2(T2, R, invD):
    T2["oovv"] += np.einsum("ijab,ijab->ijab", R, invD["oovv"])


def compute_energy(T2, V):
    energy = 0.25 * np.einsum("ijab,ijab->", V["oovv"], T2["oovv"])
    return energy

Emp2 = -0.06623692118909  # reference value from psi4

T2 = {}
T2["oovv"] = np.zeros((nocc, nocc, nvir, nvir))

F0 = {"oo": H["oo"], "vv": H["vv"]}

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

start = time.perf_counter()

maxiter = 10
for i in range(maxiter):
    Roovv = evaluate_residual(F0, T2, H)
    Roovv = antisymmetrize_residual(Roovv)
    update_amplitudes_mp2(T2, Roovv, invD)
    Emp2_wicked = compute_energy(T2, H)
    Etot_w = Eref + Emp2_wicked

    # check for convergence
    norm_R = np.linalg.norm(Roovv)
    print(f"{i:3d}    {Etot_w:+.12f}    {Emp2_wicked:+.12f}    {norm_R:e}")
    if norm_R < 1.0e-9:
        break

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

print("-" * len(header))

print(f"MP2 corr. energy: {Emp2_wicked:+.12f} Eh")
print(f"MP2 corr. energy: {Emp2:+.12f} Eh (psi4)")
print(f"Err corr. energy: {Emp2_wicked - Emp2:+.12f} Eh")

assert np.isclose(Emp2_wicked, Emp2)

"""
Explanation: Here we define the operator and derive the residual equations
End of explanation
"""


