import numpy as np
import scipy.linalg as slg

from frankenstein import scf, mp
from frankenstein.tools.io_utils import dumpMat
from frankenstein.tools.tensor_utils import ao2mo_xform_nosymm


def get_Vovov(V, C, nocc):
    # Vovov = np.einsum("pqrs,pi,qa,rj,sb->iajb", V, \
    #     C[:,:nocc],C[:,nocc:],C[:,:nocc],C[:,nocc:], optimize=True)
    nao = C.shape[0]
    nmo = C.shape[1]
    nvirt = nmo - nocc
    Vovov = np.zeros([nocc,nvirt,nocc,nvirt])
    Vaaov = np.zeros([nao,nao,nocc,nvirt])
    for p in range(nao):
        for q in range(p,nao):
            X = C[:,:nocc].T @ V[p,q] @ C[:,nocc:]
            if p == q:
                Vaaov[p,q] = X * 0.5
            else:
                Vaaov[p,q] = X
    for j in range(nocc):
        for b in range(nvirt):
            Y = Vaaov[:,:,j,b]
            Y += Y.T
            Vovov[:,:,j,b] = C[:,:nocc].T @ Y @ C[:,nocc:]
    del Vaaov

    return Vovov


def get_Diajb(moe, nocc):
    nao = moe.shape[0]
    nvirt = nao - nocc
    Diajb = moe[:nocc].reshape(nocc,1,1,1) + \
        moe[:nocc].reshape(nocc,1) - \
        moe[nocc:].reshape(nvirt,1,1) - \
        moe[nocc:]
    return Diajb


def get_mp2_t2(V, C, moe, nocc):
    nao = C.shape[0]
    nvirt = nao - nocc
    Vovov = get_Vovov(V, C, nocc)
    Diajb = get_Diajb(moe, nocc)
    t2 = Vovov / Diajb
    return t2


def get_emp2_from_eri(V, C, moe, nocc):
    Vovov = get_Vovov(V, C, nocc)
    t2 = get_mp2_t2(V, C, moe, nocc)
    return np.sum((2.*Vovov - Vovov.transpose(0,3,2,1)) * t2)


def get_emp2_from_t2(V, C, t2, nocc):
    t2 = t2.transpose(0,2,1,3)  # phys (oovv) --> chem (ovov)
    Vovov = np.einsum("pqrs,pi,qa,rj,sb->iajb", V, \
        C[:,:nocc],C[:,nocc:],C[:,:nocc],C[:,nocc:], optimize=True)
    return np.sum((2.*Vovov - Vovov.transpose(0,3,2,1)) * t2)


def get_mp2_rdm1(V, C, moe, nocc):
    t2 = get_mp2_t2(V, C, moe, nocc)
    t2_ = t2.transpose(0,3,2,1)
    dm1oo = -np.einsum("iakb,jakb->ij", t2,2.*t2-t2_)
    dm1vv = np.einsum("iajc,ibjc->ab", t2,2.*t2-t2_)
    return slg.block_diag(dm1oo,dm1vv)


def get_cphf_U(mf, v, vo_only=False):
    from cphf_utils import cphf_kernel, get_full_u
    u = cphf_kernel(mf, v)
    if vo_only:
        return u
    else:
        U = get_full_u(mf, v, u)
        return U


def get_scf_rdm1_mo_lr(mf, v):
    U = get_cphf_U(mf, v)
    nao = mf.nao
    nocc = mf.nocc
    rho = np.zeros([nao,nao])
    rho[np.diag_indices(nocc)] = 1.
    return U@rho + rho@U.T


def get_mp2_rdm1_mo_lr1(mf, v, rdm1_mo_mp2):
    U = get_cphf_U(mf, v)
    return U@rdm1_mo_mp2 + rdm1_mo_mp2@U.T


def get_Vovov_lr(V, C, U, nocc):
    def get_Vovov_gen(V, C1, C2, C3, C4, nocc):
        Vovov = np.einsum("pqrs,pi,qa,rj,sb->iajb", V, \
            C1[:,:nocc],C2[:,nocc:],C3[:,:nocc],C4[:,nocc:], optimize=True)
        return Vovov
    Vovov_lr1 = get_Vovov_gen(V,C@U,C,C,C,nocc)
    Vovov_lr2 = get_Vovov_gen(V,C,C@U,C,C,nocc)
    Vovov_lr = Vovov_lr1 + Vovov_lr1.transpose(2,3,0,1) + \
        Vovov_lr2 + Vovov_lr2.transpose(2,3,0,1)
    return Vovov_lr


def get_Vovov_lr_batch(Vavov, Voaov, C, U, nocc):
    CU = C @ U
    Vovov_lr1 = np.einsum("pajb,pi->iajb", Vavov, CU[:,:nocc], optimize=True)
    Vovov_lr2 = np.einsum("ipjb,pa->iajb", Voaov, CU[:,nocc:], optimize=True)
    Vovov_lr = Vovov_lr1 + Vovov_lr1.transpose(2,3,0,1) + \
        Vovov_lr2 + Vovov_lr2.transpose(2,3,0,1)
    return Vovov_lr


