"""
This script tests the cisd module.
"""


import os
import numpy as np
import time
from frankenstein import scf, ci
from frankenstein.tools.lat_utils import get_mol


def test_cisd_ch4():
    nocc = 5
    basis = "sto-3g"
    if __name__ == "__main__":
        xyz = "geom/ch4.zmat"
    else:
        xyz = os.path.dirname(__file__) + "/geom/ch4.zmat"
    h, V, e_nuc = get_mol(xyz, basis, True)
    mf = scf.RHF(h=h, V=V, nocc=nocc, e_nuc=e_nuc)
    mf.kernel(verbose="mute")
    mci = ci.CISD(mf)
    start = time.time()
    mci.kernel()
    end = time.time()
    dt1 = end - start

    print("SCF energy                : % .10f" % (mf.e_tot))
    print("CISD correlation energy   : % .10f" % (mci.e_corr))
    print("CISD total energy         : % .10f" % (mci.e_tot))

    # check with pyscf
    try:
        from pyscf import gto
        from pyscf import scf as pscf
        from pyscf import ci as pci
    except:
        raise RuntimeError("Failed to import gto, scf and "
            "ci from pyscf.")

    with open(xyz, "r") as f:
        f.readline()
        f.readline()
        atom = ""
        for line in f:
            atom += line

    mol = gto.M(atom=atom, basis=basis)
    pmf = pscf.RHF(mol).run()
    pmci = pci.CISD(pmf).run()

    assert(np.allclose(mf.e_tot, pmf.e_tot))
    assert(np.allclose(mci.e_corr, pmci.e_corr))
    assert(np.allclose(mci.e_tot, pmci.e_tot))

    # check rdm
    mci.kernel(rdm_level=2, ret_ci=True)
    e1 = 2. * np.sum(h*mci.rdm1)
    e2 = np.einsum("pqrs,pqrs->", V, mci.rdm2, optimize=True)
    etot = e1+e2+e_nuc
    print("CISD 1e energy            : % .10f" % e1)
    print("CISD 2e energy            : % .10f" % e2)
    print("CISD total energy from rdm: % .10f" % etot)

    assert(np.allclose(mci.e_tot, etot))

    # re-run with t1 and t2, should be faster
    start = time.time()
    mci.kernel(ci0=mci.ci)
    end = time.time()
    dt2 = end - start

    print("dt1: % .3f sec" % dt1)
    print("dt2: % .3f sec" % dt2)

    assert(dt2 < dt1)

if __name__ == "__main__":
    test_cisd_ch4()
