from functools import reduce
import unittest
import numpy
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo
from pyscf import fci

norb = 6
nelec = 6
na = fci.cistring.num_strings(norb, nelec//2)
numpy.random.seed(1)
ci0 = numpy.random.random((na,na))
ci0 = ci0 + ci0.T
rdm1, rdm2 = fci.direct_spin1.make_rdm12(ci0, norb, nelec)

class KnowValues(unittest.TestCase):
    def test_rdm3(self):
        dm3ref = make_dm3_o0(ci0, norb, nelec)
        dm1, dm2, dm3 = fci.rdm.make_dm123('FCI3pdm_kern_spin0', ci0, ci0, norb, nelec)
        self.assertTrue(numpy.allclose(dm3ref, dm3))

        dm3a = reorder_dm123_o0(dm1, dm2, dm3, False)[2]
        dm3b = fci.rdm.reorder_dm123(dm1, dm2, dm3, False)[2]
        self.assertTrue(numpy.allclose(dm3a, dm3b))
        dm3 = dm3b
        fac = 1. / (nelec-2)
        self.assertTrue(numpy.allclose(rdm2, numpy.einsum('ijklmm->ijkl',dm3)*fac))
        self.assertTrue(numpy.allclose(rdm2, numpy.einsum('ijmmkl->ijkl',dm3)*fac))
        self.assertTrue(numpy.allclose(rdm2, numpy.einsum('mmijkl->ijkl',dm3)*fac))

        dm3 = fci.rdm.make_dm123('FCI3pdm_kern_sf', ci0, ci0, norb, nelec)[2]
        dm2 = fci.direct_spin1.make_rdm12(ci0, norb, nelec, reorder=False)[1]
        self.assertTrue(numpy.allclose(dm2, numpy.einsum('mmijkl->ijkl',dm3)/nelec))

        numpy.random.seed(2)
        na = fci.cistring.num_strings(norb, 5)
        nb = fci.cistring.num_strings(norb, 3)
        ci1 = numpy.random.random((na,nb))
        dm3ref = make_dm3_o0(ci1, norb, (5,3))
        dm3 = fci.rdm.make_dm123('FCI3pdm_kern_sf', ci1, ci1, norb, (5,3))[2]
        self.assertTrue(numpy.allclose(dm3ref, dm3))

    def test_dm4(self):
        dm4ref = make_dm4_o0(ci0, norb, nelec)
        dm4 = fci.rdm.make_dm1234('FCI4pdm_kern_sf', ci0, ci0, norb, nelec)[3]
        self.assertTrue(numpy.allclose(dm4ref, dm4))

        numpy.random.seed(2)
        na = fci.cistring.num_strings(norb, 5)
        nb = fci.cistring.num_strings(norb, 3)
        ci1 = numpy.random.random((na,nb))
        dm4ref = make_dm4_o0(ci1, norb, (5,3))
        dm1, dm2, dm3, dm4 = fci.rdm.make_dm1234('FCI4pdm_kern_sf', ci1, ci1, norb, (5,3))
        self.assertTrue(numpy.allclose(dm4ref, dm4))
        self.assertTrue(numpy.allclose(dm3, numpy.einsum('ppmnijkl->mnijkl',dm4)/8))
        self.assertTrue(numpy.allclose(dm3, numpy.einsum('mnppijkl->mnijkl',dm4)/8))
        self.assertTrue(numpy.allclose(dm3, numpy.einsum('mnijppkl->mnijkl',dm4)/8))
        self.assertTrue(numpy.allclose(dm3, numpy.einsum('mnijklpp->mnijkl',dm4)/8))

        dm3a, dm4a = reorder_dm1234_o0(dm1, dm2, dm3, dm4, False)[2:]
        dm4b = fci.rdm.reorder_dm1234(dm1, dm2, dm3, dm4, False)[3]
        self.assertTrue(numpy.allclose(dm4a, dm4b))
        self.assertTrue(numpy.allclose(dm3a, numpy.einsum('ppmnijkl->mnijkl',dm4b)/5))
        self.assertTrue(numpy.allclose(dm3a, numpy.einsum('mnppijkl->mnijkl',dm4b)/5))
        self.assertTrue(numpy.allclose(dm3a, numpy.einsum('mnijppkl->mnijkl',dm4b)/5))
        self.assertTrue(numpy.allclose(dm3a, numpy.einsum('mnijklpp->mnijkl',dm4b)/5))

    def test_tdm2(self):
        dm1 = numpy.einsum('ij,ijkl->kl', ci0, _trans1(ci0, norb, nelec))
        self.assertTrue(numpy.allclose(rdm1, dm1))

        dm2 = numpy.einsum('ij,ijklmn->klmn', ci0, _trans2(ci0, norb, nelec))
        dm2 = fci.rdm.reorder_rdm(rdm1, dm2)[1]
        self.assertTrue(numpy.allclose(rdm2,dm2))

        na = ci0.shape[0]
        numpy.random.seed(1)
        ci = numpy.random.random((na,na))
        ci1 = numpy.random.random((na,na))
        dm1, dm2 = fci.direct_spin1.trans_rdm12(ci, ci1, norb, nelec)
        numpy.random.seed(2)
        self.assertAlmostEqual(numpy.dot(dm2.flatten(),numpy.random.random(dm2.size)),
                               3790.8867819690477, 7)
        self.assertTrue(numpy.allclose(dm2, dm2.transpose(2,3,0,1)))

        t1 = _trans1(ci1, norb, nelec)
        t2 = _trans2(ci1, norb, nelec)
        dm1a = numpy.einsum('ij,ijpq->pq', ci, t1)
        dm2a = numpy.einsum('ij,ijpqrs->pqrs', ci, t2)
        self.assertTrue(numpy.allclose(dm1a, dm1))
        dm1a, dm2a = fci.rdm.reorder_rdm(dm1a, dm2a)
        self.assertTrue(numpy.allclose(dm2a,dm2a.transpose(2,3,0,1)))

# (6o,6e)   ~ 4MB
# (8o,8e)   ~ 153MB
# (10o,10e) ~ 4.8GB
# t2(*,ij,kl) = E_i^j E_k^l|0>
def _trans2(fcivec, norb, nelec):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelecb = nelec//2
    else:
        neleca, nelecb = nelec
    link_indexa = fci.cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = fci.cistring.gen_linkstr_index(range(norb), nelecb)
    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    fcivec = fcivec.reshape(na,nb)
    t1 = _trans1(fcivec, norb, nelec)
    t2 = numpy.zeros((na,nb,norb,norb,norb,norb))
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            t2[str1,:,a,i] += sign * t1[str0]
    for k in range(na):
        for str0, tab in enumerate(link_indexb):
            for a, i, str1, sign in tab:
                t2[k,str1,a,i] += sign * t1[k,str0]
    return t2
def _trans1(fcivec, norb, nelec):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelecb = nelec//2
    else:
        neleca, nelecb = nelec
    link_indexa = fci.cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = fci.cistring.gen_linkstr_index(range(norb), nelecb)
    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    fcivec = fcivec.reshape(na,nb)
    t1 = numpy.zeros((na,nb,norb,norb))
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            t1[str1,:,a,i] += sign * fcivec[str0]
    for k in range(na):
        for str0, tab in enumerate(link_indexb):
            for a, i, str1, sign in tab:
                t1[k,str1,a,i] += sign * fcivec[k,str0]
    return t1

#
# NOTE: this rdm3 is defined as
# rdm3(p,q,r,s,t,u) = <p^+ q r^+ s t^+ u>
def make_dm3_o0(fcivec, norb, nelec):
    # <0|p^+ q r^+ s|i> <i|t^+ u|0>
    t1 = _trans1(fcivec, norb, nelec)
    t2 = _trans2(fcivec, norb, nelec)
    na, nb = t1.shape[:2]
    rdm3 = numpy.dot(t1.reshape(na*nb,-1).T, t2.reshape(na*nb,-1))
    return rdm3.reshape((norb,)*6).transpose(1,0,2,3,4,5)

def make_dm4_o0(fcivec, norb, nelec):
    # <0|p^+ q r^+ s|i> <i|t^+ u|0>
    t2 = _trans2(fcivec, norb, nelec)
    na, nb = t2.shape[:2]
    rdm4 = numpy.dot(t2.reshape(na*nb,-1).T, t2.reshape(na*nb,-1))
    return rdm4.reshape((norb,)*8).transpose(3,2,1,0,4,5,6,7)

# <p^+ q r^+ s t^+ u> => <p^+ r^+ t^+ u s q>
# rdm2 is <p^+ q r^+ s>
def reorder_dm123_o0(rdm1, rdm2, rdm3, inplace=True):
    rdm1, rdm2 = fci.rdm.reorder_rdm(rdm1, rdm2, inplace)
    if not inplace:
        rdm3 = rdm3.copy()
    norb = rdm1.shape[0]
    for p in range(norb):
        for q in range(norb):
            for s in range(norb):
                rdm3[p,q,q,s] += -rdm2[p,s]
            for u in range(norb):
                rdm3[p,q,:,:,q,u] += -rdm2[p,u]
            for s in range(norb):
                rdm3[p,q,:,s,s,:] += -rdm2[p,q]
    for q in range(norb):
        for s in range(norb):
            rdm3[:,q,q,s,s,:] += -rdm1
    return rdm1, rdm2, rdm3

# <p^+ q r^+ s t^+ u w^+ v> => <p^+ r^+ t^+ w^+ v u s q>
# rdm2, rdm3 are the (reordered) standard 2-pdm and 3-pdm
def reorder_dm1234_o0(rdm1, rdm2, rdm3, rdm4, inplace=True):
    rdm1, rdm2, rdm3 = fci.rdm.reorder_dm123(rdm1, rdm2, rdm3, inplace)
    if not inplace:
        rdm4 = rdm4.copy()
    norb = rdm1.shape[0]
    delta = numpy.eye(norb)
    rdm4 -= numpy.einsum('qv,pwrstu->pqrstuvw', delta, rdm3)
    rdm4 -= numpy.einsum('sv,pqrwtu->pqrstuvw', delta, rdm3)
    rdm4 -= numpy.einsum('uv,pqrstw->pqrstuvw', delta, rdm3)
    rdm4 -= numpy.einsum('qt,pursvw->pqrstuvw', delta, rdm3)
    rdm4 -= numpy.einsum('st,pqruvw->pqrstuvw', delta, rdm3)
    rdm4 -= numpy.einsum('qr,pstuvw->pqrstuvw', delta, rdm3)
    rdm4 -= numpy.einsum('qr,sv,pwtu', delta, delta, rdm2)
    rdm4 -= numpy.einsum('qr,uv,pstw', delta, delta, rdm2)
    rdm4 -= numpy.einsum('qt,uv,pwrs', delta, delta, rdm2)
    rdm4 -= numpy.einsum('qt,sv,purw', delta, delta, rdm2)
    rdm4 -= numpy.einsum('st,qv,pwru', delta, delta, rdm2)
    rdm4 -= numpy.einsum('st,uv,pqrw', delta, delta, rdm2)
    rdm4 -= numpy.einsum('qr,st,puvw', delta, delta, rdm2)
    rdm4 -= numpy.einsum('qr,st,uv,pw->pqrstuvw', delta, delta, delta, rdm1)
    return rdm1, rdm2, rdm3, rdm4


if __name__ == "__main__":
    print("Full Tests for fci.rdm")
    unittest.main()