def get_moe_lr(V, C, v, U, nocc):
    nao = V.shape[0]
    nvirt = nao - nocc
    u = U[nocc:,:nocc]
    moe_lr = np.zeros(nao)
    g1 = np.einsum("pqrs,pj,qj,ra,si->jai", \
        V, C, C, C[:,nocc:], C[:,:nocc], optimize=True)
    g2 = np.einsum("pqrs,pj,qi,ra,sj->jia", \
        V, C, C[:,:nocc], C[:,nocc:], C, optimize=True)

    for p in range(nao):
        moe_lr[p] += C[:,p].T @ v @ C[:,p]
        moe_lr[p] += 4.*np.sum(g1[p]*u) - 2.*np.sum(g2[p].T*u)
    return moe_lr


def get_moe_lr_batch(V, C, vs, Us, nocc):
    nao = V.shape[0]
    nvirt = nao - nocc
    g1 = np.einsum("pqrs,pj,qj,ra,si->jai", \
        V, C, C, C[:,nocc:], C[:,:nocc], optimize=True)
    g2 = np.einsum("pqrs,pj,qi,ra,sj->jia", \
        V, C, C[:,:nocc], C[:,nocc:], C, optimize=True)

    moes_lr = []
    for v, U in zip(vs, Us):
        u = U[nocc:,:nocc]
        moe_lr = np.zeros(nao)
        for p in range(nao):
            moe_lr[p] += C[:,p].T @ v @ C[:,p]
            moe_lr[p] += 4.*np.sum(g1[p]*u) - 2.*np.sum(g2[p].T*u)
        moes_lr.append(moe_lr)
    return moes_lr


def get_mp2_rdm1_mo_lr2(mf, v):
    nao = mf.nao
    nocc = mf.nocc
    nvirt = nao - nocc
    V = mf.mol.V
    C = mf.mo_coeff
    moe = mf.mo_energy

    U = get_cphf_U(mf, v)
    t2 = get_mp2_t2(V, C, moe, nocc)
    Vovov_lr = get_Vovov_lr(V, C, U, nocc)
    Diajb = get_Diajb(moe, nocc)
    moe_lr = get_moe_lr(V, C, v, U, nocc)
    Diajb_lr = get_Diajb(moe_lr, nocc)
    t2_lr = (Vovov_lr - t2*Diajb_lr) / Diajb

    dm1oo_lr = -np.einsum("iakb,jakb->ij", t2_lr, 2.*t2-t2.transpose(0,3,2,1))
    dm1oo_lr += dm1oo_lr.T
    dm1vv_lr = np.einsum("iajc,ibjc->ab", t2_lr, 2.*t2-t2.transpose(0,3,2,1))
    dm1vv_lr += dm1vv_lr.T

    return slg.block_diag(dm1oo_lr, dm1vv_lr)


def get_tot_rdm1_mo_an(mf, v):
    V = mf.mol.V
    C = mf.mo_coeff
    moe = mf.mo_energy
    nao = mf.nao
    nocc = mf.nocc

    U = get_cphf_U(mf, v)

    # scf contribution
    rho = np.zeros([nao,nao])
    rho[np.diag_indices(nocc)] = 1.
    rdm1_scf_mo_lr = U@rho + rho@U.T

    # mp2 contribution 1
    Vovov = get_Vovov(V, C, nocc)
    Diajb = get_Diajb(moe, nocc)
    t2 = Vovov / Diajb

    rdm1_mo_mp2 = np.zeros([nao,nao])
    rdm1_mo_mp2[:nocc,:nocc] = -np.einsum("iakb,jakb->ij", \
        t2, 2.*t2-t2.transpose(0,3,2,1))
    rdm1_mo_mp2[nocc:,nocc:] = np.einsum("iajc,ibjc->ab", \
        t2, 2.*t2-t2.transpose(0,3,2,1))
    rdm1_mp2_mo_lr1 = U@rdm1_mo_mp2 + rdm1_mo_mp2@U.T

    # mp2 contribution 2
    Vovov_lr = get_Vovov_lr(V, C, U, nocc)
    moe_lr = get_moe_lr(V, C, v, U, nocc)
    Diajb_lr = get_Diajb(moe_lr, nocc)
    t2_lr = (Vovov_lr - t2*Diajb_lr) / Diajb

    rdm1_mp2_mo_lr2 = np.zeros([nao,nao])
    rdm1_mp2_mo_lr2[:nocc,:nocc] = \
        -np.einsum("iakb,jakb->ij", t2_lr, 2.*t2-t2.transpose(0,3,2,1))
    rdm1_mp2_mo_lr2[nocc:,nocc:] = \
        np.einsum("iajc,ibjc->ab", t2_lr, 2.*t2-t2.transpose(0,3,2,1))
    rdm1_mp2_mo_lr2 += rdm1_mp2_mo_lr2.T

    # collect contributions
    rdm1_tot_mo_lr = rdm1_scf_mo_lr + rdm1_mp2_mo_lr1 + rdm1_mp2_mo_lr2

    return rdm1_tot_mo_lr


