import wicked as w
from IPython.display import display, Math, Latex

def latex(expr):
    """Function to render any object that has a member latex() function"""
    display(Math(expr.latex()))
    
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']) 
wt = w.WickTheorem()

"""
Explanation: Generating and implementing many-body equations
Preliminaries
Once more, let's start by importing wick&d and defining a Slater determinant reference
End of explanation
"""


E0 = w.op("E_0",[""])
F = w.utils.gen_op('f',1,'ov','ov')
V = w.utils.gen_op('v',2,'ov','ov')
H = E0 + F + V
T = w.op("t",["v+ o", "v+ v+ o o"])
Hbar = w.bch_series(H,T,2)
expr = wt.contract(Hbar,0,0)
expr

"""
Explanation: Generating equations for fully contracted terms
In the previous notebook, we computed the coupled cluster energy expression
\begin{equation}
E = \langle \Phi | e^{-\hat{T}} \hat{H} e^{\hat{T}} | \Phi \rangle
= E_0 + \sum_{i}^\mathbb{O} \sum_{a}^\mathbb{V} f^{a}{i} t^{i}{a} + 
\frac{1}{4} \sum_{ij}^\mathbb{O} \sum_{ab}^\mathbb{V}
 (t^{i j}{a b} + 2 t^{i}{a} t^{j}{b}) v^{a b}{i j} 
\end{equation}
with the following code
End of explanation
"""


mbeq = expr.to_manybody_equations('r')
mbeq

"""
Explanation: First we convert the expression derived into a set of equations. You get back a dictionary that shows all the components to the equations. The vertical bar (|) in the key separates the lower (left) and upper (right) indices in the resulting expression
End of explanation
"""


for eq in mbeq['|']:
    print(eq.compile('einsum'))

"""
Explanation: Converting equations to code
From the equations generated above, you can get tensor contractions by calling the compile function on each individual term in the equations. Here we generate python code that uses numpy's einsum function to evaluate contractions. To use this code you will need to import einsum
python
from numpy import einsum
and you will need to define a dictionary of tensors (f["vo"],v["vvoo"],t["ov"],...) of appropriate dimensions:
End of explanation
"""


F = w.utils.gen_op('f',1,'ov','ov')
T1 = w.op("t",["v+ o"])
expr = wt.contract(w.commutator(F,T1),2,2)
latex(expr)

"""
Explanation: Many-body equations
Suppose we want to compute the contributions to the coupled cluster residual equations
\begin{equation}
r^{i}{a} = \langle \Phi| { \hat{a}^\dagger{i} \hat{a}a } [\hat{F},\hat{T}_1] | \Phi \rangle
\end{equation}
Wick&d can compute this quantity using the corresponding many-body representation of the operator $[\hat{F},\hat{T}_1]$.
If you expand the operator $[\hat{F},\hat{T}_1]$ into its second quantized operator components we can identify a particle-hole excitation term:
\begin{equation}
[\hat{F},\hat{T}_1] = g^{j}{b} { \hat{a}^\dagger_{b} \hat{a}j } + \cdots
\end{equation}
From this expression we see that the residual $r{a}^{i}$ is precisely the quantity we need to evaluate since
\begin{equation}
r^{i}{a} = \langle \Phi| { \hat{a}^\dagger{i} \hat{a}a } [\hat{F},\hat{T}_1] | \Phi \rangle
 = g^{j}{b}  \langle \Phi| { \hat{a}^\dagger_{i} \hat{a}a }  { \hat{a}^\dagger{b} \hat{a}j }  | \Phi \rangle = g^{i}{a}
\end{equation}
where in the last step we applied Wick's theorem to evaluate the expectation value.
Let's start by computing $[\hat{F},\hat{T}_1]$ with Wick's theorem:
End of explanation
"""


mbeq = expr.to_manybody_equations('g')
print(mbeq)

"""
Explanation: Next, we call to_manybody_equations to generate many-body equations
End of explanation
"""


mbeq_ov = mbeq["o|v"]

for eq in mbeq_ov:
    latex(eq)

"""
Explanation: Out of all the terms, we select the terms that multiply the excitation operator ${ \hat{a}^\dagger_{a} \hat{a}_i }$ ("o|v")
End of explanation
"""


