"""
ROHF-MP2
"""


import numpy as np

from pyscf import scf, mp, lib


def lorentz_regularization(denom, alpha, deg=2):
    return denom + alpha**deg / denom


def kernel1(mp, mo_energy, mo_coeff, eris, with_t2, thr_zero, alpha, verbose):

    if mo_energy is None or mo_coeff is None:
        moidx = mp.get_frozen_mask()
        mo_coeff = None
        mo_energy = (mp.mo_energy[0][moidx[0]], mp.mo_energy[1][moidx[1]])
    else:
        # For backward compatibility.  In pyscf-1.4 or earlier, mp.frozen is
        # not supported when mo_energy or mo_coeff is given.
        assert(mp.frozen is 0 or mp.frozen is None)

    if eris is None: eris = mp.ao2mo(mo_coeff)

    nocca, noccb = mp.get_nocc()
    nmoa, nmob = mp.get_nmo()
    nvira, nvirb = nmoa-nocca, nmob-noccb
    mo_ea, mo_eb = mo_energy
    eia_a = mo_ea[:nocca,None] - mo_ea[None,nocca:]
    eia_b = mo_eb[:noccb,None] - mo_eb[None,noccb:]

    if with_t2:
        dtype = eris.ovov.dtype
        t2aa = np.empty((nocca,nocca,nvira,nvira), dtype=dtype)
        t2ab = np.empty((nocca,noccb,nvira,nvirb), dtype=dtype)
        t2bb = np.empty((noccb,noccb,nvirb,nvirb), dtype=dtype)
        t2 = (t2aa,t2ab,t2bb)
    else:
        t2 = None

    emp2 = 0.0
    for i in range(nocca):
        eris_ovov = np.asarray(eris.ovov[i*nvira:(i+1)*nvira])
        eris_ovov = eris_ovov.reshape(nvira,nocca,nvira).transpose(1,0,2)
        denom = lib.direct_sum('a+jb->jab', eia_a[i], eia_a)
        t2i = eris_ovov.conj() / lorentz_regularization(denom, alpha)
        emp2 += np.einsum('jab,jab', t2i, eris_ovov) * .5
        emp2 -= np.einsum('jab,jba', t2i, eris_ovov) * .5
        if with_t2:
            t2aa[i] = t2i - t2i.transpose(0,2,1)

        eris_ovov = np.asarray(eris.ovOV[i*nvira:(i+1)*nvira])
        eris_ovov = eris_ovov.reshape(nvira,noccb,nvirb).transpose(1,0,2)
        denom = lib.direct_sum('a+jb->jab', eia_a[i], eia_b)
        if i == nocca-1 and np.abs(denom[-1,0,0]) < thr_zero:
            denom[-1,0,0] = 1.E20
        t2i = eris_ovov.conj() / lorentz_regularization(denom, alpha)
        emp2 += np.einsum('JaB,JaB', t2i, eris_ovov)
        if with_t2:
            t2ab[i] = t2i

    for i in range(noccb):
        eris_ovov = np.asarray(eris.OVOV[i*nvirb:(i+1)*nvirb])
        eris_ovov = eris_ovov.reshape(nvirb,noccb,nvirb).transpose(1,0,2)
        denom = lib.direct_sum('a+jb->jab', eia_b[i], eia_b)
        t2i = eris_ovov.conj() / lorentz_regularization(denom, alpha)
        emp2 += np.einsum('jab,jab', t2i, eris_ovov) * .5
        emp2 -= np.einsum('jab,jba', t2i, eris_ovov) * .5
        if with_t2:
            t2bb[i] = t2i - t2i.transpose(0,2,1)

    return emp2.real, t2


class ROMP2(mp.ump2.UMP2):

    def kernel(self, mo_energy=None, mo_coeff=None, eris=None, with_t2=True,
        thr_zero=1.E-10, alpha=0.0):
        return kernel1(self, mo_energy, mo_coeff, eris, with_t2, thr_zero,
            alpha, self.verbose)