def get_tot_rdm1_mo_an_batch(mf, vs, Us):
    V = mf.mol.V
    C = mf.mo_coeff
    moe = mf.mo_energy
    nao = mf.nao
    nocc = mf.nocc
    nv = len(vs)

    # intermediates used by all Us/vs
    rho = np.zeros([nao,nao])
    rho[np.diag_indices(nocc)] = 1.

    Vovov = get_Vovov(V, C, nocc)
    Diajb = get_Diajb(moe, nocc)
    t2 = Vovov / Diajb
    rdm1_mo_mp2 = np.zeros([nao,nao])
    rdm1_mo_mp2[:nocc,:nocc] = -np.einsum("iakb,jakb->ij", \
        t2, 2.*t2-t2.transpose(0,3,2,1))
    rdm1_mo_mp2[nocc:,nocc:] = np.einsum("iajc,ibjc->ab", \
        t2, 2.*t2-t2.transpose(0,3,2,1))

    moes_lr = get_moe_lr_batch(V, C, vs, Us, nocc)
    Vavov = np.einsum("pqrs,qa,rj,sb->pajb", V, \
        C[:,nocc:], C[:,:nocc], C[:,nocc:], optimize=True)
    Voaov = np.einsum("pqrs,pi,rj,sb->iqjb", V, \
        C[:,:nocc], C[:,:nocc], C[:,nocc:], optimize=True)

    # determine lr for each U/v
    rdm1s_tot_mo_lr = [None for i in range(nv)]
    ind = -1
    for v, U, moe_lr in zip(vs, Us, moes_lr):
        ind += 1

        # scf contribution
        rdm1_scf_mo_lr = U@rho + rho@U.T

        # mp2 contribution 1
        rdm1_mp2_mo_lr1 = U@rdm1_mo_mp2 + rdm1_mo_mp2@U.T

        # mp2 contribution 2
        Vovov_lr = get_Vovov_lr_batch(Vavov, Voaov, C, U, nocc)
        Diajb_lr = get_Diajb(moe_lr, nocc)
        t2_lr = (Vovov_lr - t2*Diajb_lr) / Diajb

        rdm1_mp2_mo_lr2 = np.zeros([nao,nao])
        rdm1_mp2_mo_lr2[:nocc,:nocc] = \
            -np.einsum("iakb,jakb->ij", t2_lr, 2.*t2-t2.transpose(0,3,2,1))
        rdm1_mp2_mo_lr2[nocc:,nocc:] = \
            np.einsum("iajc,ibjc->ab", t2_lr, 2.*t2-t2.transpose(0,3,2,1))
        rdm1_mp2_mo_lr2 += rdm1_mp2_mo_lr2.T

        # collect contributions
        rdm1s_tot_mo_lr[ind] = rdm1_scf_mo_lr + rdm1_mp2_mo_lr1 + rdm1_mp2_mo_lr2

    return rdm1s_tot_mo_lr


def get_tot_rdm1_mo_fd(h, V, nocc, v, eps=1.E-2):
    def get_tot_rdm1_ao(h, V, nocc, D0):
        mf = scf.RHF(h=h, V=V, nocc=nocc, verbose="mute")
        mf.kernel(D0=D0, verbose="mute")
        mmp = mp.MP2(mf)
        mmp.kernel(rdm_level=1)
        return mmp.rdm1

    mf = scf.RHF(h=h, V=V, nocc=nocc)
    mf.kernel(verbose="mute")
    C = mf.mo_coeff
    rdm1_ao_ff = get_tot_rdm1_ao(h+2*eps*v, V, nocc, mf.rdm1)
    rdm1_ao_f = get_tot_rdm1_ao(h+eps*v, V, nocc, mf.rdm1)
    rdm1_ao_b = get_tot_rdm1_ao(h-eps*v, V, nocc, mf.rdm1)
    rdm1_ao_bb = get_tot_rdm1_ao(h-2*eps*v, V, nocc, mf.rdm1)
    rdm1_ao_lr_fd = (-rdm1_ao_ff + 8.*rdm1_ao_f - \
        8.*rdm1_ao_b + rdm1_ao_bb) / 12. / eps
    rdm1_mo_lr_fd = C.T @ rdm1_ao_lr_fd @ C
    return rdm1_mo_lr_fd


