#!/usr/bin/env python3
## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab
## ---------------------------------------------------------------------
##
## Copyright (C) 2019 by the adcc authors
##
## This file is part of adcc.
##
## adcc is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published
## by the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## adcc is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with adcc. If not, see <http://www.gnu.org/licenses/>.
##
## ---------------------------------------------------------------------
import libadcc
import numpy as np

from .functions import direct_sum, evaluate, einsum, zeros_like
from .misc import cached_property, cached_member_function
from .ReferenceState import ReferenceState
from .OneParticleDensity import OneParticleDensity
from .TwoParticleDensity import TwoParticleDensity
from .NParticleOperator import product_trace, OperatorSymmetry
from .Intermediates import register_as_intermediate
from .timings import Timer, timed_member_call
from .MoSpaces import MoSpaces, split_spaces
from . import block as b


class LazyMp:
    def __init__(self, hf):
        """
        Initialise the class dealing with the Møller-Plesset ground state.
        """
        if isinstance(hf, libadcc.HartreeFockSolution_i):
            hf = ReferenceState(hf)
        if not isinstance(hf, ReferenceState):
            raise TypeError("hf needs to be a ReferenceState "
                            "or a HartreeFockSolution_i")
        self.reference_state: ReferenceState = hf
        self.mospaces: MoSpaces = hf.mospaces
        self.timer: Timer = Timer()
        self.has_core_occupied_space: bool = hf.has_core_occupied_space

    def __getattr__(self, attr):
        # Shortcut some quantities, which are needed most often
        if attr.startswith("t2") and len(attr) == 4:  # t2oo, t2oc, t2cc
            xxvv = b.__getattr__(attr[2:4] + "vv")
            return self.t2(xxvv)
        else:
            raise AttributeError

    @cached_member_function()
    def df(self, space):
        """Delta Fock matrix"""
        hf = self.reference_state
        s1, s2 = split_spaces(space)
        fC = hf.fock(s1 + s1).diagonal()
        fv = hf.fock(s2 + s2).diagonal()
        return direct_sum("-i+a->ia", fC, fv)

    @cached_member_function()
    def t2(self, space):
        """T2 amplitudes"""
        hf = self.reference_state
        sp = split_spaces(space)
        assert all(s == b.v for s in sp[2:])
        eia = self.df(sp[0] + b.v)
        ejb = self.df(sp[1] + b.v)
        return (
            hf.eri(space) / direct_sum("ia+jb->ijab", eia, ejb).symmetrise((2, 3))
        )

    @cached_member_function()
    def td2(self, space):
        """Return the T^D_2 term"""
        if space != b.oovv:
            raise NotImplementedError("T^D_2 term not implemented "
                                      f"for space {space}.")
        t2erit = self.t2eri(b.oovv, b.ov).transpose((1, 0, 2, 3))
        denom = direct_sum(
            'ia,jb->ijab', self.df(b.ov), self.df(b.ov)
        ).symmetrise(0, 1)
        return (
            + 4.0 * t2erit.antisymmetrise(2, 3).antisymmetrise(0, 1)
            - 0.5 * self.t2eri(b.oovv, b.vv)
            - 0.5 * self.t2eri(b.oovv, b.oo)
        ) / denom

    @cached_member_function()
    def tt2(self, space: str):
        """
        Return the second order MP triples amplitudes for the given space
        (e.g. o1o1o1v1v1v1).
        """
        if space != b.ooovvv:
            raise NotImplementedError("Second order MP triples amplitudes not "
                                      f"implemented for space {space}.")
        hf = self.reference_state
        df = self.df(b.ov)
        t2_1 = self.t2(b.oovv)
        # denom = a + b + c - i - j - k   //   df = i - a
        denom = - 1 * direct_sum(
            "ia,jkbc->ijkabc", df, direct_sum("jb,kc->jkbc", df, df)
        ).symmetrise(0, 1, 2).symmetrise(3, 4, 5)
        # prefactor of 9, because we have 9 terms each in the expression, while the
        # antisymmetrisation generates 36 terms and introduces a prefactor of 1/36.
        # Each of the 9 terms is therefore generated 4 times.
        # The scaling in the comments is given as: [comp_scaling] / [mem_scaling]
        numerator = (
            # N^7: O^3V^4 / N^6: O^3V^3
            + 9 * einsum('idab,jkcd->ijkabc', hf.ovvv, t2_1)
            # N^7: O^4V^3 / N^6: O^3V^3
            + 9 * einsum('ijla,klbc->ijkabc', hf.ooov, t2_1)
        ).antisymmetrise(0, 1, 2).antisymmetrise(3, 4, 5)
        return numerator / denom

    @cached_member_function()
    def t2eri(self, space, contraction):
        """
        Return the T2 tensor with ERI tensor contraction intermediates.
        These are called pi1 to pi7 in libadc.
        """
        hf = self.reference_state
        key = space + contraction
        expressions = {
            # space + contraction
            b.ooov + b.vv: ('ijbc,kabc->ijka', b.ovvv),
            b.ooov + b.ov: ('ilab,lkjb->ijka', b.ooov),
            b.oovv + b.oo: ('klab,ijkl->ijab', b.oooo),
            b.oovv + b.ov: ('jkac,kbic->ijab', b.ovov),
            b.oovv + b.vv: ('ijcd,abcd->ijab', b.vvvv),
            b.ovvv + b.oo: ('jkbc,jkia->iabc', b.ooov),
            b.ovvv + b.ov: ('ijbd,jcad->iabc', b.ovvv),
        }
        if key not in expressions:
            raise NotImplementedError("t2eri intermediate not implemented "
                                      f"for space '{space}' and contraction "
                                      f"'{contraction}'.")
        contraction_str, eri_block = expressions[key]
        return einsum(contraction_str, self.t2oo, hf.eri(eri_block))

    @cached_property
    @timed_member_call(timer="timer")
    def mp2_diffdm(self):
        """
        Return the MP2 difference density in the MO basis.
        """
        hf = self.reference_state
        ret = OneParticleDensity(self.mospaces,
                                 symmetry=OperatorSymmetry.HERMITIAN)
        # NOTE: the following 3 blocks are equivalent to the cvs_p0 intermediates
        # defined at the end of this file
        ret.oo = -0.5 * einsum("ikab,jkab->ij", self.t2oo, self.t2oo)
        ret.ov = -0.5 * (
            + einsum("ijbc,jabc->ia", self.t2oo, hf.ovvv)
            + einsum("jkib,jkab->ia", hf.ooov, self.t2oo)
        ) / self.df(b.ov)
        ret.vv = 0.5 * einsum("ijac,ijbc->ab", self.t2oo, self.t2oo)

        if self.has_core_occupied_space:
            # additional terms to "revert" CVS for ground state density
            ret.oo += -0.5 * einsum("iLab,jLab->ij", self.t2oc, self.t2oc)
            ret.ov += -0.5 * (
                + einsum("jMib,jMab->ia", hf.ocov, self.t2oc)
                + einsum("iLbc,Labc->ia", self.t2oc, hf.cvvv)
                + einsum("kLib,kLab->ia", hf.ocov, self.t2oc)
                + einsum("iMLb,LMab->ia", hf.occv, self.t2cc)
                - einsum("iLMb,LMab->ia", hf.occv, self.t2cc)
            ) / self.df(b.ov)
            ret.vv += (
                + 0.5 * einsum("IJac,IJbc->ab", self.t2cc, self.t2cc)
                + 1.0 * einsum("kJac,kJbc->ab", self.t2oc, self.t2oc)
            )
            # compute extra CVS blocks
            ret.cc = -0.5 * (
                + einsum("kIab,kJab->IJ", self.t2oc, self.t2oc)
                + einsum('LIab,LJab->IJ', self.t2cc, self.t2cc)
            )
            ret.oc = -0.5 * (
                + einsum("kIab,kjab->jI", self.t2oc, self.t2oo)
                + einsum("ILab,jLab->jI", self.t2cc, self.t2oc)
            )
            ret.cv = -0.5 * (
                - einsum("jIbc,jabc->Ia", self.t2oc, hf.ovvv)
                + einsum("jkIb,jkab->Ia", hf.oocv, self.t2oo)
                + einsum("jMIb,jMab->Ia", hf.occv, self.t2oc)
                + einsum("ILbc,Labc->Ia", self.t2cc, hf.cvvv)
                + einsum("kLIb,kLab->Ia", hf.occv, self.t2oc)
                + einsum("LMIb,LMab->Ia", hf.cccv, self.t2cc)
            ) / self.df(b.cv)
        ret.reference_state = self.reference_state
        return evaluate(ret)

    def density(self, level=2):
        """
        Return the MP density in the MO basis with all corrections
        up to the specified order of perturbation theory
        """
        if level in [0, 1]:
            return self.reference_state.density
        elif level == 2:
            return self.reference_state.density + self.mp2_diffdm
        else:
            raise NotImplementedError("Only densities for level 0, 1 and 2"
                                      " are implemented.")

    @cached_property
    @timed_member_call(timer="timer")
    def mp1_dm_correction_2p(self) -> TwoParticleDensity:
        """
        Return the two-particle MP1 difference density correction in the MO basis.
        """
        ret = TwoParticleDensity(self.mospaces,
                                 symmetry=OperatorSymmetry.HERMITIAN)
        ret.oovv = -1.0 * self.t2oo
        return ret.evaluate()

    @cached_property
    @timed_member_call(timer="timer")
    def mp2_dm_correction_2p(self) -> TwoParticleDensity:
        """
        Return the two-particle MP2 difference density correction in the MO basis.
        """
        hf: ReferenceState = self.reference_state
        ret = TwoParticleDensity(self.mospaces,
                                 symmetry=OperatorSymmetry.HERMITIAN)
        p0: OneParticleDensity = self.mp2_diffdm

        # constuct Kronecker Delta
        d_oo = zeros_like(hf.foo)
        d_oo.set_mask("ii", 1)

        ret.oooo = (
            + 4.0 * einsum("ik,jl->ijkl", p0.oo, d_oo)
            .antisymmetrise(0, 1).antisymmetrise(2, 3)
            + 0.5 * einsum("ijab,klab->ijkl", self.t2oo, self.t2oo)
        )
        ret.ooov = (
            + 2.0 * einsum("ja,ik->ijka", p0.ov, d_oo).antisymmetrise(0, 1)
        )
        ret.oovv = (
            - 1.0 * self.td2(b.oovv)
        )
        ret.ovov = (
            + 1.0 * einsum("ab,ij->iajb", p0.vv, d_oo)
            - 1.0 * einsum("jkac,ikbc->iajb", self.t2oo, self.t2oo)
        )
        ret.vvvv = (
            + 0.5 * einsum("ijab,ijcd->abcd", self.t2oo, self.t2oo)
        )
        return evaluate(ret)

    def diffdm_2p(self, level=2) -> TwoParticleDensity:
        """
        Return the two-particle MP difference density in the MO basis with all
        corrections up to the specified order of perturbation theory.
        """
        if level == 1:
            return self.mp1_dm_correction_2p
        elif level == 2:
            return (self.mp1_dm_correction_2p
                    + self.mp2_dm_correction_2p)
        else:
            raise NotImplementedError("Only first and second-order two-particle "
                                      "density corrections are implemented.")

    def density_2p(self, level=2) -> TwoParticleDensity:
        """
        Return the two-particle MP density in the MO basis with all corrections
        up to the specified order of perturbation theory.
        """
        if level == 0:
            return self.reference_state.density_2p
        diffdm = self.diffdm_2p(level)
        return self.reference_state.density_2p + diffdm

    def dipole_moment(self, level=2):
        """
        Return the MP dipole moment at the specified level of
        perturbation theory.
        """
        if level in [0, 1]:
            return self.reference_state.dipole_moment
        elif level == 2:
            return self.mp2_dipole_moment
        else:
            raise NotImplementedError("Only dipole moments for level 0, 1 and 2"
                                      " are implemented.")

    @cached_member_function()
    def energy_correction(self, level=2):
        """Obtain the MP energy correction at a particular level"""
        if level > 3:
            raise NotImplementedError(f"MP({level}) energy correction "
                                      "not implemented.")
        if level < 2:
            return 0.0
        hf = self.reference_state
        is_cvs = self.has_core_occupied_space
        if level == 2 and not is_cvs:
            terms = [(1.0, hf.oovv, self.t2oo)]
        elif level == 2 and is_cvs:
            terms = [(1.0, hf.oovv, self.t2oo),
                     (2.0, hf.ocvv, self.t2oc),
                     (1.0, hf.ccvv, self.t2cc)]
        elif level == 3 and not is_cvs:
            terms = [(1.0, hf.oovv, self.td2(b.oovv))]
        elif level == 3 and is_cvs:
            raise NotImplementedError("CVS-MP3 energy correction not implemented.")
        return sum(
            -0.25 * pref * eri.dot(t2)
            for pref, eri, t2 in terms
        )

    def energy(self, level=2):
        """
        Obtain the total energy (SCF energy plus all corrections)
        at a particular level of perturbation theory.
        """
        if level == 0:
            # Sum of orbital energies ...
            raise NotImplementedError("Total MP(0) energy not implemented.")

        # Accumulator for all energy terms
        energies = [self.reference_state.energy_scf]
        for il in range(2, level + 1):
            energies.append(self.energy_correction(il))
        return sum(energies)

    def to_qcvars(self, properties=False, recurse=False, maxlevel=2):
        """
        Return a dictionary with property keys compatible to a Psi4 wavefunction
        or a QCEngine Atomicresults object.
        """
        qcvars = {}
        for level in range(2, maxlevel + 1):
            try:
                mpcorr = self.energy_correction(level)
                qcvars[f"MP{level} CORRELATION ENERGY"] = mpcorr
                qcvars[f"MP{level} TOTAL ENERGY"] = self.energy(level)
            except NotImplementedError:
                pass
            except ValueError:
                pass

        if properties:
            for level in range(2, maxlevel + 1):
                try:
                    qcvars["MP2 DIPOLE"] = self.dipole_moment(level)
                except NotImplementedError:
                    pass

        if recurse:
            qcvars.update(self.reference_state.to_qcvars(properties, recurse))
        return qcvars

    @property
    def mp2_density(self):
        return self.density(2)

    @cached_property
    def mp2_dipole_moment(self):
        refstate = self.reference_state
        dipole_integrals = refstate.operators.electric_dipole
        mp2corr = np.array([product_trace(comp, self.mp2_diffdm)
                            for comp in dipole_integrals])
        return refstate.dipole_moment + mp2corr

    @cached_member_function()
    def ssq(self, level=2):
        """
        Return <S^2> of the ground state.
        """
        if self.reference_state.restricted:
            raise NotImplementedError(
                "<S^2> is not implemented for restricted HF references."
            )
        ssq_1p_op = self.reference_state.operators.ssq_1p
        ssq_2p_op = self.reference_state.operators.ssq_2p
        # the trace of the second-order (and higher) correction to the RDM1
        # is zero -> no influence on top of HF density for ground state
        ssq_1p = product_trace(ssq_1p_op, self.density(0))
        ssq_2p = product_trace(ssq_2p_op, self.density_2p(level))
        return (ssq_1p + ssq_2p)


#
# Register cvs_p0 intermediate
#
@register_as_intermediate
def cvs_p0(hf, mp, intermediates):
    # NOTE: equal to mp2_diffdm if CVS applied for the density
    ret = OneParticleDensity(hf.mospaces, symmetry=OperatorSymmetry.HERMITIAN)
    ret.oo = -0.5 * einsum("ikab,jkab->ij", mp.t2oo, mp.t2oo)
    ret.ov = -0.5 * (+ einsum("ijbc,jabc->ia", mp.t2oo, hf.ovvv)
                     + einsum("jkib,jkab->ia", hf.ooov, mp.t2oo)) / mp.df(b.ov)
    ret.vv = 0.5 * einsum("ijac,ijbc->ab", mp.t2oo, mp.t2oo)
    return ret
