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


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

wt = w.WickTheorem()

expr = wt.contract(w.commutator(F0, T2), 0, 4)
expr += wt.contract(V1, 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)

import psi4
import forte
import forte.utils
from forte import forte_options
import numpy as np
from math import isclose

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


# setup xyz geometry for linear H4
geometry = """
H 0.0 0.0 0.0
H 0.0 0.0 1.0
H 0.0 0.0 2.0
H 0.0 0.0 3.0
symmetry c1
"""

(Escf, psi4_wfn) = forte.utils.psi4_scf(geometry,
                                        basis='sto-3g',
                                        reference='rhf',
                                        options={'E_CONVERGENCE' : 1.e-12})

psi4.set_options({'mp2_type': 'conv','freeze_core': True})
Emp2_psi4 = psi4.energy('mp2')
print(f"RHF energy:       {Escf:.12f} Eh")
print(f"MP2 energy:       {Emp2_psi4:.12f} Eh")
print(f"MP2 corr. energy: {Emp2_psi4 - Escf:.12f} Eh")

"""
Explanation: Compute the Hartree–Fock and MP2 energy
End of explanation
"""


# Define the orbital spaces
mo_spaces = {'RESTRICTED_DOCC': [2],'RESTRICTED_UOCC': [2]}

# pass Psi4 options to Forte
options = psi4.core.get_options()
options.set_current_module('FORTE')
forte_options.get_options_from_psi4(options)

# Grab the number of MOs per irrep
nmopi = psi4_wfn.nmopi()
# Grab the point group symbol (e.g. "C2V")
point_group = psi4_wfn.molecule().point_group().symbol()
# create a MOSpaceInfo object
mo_space_info = forte.make_mo_space_info_from_map(nmopi, point_group,mo_spaces, [])
# make a ForteIntegral object
ints = forte.make_ints_from_psi4(psi4_wfn, forte_options, mo_space_info)

"""
Explanation: Prepare integrals for Forte
End of explanation
"""


occmos = mo_space_info.corr_absolute_mo('RESTRICTED_DOCC')
virmos = mo_space_info.corr_absolute_mo('RESTRICTED_UOCC')
allmos = mo_space_info.corr_absolute_mo('CORRELATED')
nocc = 2 * len(occmos)
nvir = 2 * len(virmos)

"""
Explanation: Define orbital spaces and dimensions
End of explanation
"""


# Build the Fock matrix blocks
F = {'oo': forte.spinorbital_oei(ints,occmos, occmos),
     'vv': forte.spinorbital_oei(ints,virmos, virmos),
     'ov': forte.spinorbital_oei(ints,occmos, virmos)}

# OO block
v = forte.spinorbital_tei(ints,occmos,occmos,occmos, occmos)
F['oo'] += np.einsum('piqi->pq', v)

# VV block
v = forte.spinorbital_tei(ints,virmos, occmos, virmos, occmos)
F['vv'] += np.einsum('piqi->pq', v)

# OV block
v = forte.spinorbital_tei(ints, occmos, occmos, virmos, occmos)
F['ov'] += np.einsum('piqi->pq', v)

# Build the diagonal orbital energies
Fdiag = {'oo': np.diag(F['oo']), 'vv': np.diag(F['vv'])}
F0 = {'oo' : np.diag(Fdiag['oo']),'vv' : np.diag(Fdiag['vv']) }

# Build the two-electron integrals
V = {}
V["oovv"] = forte.spinorbital_tei(ints,occmos,occmos,virmos,virmos)

"""
Explanation: Build the Fock matrix and the zeroth-order Fock matrix
End of explanation
"""


d2 = np.zeros((nocc,nocc,nvir,nvir))

fo = Fdiag['oo']
fv = Fdiag['vv']
for i in range(nocc):
    for j in range(nocc):
        for a in range(nvir):
            for b in range(nvir):
                si = i % 2
                sj = j % 2
                sa = a % 2
                sb = b % 2
                if si == sj == sa == sb:
                    d2[i][j][a][b] = 1.0 / (fo[i] + fo[j] - fv[a] - fv[b])
                if si == sa and sj == sb and si != sj:
                    d2[i][j][a][b] = 1.0 / (fo[i] + fo[j] - fv[a] - fv[b])
                if si == sb and sj == sa and si != sj:
                    d2[i][j][a][b] = 1.0 / (fo[i] + fo[j] - fv[a] - fv[b])

# Compute the MP2 correlation energy
Emp2 = 0.0
for i in range(nocc):
    for j in range(nocc):
        for a in range(nvir):
            for b in range(nvir):
                Emp2 += 0.25 * V["oovv"][i][j][a][b] ** 2 / (fo[i] + fo[j] - fv[a] - fv[b])
print(f"MP2 corr. energy: {Emp2:.12f} Eh")

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(T2,R,d2):
    T2["oovv"] += np.einsum("ijab,ijab->ijab",R,d2)
    
def compute_energy(T2,V):
    energy = 0.25 * np.einsum('ijab,ijab->', V["oovv"], T2["oovv"])
    return energy

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

maxiter = 10
for i in range(maxiter):
    Roovv = evaluate_residual(F0,T2,V)
    Roovv = antisymmetrize_residual(Roovv)
    update_amplitudes(T2,Roovv,d2)
    Emp2_wicked = compute_energy(T2,V)
    
    # check for convergence
    norm_R = np.linalg.norm(Roovv)
    print(f"{i} {Emp2_wicked:+.12f} {norm_R}") 
    
    if norm_R < 1.0e-9:
        break
    
print(f"MP2 corr. energy: {Emp2_wicked:+.12f} Eh")
print(f"Err corr. energy: {Emp2_wicked - Emp2:+.12f} Eh")

assert isclose(Emp2_wicked,Emp2)

"""
Explanation: Build the MP denominators
End of explanation
"""