def test_mp2_t2():
    np.random.seed(17)
    nao = 10
    nocc = 4
    nvirt = nao - nocc
    from frankenstein.tools.lat_utils import get_random_hV
    h, V = get_random_hV(nao)

    mf = scf.RHF(h=h, V=V, nocc=nocc)
    mf.kernel(verbose="mute")
    C = mf.mo_coeff

    v = np.zeros([nao,nao])
    n0 = 3
    shift = 0

    vs = []
    for i in range(n0):
        for j in range(i,n0):
            v[i+shift,j+shift] = v[j+shift,i+shift] = 1.
            vs.append(v.copy())
            v[i+shift,j+shift] = v[j+shift,i+shift] = 0.

    from cphf_utils import cphf_kernel_batch, get_full_u_batch
    us = cphf_kernel_batch(mf, vs)
    Us = get_full_u_batch(mf, vs, us)
    rdm1s_mo_tot_lr = get_tot_rdm1_mo_an_batch(mf, vs, Us)

    errs = np.zeros([n0,n0])
    errs_norm = np.zeros([n0,n0])
    errs_ao = np.zeros([n0,n0])
    ind = -1
    for i in range(n0):
        for j in range(i,n0):
            ind += 1
            v = vs[ind]

            rdm1_mo_tot_lr_fd = get_tot_rdm1_mo_fd(h, V, nocc, v)

            err = np.max(np.abs(rdm1s_mo_tot_lr[ind] - rdm1_mo_tot_lr_fd))
            err_norm = np.linalg.norm(rdm1s_mo_tot_lr[ind] - rdm1_mo_tot_lr_fd)
            err_ao = np.max(np.abs(C@(rdm1s_mo_tot_lr[ind] - \
                rdm1_mo_tot_lr_fd)@C.T))

            errs[i,j] = errs[j,i] = err
            errs_norm[i,j] = errs_norm[j,i] = err_norm
            errs_ao[i,j] = errs_ao[j,i] = err_ao

    dumpMat(errs, "MO basis, max error", fmt="%.1E")
    dumpMat(errs_norm, "MO basis, rms error", fmt="%.1E")
    dumpMat(errs_ao, "AO basis, rms error", fmt="%.1E")


##################################
# The real thing starts here
##################################
class MP2_ERIS:
    def __init__(self, mf):
        nao = mf.nao
        nocc = mf.nocc
        nvirt = nao - nocc
        V = mf.mol.V
        C = mf.mo_coeff
        Co = C[:,:nocc]
        Cv = C[:,nocc:]
        moe = mf.mo_energy

        self.nao = nao
        self.nocc = nocc
        self.nvirt = nvirt
        self.C = C
        self.moe = moe

        self.ovov = np.einsum("pqrs,pi,qa,rj,sb->iajb", V, \
            Co, Cv, Co, Cv, optimize=True)
        self.vovo = self.ovov.transpose(1,0,3,2)
        self.vvoo = np.einsum("pqrs,pa,qb,ri,sj->abij", V, \
            Cv, Cv, Co, Co, optimize=True)
        self.oovo = np.einsum("pqrs,pi,qj,ra,sk->ijak", V, \
            Co, Co, Cv, Co, optimize=True)
        self.vvvo = np.einsum("pqrs,pa,qb,rc,si->abci", V, \
            Cv, Cv, Cv, Co, optimize=True)
        self.Dia = moe[:nocc].reshape(nocc,1) - moe[nocc:]
        self.Diajb = self.Dia.reshape(nocc,nvirt,1,1) + self.Dia
        self.t2 = self.ovov / self.Diajb


def get_mp2_rdm2(eris, dm1oo, dm1vv, dm1vo):
    nao = eris.nao
    nocc = eris.nocc
    t2 = eris.t2
    Ioo = np.eye(nocc)

    rdm2 = np.zeros([nao,nao,nao,nao])
    # scf contribution
    rdm2[:nocc,:nocc,:nocc,:nocc] = \
        2.*np.einsum("ij,kl->ijkl", Ioo, Ioo) - \
        np.einsum("il,kj->ijkl", Ioo, Ioo)
    # oooo, direct + exchange
    rdm2[:nocc,:nocc,:nocc,:nocc] += \
        2.*(np.einsum("ij,kl->ijkl", dm1oo, Ioo) + \
        np.einsum("ij,kl->ijkl", Ioo, dm1oo)) - (\
        np.einsum("il,kj->ijkl", dm1oo, Ioo) +
        np.einsum("il,kj->ijkl", Ioo, dm1oo))
    # vvoo/oovv, direct
    rdm2[nocc:,nocc:,:nocc,:nocc] = 2.*np.einsum("ab,kl->abkl", dm1vv, Ioo)
    rdm2[:nocc,:nocc,nocc:,nocc:] = 2.*np.einsum("kl,ab->klab", Ioo, dm1vv)
    # vvoo/oovv, exchange
    rdm2[nocc:,:nocc,:nocc,nocc:] = -np.einsum("ab,kl->alkb", dm1vv, Ioo)
    rdm2[:nocc,nocc:,nocc:,:nocc] = -np.einsum("ab,kl->kbal", dm1vv, Ioo)
    # ovov/vovo
    rdm2[:nocc,nocc:,:nocc,nocc:] = 2.*t2 - t2.transpose(0,3,2,1)
    rdm2[nocc:,:nocc,nocc:,:nocc] = 2.*t2.transpose(1,0,3,2) - t2.transpose(1,2,3,0)
    # orbital relaxation contributions
    rdm2[nocc:,:nocc,:nocc,:nocc] += 2.*np.einsum("ai,kl->aikl", dm1vo, Ioo)
    rdm2[:nocc,nocc:,:nocc,:nocc] += 2.*np.einsum("ai,kl->iakl", dm1vo, Ioo)
    rdm2[:nocc,:nocc,nocc:,:nocc] += 2.*np.einsum("ai,kl->klai", dm1vo, Ioo)
    rdm2[:nocc,:nocc,:nocc,nocc:] += 2.*np.einsum("ai,kl->klia", dm1vo, Ioo)
    rdm2[nocc:,:nocc,:nocc,:nocc] -= np.einsum("al,ki->aikl", dm1vo, Ioo)
    rdm2[:nocc,:nocc,:nocc,nocc:] -= np.einsum("al,ki->lika", dm1vo, Ioo)
    rdm2[:nocc,:nocc,nocc:,:nocc] -= np.einsum("al,ki->klai", dm1vo, Ioo)
    rdm2[:nocc,nocc:,:nocc,:nocc] -= np.einsum("al,ki->kali", dm1vo, Ioo)

    return rdm2


