"""Automatic derivation of resitrict CCDF theory.

This script serves as an example of using drudge for complex symbolic
manipulations.  The derivation here is going to be based on the approach in GE
Scuseria et al, J Chem Phys 89 (1988) 7382 (10.1063/1.455269).

"""

from pyspark import SparkConf, SparkContext
from sympy import IndexedBase, Rational, symbols

from drudge import SpinOneHalfPartHoleDrudge, Vec, UP, DOWN, TimeStamper

# Environment setting up.

conf = SparkConf().setAppName('rccsd')
ctx = SparkContext(conf=conf)
dr = SpinOneHalfPartHoleDrudge(ctx)
dr.full_simplify = False

p = dr.names
c_dag = p.c_dag
c_ = p.c_
a, b, c, d = p.V_dumms[:4]
i, j, k, l = p.O_dumms[:4]

#
# Cluster excitation operator
# 
# Here, we first write the cluster excitation operator in terms of the
# unitary group generator.  Then they will be substituted by their fermion
# operator definition.
#

t1 = IndexedBase('t^1')
t2 = IndexedBase('t^2')
e_ = Vec('E')

cluster_e = dr.einst(
    t1[a, i] * e_[a, i] +
    Rational(1, 2) * t2[a, b, i, j] * e_[a, i] * e_[b, j]
)

cluster = cluster_e.subst(
    e_[a, i], c_dag[a, UP] * c_[i, UP] + c_dag[a, DOWN] * c_[i, DOWN]
)
dr.set_n_body_base(t2, 2)
cluster = cluster.simplify()
cluster.cache()

#
# Similarity transform of the Hamiltonian
# 

stamper = TimeStamper()

curr = dr.ham
h_bar = dr.ham
for order in range(4):
    curr = (curr | cluster).simplify() * Rational(1, order + 1)
    stamper.stamp('Commutator order {}'.format(order + 1), curr)
    h_bar += curr
    continue

h_bar = h_bar.simplify()
h_bar.repartition(cache=True)
stamper.stamp('H-bar assembly', h_bar)

en_eqn = h_bar.eval_fermi_vev().simplify()
stamper.stamp('Energy equation', en_eqn)

dr.wick_parallel = 1

e_dag = Vec(r'E^\dagger')
beta, gamma, u, v = symbols('beta gamma u v')
spin = symbols('spin')
e_dag_def = dr.sum((spin, UP, DOWN), c_dag[i, spin] * c_[a, spin])
projs = [e_dag_def.act(e_dag[a, i], p) for p in [
    e_dag[beta, u],
    e_dag[beta, u] * e_dag[gamma, v]
]]

#
# Dump the result to a simple report.
#

amp_eqns = []
for order, proj in enumerate(projs):
    eqn = (proj * h_bar).eval_fermi_vev().simplify()
    stamper.stamp('T{} equation'.format(order + 1), eqn)
    amp_eqns.append(eqn)
    continue

with dr.report('rCCSD.html', 'restricted CCSD theory') as rep:
    rep.add('Energy equation', en_eqn)
    for i, v in enumerate(amp_eqns):
        rep.add(r'\(T^{}\) amplitude equation'.format(i + 1), v)
        continue

@@ -417,27 +417,25 @@
 , DOWN, 
-TimeStamper
+Stopwatch
 %0A%0A# Envi
@@ -888,17 +888,16 @@
 on.%0A#%0A%0At
-1
  = Index
@@ -909,49 +909,103 @@
 e('t
-%5E1
 ')%0A
-t2 = IndexedBase('t%5E2')%0Ae_ = Vec('E'
+e_ = dr.define(%0A    Vec('E')%5Ba, i%5D, c_dag%5Ba, UP%5D * c_%5Bi, UP%5D + c_dag%5Ba, DOWN%5D * c_%5Bi, DOWN%5D%0A
 )%0A%0Ac
@@ -1010,18 +1010,16 @@
 %0Acluster
-_e
  = dr.ei
@@ -1028,17 +1028,16 @@
 t(%0A    t
-1
 %5Ba, i%5D *
@@ -1070,17 +1070,16 @@
 , 2) * t
-2
 %5Ba, b, i
@@ -1112,107 +1112,8 @@
 %0A)%0A%0A
-cluster = cluster_e.subst(%0A    e_%5Ba, i%5D, c_dag%5Ba, UP%5D * c_%5Bi, UP%5D + c_dag%5Ba, DOWN%5D * c_%5Bi, DOWN%5D%0A)%0A
 dr.s
@@ -1128,17 +1128,16 @@
 y_base(t
-2
 , 2)%0Aclu
@@ -1233,27 +1233,27 @@
 %0A%0Ast
-amper = TimeStamper
+opwatch = Stopwatch
 ()%0A%0A
@@ -1374,27 +1374,28 @@
 )%0A    st
-amper.stamp
+opwatch.tock
 ('Commut
@@ -1523,27 +1523,28 @@
 True)%0Ast
-amper.stamp
+opwatch.tock
 ('H-bar 
@@ -1607,27 +1607,28 @@
 ify()%0Ast
-amper.stamp
+opwatch.tock
 ('Energy
@@ -1674,34 +1674,8 @@
  1%0A%0A
-e_dag = Vec(r'E%5E%5Cdagger')%0A
 beta
@@ -1748,14 +1748,54 @@
 _dag
-_def =
+ = dr.define(%0A    Vec(r'E%5E%5Cdagger')%5Ba, i%5D,%0A   
  dr.
@@ -1846,16 +1846,18 @@
  spin%5D)%0A
+)%0A
 projs = 
@@ -1861,48 +1861,8 @@
  = %5B
-e_dag_def.act(e_dag%5Ba, i%5D, p) for p in %5B
 %0A   
@@ -1916,17 +1916,16 @@
 ma, v%5D%0A%5D
-%5D
 %0A%0A#%0A# Du
@@ -2075,19 +2075,20 @@
   st
-amper.stamp
+opwatch.tock
 ('T%7B

Update the rCCSD example to latest drudge