for eq in mbeq_ov:
    print(eq.compile('einsum'))

"""
Explanation: Lastly, we can compile these equations into code
End of explanation
"""


T2 = w.op("t", ["v+ v+ o o"])
Vovov = w.op("v", ["o+ v+ v o"])
expr = wt.contract(w.commutator(Vovov, T2), 4, 4)
latex(expr)

"""
Explanation: Antisymmetrization of uncontracted operator indices
To gain efficiency, Wick&d treats contractions involving inequivalent lines in a special way. Consider the following term contributing to the CCSD doubles amplitude equations that arises from $[\hat{V}\mathrm{ovov},\hat{T}_2]$ (see the sixth term in Eq. (153) of Crawford and Schaefer, https://doi.org/10.1002/9780470125915.ch2)
\begin{equation}
r^{ij}{ab} \leftarrow \langle \Phi| { \hat{a}^\dagger_{i}\hat{a}^\dagger_{j} \hat{a}b \hat{a}_a } [\hat{V}\mathrm{ovov},\hat{T}2] | \Phi \rangle = - P(ij)P(ab) \sum{kc} \langle kb \| jc \rangle t^{ik}_{ac}
\end{equation}
where $P(pq)$ is an antisymmetric permutation operator [$P(pq)f(p,q) = f(p,q) - f(q,p)$].
This expression corresponds to a single diagram, but algebraically it consists of four terms obtained by index permutations $i \leftrightarrow j$ and $a \leftrightarrow b$, so that the residual is antisymmetric with respect to separate permutations of upper and lower indices.
Let's first take a look at what happens when we apply Wick's theorem with wick&d to the quantity $[\hat{V}_\mathrm{ovov},\hat{T}_2]$
End of explanation
"""


for eq in expr.to_manybody_equations('g')['oo|vv']:
    print(eq.compile('einsum'))

"""
Explanation: In wick&d the two-body part of $[\hat{V}\mathrm{ovov},\hat{T}_2]$ gives us only a single term
\begin{equation}
[\hat{V}\mathrm{ovov},\hat{T}2]\text{2-body} = - \sum_{abcijk} \langle kb \| jc \rangle t^{ik}{ac} { \hat{a}^{ab}{ij} } = \sum_{abij} g^{ij}{ab}  { \hat{a}^{ab}{ij} }
\end{equation}
where the tensor $g^{ij}{ab}$ is defined as
\begin{equation}
g^{ij}{ab}  = -\sum_{kc} \langle kb \| jc \rangle t^{ik}{ac}
\end{equation}
Note that contrary to $r^{ij}_{ab}$, the tensor $g^{ij}{ab}$ does not have any specific index symmetry. In other words, you need to enforce the antisymmetry.
<!-- In particular, the many-body tensors generated by wick&d are not guaranteed to be antisymmetric, i -->

This quantity is related to the CCSD residual contribution reported above in the following way
\begin{equation}
r^{ij}{ab} \leftarrow \langle \Phi| { \hat{a}^\dagger{i}\hat{a}^\dagger_{j} \hat{a}b \hat{a}_a } [\hat{V}\mathrm{ovov},\hat{T}2] | \Phi \rangle = g^{ij}{ab} - g^{ji}{ab} - g^{ij}{ba} + g^{ji}{ba} = P(ij)P(ab) g^{ij}{ab}
\end{equation}
Therefore, this example shows an important distinction between the traditional projective equation (which yields $P(ij)P(ab) g^{ij}{ab}$) vs. the many-body approach (which yields $g^{ij}{ab}$).
How are the difference between these two approaches reconciled in practic? When you solve the many-body equations, you must enforce the antisymmetry of the equations, which means that the residual contribution should be written as
\begin{equation}
\sum_{abij} g^{ij}{ab}  { \hat{a}^{ab}{ij} }
= \frac{1}{4} \sum_{abij} (P(ij)P(ab) g^{ij}{ab})  { \hat{a}^{ab}{ij} }
\end{equation}
The factor $\frac{1}{4}$ now brings this term in a form consistent with the prefactor we associate with the operator $ { \hat{a}^{ab}_{ij} }$.
When you ask wick&d to compile the many-body equation we again get a single term
End of explanation
"""