def get_mp2_B_eris(eris):
    """B_aibj = delta_aibj (ei-ea) - 4*Vaibj + Vajbi + Vabji
    """
    nao = eris.nao
    nocc = eris.nocc
    nvirt = nao - nocc

    # xform
    Vvovo = eris.vovo
    Vvvoo = eris.vvoo

    # diag
    diag = eris.Dia.T.reshape(nvirt*nocc)

    # compute
    B = (-4.*Vvovo + Vvovo.transpose(0,3,2,1) + \
        Vvvoo.transpose(0,2,1,3)).reshape(nvirt*nocc,nvirt*nocc)
    B[np.diag_indices(nvirt*nocc)] += diag

    return B


def get_mp2_X_eris(eris, dm1oo, dm1vv):
    nao = eris.nao
    nocc = eris.nocc
    nvirt = nao - nocc

    Voovo = eris.oovo
    Vvvvo = eris.vvvo
    t2 = eris.t2
    X = np.einsum("kj,kjai->ai", dm1oo, 2.*Voovo-Voovo.transpose(0,3,2,1)) + \
        np.einsum("bc,bcai->ai", dm1vv, 2.*Vvvvo-Vvvvo.transpose(2,1,0,3)) + \
        np.einsum("icjb,acbj->ai", t2, \
            2.*Vvvvo-Vvvvo.transpose(0,2,1,3), optimize=True) - \
        np.einsum("kajb,kibj->ai", t2, \
            2.*Voovo-Voovo.transpose(3,1,2,0), optimize=True)
    X = X.reshape(nvirt*nocc)

    return X


def cphf_kernel_batch(eris, B, vs, thresh=1E8):
    nao = eris.nao
    nocc = eris.nocc
    nvirt = nao - nocc
    C = eris.C
    moe = eris.moe
    nv = len(vs)

    # compute all rhs's
    Qs = np.zeros([nvirt*nocc, nv])
    for i in range(nv):
        Qs[:,i] = (C[:,nocc:].T@vs[i]@C[:,:nocc]).reshape(nvirt*nocc)

    # vo blocks
    us = np.linalg.solve(B, Qs)

    # full U's
    Voovo = eris.oovo
    Vvvvo = eris.vvvo

    Us = [None for i in range(nv)]
    for i in range(nv):
        vmo = C.T @ vs[i] @ C

        uvo = us[:,i].reshape(nvirt, nocc)
        U = np.zeros([nao,nao])
        U[nocc:,:nocc] = uvo
        U[:nocc,nocc:] = -uvo.T

        # occupied-occupied block
        # ( Qjk + sum_ai (4*Vjkai-Vjiak-Vikaj)*uai ) / (ek - ej)
        denom_oo = moe[:nocc].reshape(nocc,1) - moe[:nocc]
        denom_oo[np.diag_indices(nocc)] = 1.
        if np.sum(np.abs(denom_oo**-1) > thresh) > 0:
            raise RuntimeError
        U[:nocc,:nocc] = -(vmo[:nocc,:nocc] + \
            np.einsum("jkai,ai->jk", 4.*Voovo - \
            Voovo.transpose(0,3,2,1) - \
            Voovo.transpose(3,1,2,0), \
            uvo)) / denom_oo
        U[np.diag_indices(nocc)] = 0.

        # virtual-virtual block
        # ( Qbc + sum_ai (4*Vbcai-Vacbi-Vbaci)*uai ) / (ec - eb)
        denom_vv = moe[nocc:].reshape(nvirt,1) - moe[nocc:]
        denom_vv[np.diag_indices(nvirt)] = 1.
        if np.sum(np.abs(denom_vv**-1) > thresh) > 0:
            raise RuntimeError
        uvv = -(vmo[nocc:,nocc:] + \
            np.einsum("bcai,ai->bc", 4.*Vvvvo - \
            Vvvvo.transpose(2,1,0,3) - \
            Vvvvo.transpose(0,2,1,3), \
            uvo)) / denom_vv
        uvv[np.diag_indices(nvirt)] = 0.
        U[nocc:,nocc:] = uvv

        Us[i] = U

    return Us


def get_Vvvoo_lr_batch(Vavoo, Vvvoa, C, U, nocc):
    Vvvoo_lr1 = np.einsum("pbij,pa->abij", Vavoo, (C@U)[:,nocc:], optimize=True)
    Vvvoo_lr2 = np.einsum("abis,sj->abij", Vvvoa, (C@U)[:,:nocc], optimize=True)
    Vvvoo_lr = Vvvoo_lr1 + Vvvoo_lr1.transpose(1,0,2,3) + \
        Vvvoo_lr2 + Vvvoo_lr2.transpose(0,1,3,2)
    return Vvvoo_lr


def get_mp2_B_lr_eris(mf, v, U, moe_lr, Vvvoo_lr, Vvovo_lr):
    nao = mf.nao
    nocc = mf.nocc
    nvirt = nao - nocc

    # diag
    diag_lr = (moe_lr[:nocc]-moe_lr[nocc:].reshape(nvirt,1)).reshape(nvirt*nocc)

    # compute
    B_lr = (-4.*Vvovo_lr + Vvovo_lr.transpose(0,3,2,1) + \
        Vvvoo_lr.transpose(0,2,1,3)).reshape(nvirt*nocc,nvirt*nocc)
    B_lr[np.diag_indices(nvirt*nocc)] += diag_lr

    return B_lr


def get_Voovo_lr_batch(Vaovo, Vooao, Voova, C, U, nocc):
    Voovo_lr1 = np.einsum("pjak,pi->ijak", Vaovo, (C@U)[:,:nocc], optimize=True)
    Voovo_lr2 = np.einsum("ijrk,ra->ijak", Vooao, (C@U)[:,nocc:], optimize=True)
    Voovo_lr3 = np.einsum("ijas,sk->ijak", Voova, (C@U)[:,:nocc], optimize=True)
    Voovo_lr = Voovo_lr1 + Voovo_lr1.transpose(1,0,2,3) + \
        Voovo_lr2 + Voovo_lr3
    return Voovo_lr


def get_Vvvvo_lr_batch(Vavvo, Vvvao, Vvvva, C, U, nocc):
    Vvvvo_lr1 = np.einsum("pbci,pa->abci", Vavvo, (C@U)[:,nocc:], optimize=True)
    Vvvvo_lr2 = np.einsum("abri,rc->abci", Vvvao, (C@U)[:,nocc:], optimize=True)
    Vvvvo_lr3 = np.einsum("abcs,si->abci", Vvvva, (C@U)[:,:nocc], optimize=True)
    Vvvvo_lr = Vvvvo_lr1 + Vvvvo_lr1.transpose(1,0,2,3) + \
        Vvvvo_lr2 + Vvvvo_lr3
    return Vvvvo_lr


def get_mp2_X_lr_eris(eris, v, U, t2_lr, Voovo_lr, Vvvvo_lr, dm1oovv, dm1oovv_lr):
    nao = eris.nao
    nocc = eris.nocc
    nvirt = nao - nocc

    t2 = eris.t2

    Voovo = eris.oovo
    Vvvvo = eris.vvvo

    dm1oo = dm1oovv[:nocc,:nocc]
    dm1vv = dm1oovv[nocc:,nocc:]
    dm1oo_lr = dm1oovv_lr[:nocc,:nocc]
    dm1vv_lr = dm1oovv_lr[nocc:,nocc:]

    X_lr = np.einsum("kj,kjai->ai", dm1oo_lr, 2.*Voovo-Voovo.transpose(0,3,2,1)) + \
        np.einsum("bc,bcai->ai", dm1vv_lr, 2.*Vvvvo-Vvvvo.transpose(2,1,0,3)) + \
        np.einsum("icjb,acbj->ai", t2_lr, \
            2.*Vvvvo-Vvvvo.transpose(0,2,1,3), optimize=True) - \
        np.einsum("kajb,kibj->ai", t2_lr, \
            2.*Voovo-Voovo.transpose(3,1,2,0), optimize=True) + \
        np.einsum("kj,kjai->ai", dm1oo, 2.*Voovo_lr-Voovo_lr.transpose(0,3,2,1)) + \
        np.einsum("bc,bcai->ai", dm1vv, 2.*Vvvvo_lr-Vvvvo_lr.transpose(2,1,0,3)) + \
        np.einsum("icjb,acbj->ai", t2, \
            2.*Vvvvo_lr-Vvvvo_lr.transpose(0,2,1,3), optimize=True) - \
        np.einsum("kajb,kibj->ai", t2, \
            2.*Voovo_lr-Voovo_lr.transpose(3,1,2,0), optimize=True)
    X_lr = X_lr.reshape(nvirt*nocc)

    return X_lr


def mp2_rdm(mf, eris, ret_rdm2=False, vs=None, timing=False):
    nao = eris.nao
    nocc = eris.nocc
    nvirt = nao - nocc
    Diajb = eris.Diajb
    t2 = eris.t2
    V = mf.mol.V
    C = mf.mo_coeff
    Co = C[:,:nocc]
    Cv = C[:,nocc:]

    ###  make mp2 rdm1 ###
    # oo/vv
    if timing:
        start = time.time()
    dm1oo = -np.einsum("iakb,jakb->ij", t2, 2.*t2-t2.transpose(0,3,2,1), optimize=True)
    dm1vv = np.einsum("iajc,ibjc->ab", t2, 2.*t2-t2.transpose(0,3,2,1), optimize=True)
    if timing:
        end = time.time()
        print(" ** timing-dm1 oo/vv : {:.6f}".format(end-start), flush=True)

    # vo
    if timing:
        start = time.time()
    B = get_mp2_B_eris(eris)
    X = get_mp2_X_eris(eris, dm1oo, dm1vv)
    dm1vo = np.linalg.solve(B.T, X).reshape(nvirt,nocc)
    if timing:
        end = time.time()
        print(" ** timing-dm1 ov    : {:.6f}".format(end-start), flush=True)

    # collect
    dm1_mp2 = np.block([[dm1oo, dm1vo.T], [dm1vo, dm1vv]])

    # rhf contributions
    dm1_rhf = np.zeros([nao,nao])
    dm1_rhf[np.diag_indices(nocc)] = 1.

    # xform to AO basis
    dm1_ao = C@(dm1_mp2+dm1_rhf)@C.T

    ###  make mp2 rdm2 (rhf contributions included) ###
    if ret_rdm2:
        # directly form AO basis rdm2
        if timing:
            start = time.time()
        dm2_ao = ao2mo_xform_nosymm(\
            get_mp2_rdm2(eris, dm1oo, dm1vv, dm1vo), C.T)
        if timing:
            end = time.time()
            print(" ** timing-dm2 ao    : {:.6f}".format(end-start), flush=True)
    else:
        dm2_ao = None

    # return here if rdm1_lr is not requested
    if vs is None:
        return dm1_ao, dm2_ao

    ### cphf ###
    nv = len(vs)

    if timing:
        start = time.time()
    Us = cphf_kernel_batch(eris, B, vs)
    if timing:
        end = time.time()
        print(" ** timing-cphf Us   : {:.6f}".format(end-start), flush=True)

    ### compute common intermediates ###
    if timing:
        start = time.time()
    # for Diajb_lr
    moes_lr = get_moe_lr_batch(V, C, vs, Us, nocc)
    # for Vovov_lr
    Vavov = np.einsum("pqrs,qa,rj,sb->pajb", V, Cv, Co, Cv, optimize=True)
    Voaov = np.einsum("pqrs,pi,rj,sb->iqjb", V, Co, Co, Cv, optimize=True)
    # for Vvvoo_lr
    Vavoo = np.einsum("pqrs,qa,ri,sj->paij", V, Cv, Co, Co, optimize=True)
    Vvvoa = np.einsum("pqrs,pa,qb,ri->abis", V, Cv, Cv, Co, optimize=True)
    # for Voovo_lr
    Vaovo = Voaov.transpose(1,0,3,2)
    Vooao = np.einsum("pqrs,pi,qj,sk->ijrk", V, Co, Co, Co, optimize=True)
    Voova = Vavoo.transpose(2,3,1,0)
    # for Vvvvo_lr
    Vavvo = np.einsum("pqrs,qa,rb,si->pabi", V, Cv, Cv, Co, optimize=True)
    Vvvao = Vvvoa.transpose(0,1,3,2)
    Vvvva = np.einsum("pqrs,pa,qb,rc->abcs", V, Cv, Cv, Cv, optimize=True)
    if timing:
        end = time.time()
        print(" ** timing-loop prep : {:.6f}".format(end-start), flush=True)

    ### compute rdm1 lr ###
    # oo/vv lr
    rhss_lr = np.zeros([nvirt*nocc,nv])
    dm1s_mo_lr = [None for i in range(nv)]
    ind = -1
    ttot = 0.
    for v, U, moe_lr in zip(vs, Us, moes_lr):
        ind += 1

        if timing:
            start = time.time()
        # scf, orbital response
        dm1_or_lr = U@dm1_rhf + dm1_rhf@U.T

        # mp2, orbital response
        dm1_or_lr += U@dm1_mp2 + dm1_mp2@U.T

        # mp2, oo/vv
        Vovov_lr = get_Vovov_lr_batch(Vavov, Voaov, C, U, nocc)

        Diajb_lr = get_Diajb(moe_lr, nocc)

        t2_lr = (Vovov_lr - t2*Diajb_lr) / Diajb

        dm1oo_lr = -np.einsum("iakb,jakb->ij", \
            t2_lr, 2.*t2-t2.transpose(0,3,2,1))
        dm1vv_lr = np.einsum("iajc,ibjc->ab", \
            t2_lr, 2.*t2-t2.transpose(0,3,2,1))
        dm1oo_lr += dm1oo_lr.T
        dm1vv_lr += dm1vv_lr.T
        dm1oovv_lr = slg.block_diag(dm1oo_lr, dm1vv_lr)

        # mp2, vo
        Vvvoo_lr = get_Vvvoo_lr_batch(Vavoo, Vvvoa, C, U, nocc)
        B_lr = get_mp2_B_lr_eris(mf, v, U, moe_lr, Vvvoo_lr, Vovov_lr.transpose(1,0,3,2))
        Voovo_lr = get_Voovo_lr_batch(Vaovo, Vooao, Voova, C, U, nocc)
        Vvvvo_lr = get_Vvvvo_lr_batch(Vavvo, Vvvao, Vvvva, C, U, nocc)
        X_lr = get_mp2_X_lr_eris(eris, v, U, t2_lr, Voovo_lr, Vvvvo_lr, dm1_mp2, dm1oovv_lr)
        rhss_lr[:,ind] = X_lr-B_lr.T@dm1vo.ravel()

        # collect contributions
        dm1s_mo_lr[ind] = dm1_or_lr + dm1oovv_lr

        if timing:
            end = time.time()
            ttot += end-start

    if timing:
        print(" ** timing-per v     : {:.6f}".format(ttot/float(nv)), flush=True)
        print(" ** timing-{:4d} vs   : {:.6f}".format(nv, ttot), flush=True)

    if timing:
        start = time.time()
    dm1vos_lr = np.linalg.solve(B.T, rhss_lr)
    if timing:
        end = time.time()
        print(" ** timing-cpmp2     : {:.6f}".format(end-start), flush=True)

    for ind in range(nv):
        dm1vo_lr = dm1vos_lr[:,ind].reshape(nvirt,nocc)
        dm1s_mo_lr[ind][nocc:,:nocc] += dm1vo_lr
        dm1s_mo_lr[ind][:nocc,nocc:] += dm1vo_lr.T

    # return dm1_ao, dm2, dm1s_ao_lr
    return dm1_ao, dm1s_mo_lr


def get_mp2_lr_rdm1_lr_mo_fd(mf, v, eps=1.E-3):
    h = mf.h
    V = mf.mol.V
    nao = mf.nao
    nocc = mf.nocc
    nvirt = nao - nocc
    C = mf.mo_coeff

    mf_f = scf.RHF(h=h+eps*v, V=V, nocc=nocc)
    mf_f.kernel(verbose="mute")
    C_f = mf_f.mo_coeff
    Pmo_lr_f = get_mp2_lr_rdm1_mo(mf_f)
    Pao_lr_f = C_f @ Pmo_lr_f @ C_f.T + mf_f.rdm1
    mf_b = scf.RHF(h=h-eps*v, V=V, nocc=nocc)
    mf_b.kernel(verbose="mute")
    C_b = mf_b.mo_coeff
    Pmo_lr_b = get_mp2_lr_rdm1_mo(mf_b)
    Pao_lr_b = C_b @ Pmo_lr_b @ C_b.T + mf_b.rdm1
    dPao_lr = (Pao_lr_f-Pao_lr_b) / 2. / eps
    dPmo_lr = C.T @ dPao_lr @ C

    return dPmo_lr


import time
def test_mp2_lr_rdm1_lr_batch():
    np.random.seed(1)
    nao = 7
    nocc = 3
    nvirt = nao - nocc
    from frankenstein.tools.lat_utils import get_random_hV
    h, V = get_random_hV(nao)

    """
    geom = "geom/1.xyz"
    basis = "sto-3g"
    nocc = 21
    from frankenstein.tools.lat_utils import get_mol
    h, V, e_nuc = get_mol(geom, basis, True)
    nao = h.shape[0]
    nvirt = nao - nocc
    print("nao = {:d}".format(nao))
    """

    mf = scf.RHF(h=h, V=V, nocc=nocc)
    mf.kernel(verbose="mute")
    C = mf.mo_coeff

    n0 = 3
    vs = []
    for p in range(n0):
        for q in range(p, n0):
            v = np.zeros([nao,nao])
            v[p,q] = v[q,p] = 1.
            vs.append(v)

    start = time.time()
    eris = MP2_ERIS(mf)
    end = time.time()
    print("timing-eris : {:.6f}".format(end-start))

    # check ecorr
    mmp = mp.MP2(mf)
    mmp.kernel()
    dm1, dm2 = mp2_rdm(mf, eris, ret_rdm2=True)
    hmo = C.T @ h @ C
    Vmo = np.einsum("pqrs,pi,qj,rk,sl->ijkl", V,C,C,C,C, optimize=True)
    e1 = 2.*np.sum(hmo*dm1)
    e2 = np.sum(Vmo*dm2)
    print("etot = {: .10f}".format(mmp.e_tot))
    print("e1   = {: .10f}".format(e1))
    print("e2   = {: .10f}".format(e2))
    print("e12  = {: .10f}".format(e1+e2))
    print("etot match?", np.allclose(mmp.e_tot, e1+e2))

    start = time.time()
    dm1, dm2, dPmos = mp2_rdm(mf, eris, vs=vs)
    end = time.time()
    print("timing-dPmo : {:.6f}".format(end-start))

    pq = -1
    errs = np.zeros([n0,n0])
    for p in range(n0):
        for q in range(p,n0):
            pq += 1

            start = time.time()
            dPmo_fd = get_mp2_lr_rdm1_lr_mo_fd(mf, vs[pq])
            end = time.time()
            print("timing-dPfd : {:.6f}".format(end-start))

            err = np.linalg.norm(dPmo_fd-dPmos[pq])
            errs[p,q] = errs[q,p] = err
    dumpMat(errs, fmt="%.3E")


if __name__ == "__main__":
    # test_mp2_t2()
    test_mp2_lr_rdm1_lr_batch()
