#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import uccsd
from pyscf.cc import ccsd_lambda

einsum = lib.einsum

def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    if eris is None: eris = mycc.ao2mo()
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, nvira = t1a.shape
    noccb, nvirb = t1b.shape

    fooa = eris.focka[:nocca,:nocca]
    fova = eris.focka[:nocca,nocca:]
    fvoa = eris.focka[nocca:,:nocca]
    fvva = eris.focka[nocca:,nocca:]
    foob = eris.fockb[:noccb,:noccb]
    fovb = eris.fockb[:noccb,noccb:]
    fvob = eris.fockb[noccb:,:noccb]
    fvvb = eris.fockb[noccb:,noccb:]

    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)

    ovov = numpy.asarray(eris.ovov)
    ovov = ovov - ovov.transpose(0,3,2,1)
    OVOV = numpy.asarray(eris.OVOV)
    OVOV = OVOV - OVOV.transpose(0,3,2,1)
    ovOV = numpy.asarray(eris.ovOV)

    v1a  = fvva - einsum('ja,jb->ba', fova, t1a)
    v1b  = fvvb - einsum('ja,jb->ba', fovb, t1b)
    v1a += einsum('jcka,jkbc->ba', ovov, tauaa) * .5
    v1a -= einsum('jaKC,jKbC->ba', ovOV, tauab) * .5
    v1a -= einsum('kaJC,kJbC->ba', ovOV, tauab) * .5
    v1b += einsum('jcka,jkbc->ba', OVOV, taubb) * .5
    v1b -= einsum('kcJA,kJcB->BA', ovOV, tauab) * .5
    v1b -= einsum('jcKA,jKcB->BA', ovOV, tauab) * .5

    v2a  = fooa + einsum('ib,jb->ij', fova, t1a)
    v2b  = foob + einsum('ib,jb->ij', fovb, t1b)
    v2a += einsum('ibkc,jkbc->ij', ovov, tauaa) * .5
    v2a += einsum('ibKC,jKbC->ij', ovOV, tauab)
    v2b += einsum('ibkc,jkbc->ij', OVOV, taubb) * .5
    v2b += einsum('kcIB,kJcB->IJ', ovOV, tauab)

    ovoo = numpy.asarray(eris.ovoo)
    ovoo = ovoo - ovoo.transpose(2,1,0,3)
    OVOO = numpy.asarray(eris.OVOO)
    OVOO = OVOO - OVOO.transpose(2,1,0,3)
    OVoo = numpy.asarray(eris.OVoo)
    ovOO = numpy.asarray(eris.ovOO)
    v2a -= numpy.einsum('ibkj,kb->ij', ovoo, t1a)
    v2a += numpy.einsum('KBij,KB->ij', OVoo, t1b)
    v2b -= numpy.einsum('ibkj,kb->ij', OVOO, t1b)
    v2b += numpy.einsum('kbIJ,kb->IJ', ovOO, t1a)

    v5a  = fvoa + numpy.einsum('kc,jkbc->bj', fova, t2aa)
    v5a += numpy.einsum('KC,jKbC->bj', fovb, t2ab)
    v5b  = fvob + numpy.einsum('kc,jkbc->bj', fovb, t2bb)
    v5b += numpy.einsum('kc,kJcB->BJ', fova, t2ab)
    tmp  = fova - numpy.einsum('kdlc,ld->kc', ovov, t1a)
    tmp += numpy.einsum('kcLD,LD->kc', ovOV, t1b)
    v5a += einsum('kc,kb,jc->bj', tmp, t1a, t1a)
    tmp  = fovb - numpy.einsum('kdlc,ld->kc', OVOV, t1b)
    tmp += numpy.einsum('ldKC,ld->KC', ovOV, t1a)
    v5b += einsum('kc,kb,jc->bj', tmp, t1b, t1b)
    v5a -= einsum('lckj,klbc->bj', ovoo, t2aa) * .5
    v5a -= einsum('LCkj,kLbC->bj', OVoo, t2ab)
    v5b -= einsum('LCKJ,KLBC->BJ', OVOO, t2bb) * .5
    v5b -= einsum('lcKJ,lKcB->BJ', ovOO, t2ab)

    oooo = numpy.asarray(eris.oooo)
    OOOO = numpy.asarray(eris.OOOO)
    ooOO = numpy.asarray(eris.ooOO)
    woooo  = einsum('icjl,kc->ikjl', ovoo, t1a)
    wOOOO  = einsum('icjl,kc->ikjl', OVOO, t1b)
    wooOO  = einsum('icJL,kc->ikJL', ovOO, t1a)
    wooOO += einsum('JCil,KC->ilJK', OVoo, t1b)
    woooo += (oooo - oooo.transpose(0,3,2,1)) * .5
    wOOOO += (OOOO - OOOO.transpose(0,3,2,1)) * .5
    wooOO += ooOO.copy()
    woooo += einsum('icjd,klcd->ikjl', ovov, tauaa) * .25
    wOOOO += einsum('icjd,klcd->ikjl', OVOV, taubb) * .25
    wooOO += einsum('icJD,kLcD->ikJL', ovOV, tauab)

    v4ovvo  = einsum('jbld,klcd->jbck', ovov, t2aa)
    v4ovvo += einsum('jbLD,kLcD->jbck', ovOV, t2ab)
    v4ovvo += numpy.asarray(eris.ovvo)
    v4ovvo -= numpy.asarray(eris.oovv).transpose(0,3,2,1)
    v4OVVO  = einsum('jbld,klcd->jbck', OVOV, t2bb)
    v4OVVO += einsum('ldJB,lKdC->JBCK', ovOV, t2ab)
    v4OVVO += numpy.asarray(eris.OVVO)
    v4OVVO -= numpy.asarray(eris.OOVV).transpose(0,3,2,1)
    v4OVvo  = einsum('ldJB,klcd->JBck', ovOV, t2aa)
    v4OVvo += einsum('JBLD,kLcD->JBck', OVOV, t2ab)
    v4OVvo += numpy.asarray(eris.OVvo)
    v4ovVO  = einsum('jbLD,KLCD->jbCK', ovOV, t2bb)
    v4ovVO += einsum('jbld,lKdC->jbCK', ovov, t2ab)
    v4ovVO += numpy.asarray(eris.ovVO)
    v4oVVo  = einsum('jdLB,kLdC->jBCk', ovOV, t2ab)
    v4oVVo -= numpy.asarray(eris.ooVV).transpose(0,3,2,1)
    v4OvvO  = einsum('lbJD,lKcD->JbcK', ovOV, t2ab)
    v4OvvO -= numpy.asarray(eris.OOvv).transpose(0,3,2,1)

    woovo  = einsum('ibck,jb->ijck', v4ovvo, t1a)
    wOOVO  = einsum('ibck,jb->ijck', v4OVVO, t1b)
    wOOvo  = einsum('IBck,JB->IJck', v4OVvo, t1b)
    wOOvo -= einsum('IbcK,jb->IKcj', v4OvvO, t1a)
    wooVO  = einsum('ibCK,jb->ijCK', v4ovVO, t1a)
    wooVO -= einsum('iBCk,JB->ikCJ', v4oVVo, t1b)
    woovo += ovoo.conj().transpose(3,2,1,0) * .5
    wOOVO += OVOO.conj().transpose(3,2,1,0) * .5
    wooVO += OVoo.conj().transpose(3,2,1,0)
    wOOvo += ovOO.conj().transpose(3,2,1,0)
    woovo -= einsum('iclk,jlbc->ikbj', ovoo, t2aa)
    woovo += einsum('LCik,jLbC->ikbj', OVoo, t2ab)
    wOOVO -= einsum('iclk,jlbc->ikbj', OVOO, t2bb)
    wOOVO += einsum('lcIK,lJcB->IKBJ', ovOO, t2ab)
    wooVO -= einsum('iclk,lJcB->ikBJ', ovoo, t2ab)
    wooVO += einsum('LCik,JLBC->ikBJ', OVoo, t2bb)
    wooVO -= einsum('icLK,jLcB->ijBK', ovOO, t2ab)
    wOOvo -= einsum('ICLK,jLbC->IKbj', OVOO, t2ab)
    wOOvo += einsum('lcIK,jlbc->IKbj', ovOO, t2aa)
    wOOvo -= einsum('IClk,lJbC->IJbk', OVoo, t2ab)

    wvvvo  = einsum('jack,jb->back', v4ovvo, t1a)
    wVVVO  = einsum('jack,jb->back', v4OVVO, t1b)
    wVVvo  = einsum('JAck,JB->BAck', v4OVvo, t1b)
    wVVvo -= einsum('jACk,jb->CAbk', v4oVVo, t1a)
    wvvVO  = einsum('jaCK,jb->baCK', v4ovVO, t1a)
    wvvVO -= einsum('JacK,JB->caBK', v4OvvO, t1b)
    wvvvo += einsum('lajk,jlbc->back', .25*ovoo, tauaa)
    wVVVO += einsum('lajk,jlbc->back', .25*OVOO, taubb)
    wVVvo -= einsum('LAjk,jLcB->BAck', OVoo, tauab)
    wvvVO -= einsum('laJK,lJbC->baCK', ovOO, tauab)

    w3a  = numpy.einsum('jbck,jb->ck', v4ovvo, t1a)
    w3a += numpy.einsum('JBck,JB->ck', v4OVvo, t1b)
    w3b  = numpy.einsum('jbck,jb->ck', v4OVVO, t1b)
    w3b += numpy.einsum('jbCK,jb->CK', v4ovVO, t1a)

    wovvo  = v4ovvo
    wOVVO  = v4OVVO
    wovVO  = v4ovVO
    wOVvo  = v4OVvo
    woVVo  = v4oVVo
    wOvvO  = v4OvvO
    wovvo += lib.einsum('jbld,kd,lc->jbck', ovov, t1a, -t1a)
    wOVVO += lib.einsum('jbld,kd,lc->jbck', OVOV, t1b, -t1b)
    wovVO += lib.einsum('jbLD,KD,LC->jbCK', ovOV, t1b, -t1b)
    wOVvo += lib.einsum('ldJB,kd,lc->JBck', ovOV, t1a, -t1a)
    woVVo += lib.einsum('jdLB,kd,LC->jBCk', ovOV, t1a,  t1b)
    wOvvO += lib.einsum('lbJD,KD,lc->JbcK', ovOV, t1b,  t1a)
    wovvo -= einsum('jblk,lc->jbck', ovoo, t1a)
    wOVVO -= einsum('jblk,lc->jbck', OVOO, t1b)
    wovVO -= einsum('jbLK,LC->jbCK', ovOO, t1b)
    wOVvo -= einsum('JBlk,lc->JBck', OVoo, t1a)
    woVVo += einsum('LBjk,LC->jBCk', OVoo, t1b)
    wOvvO += einsum('lbJK,lc->JbcK', ovOO, t1a)

    if nvira > 0 and nocca > 0:
        ovvv = numpy.asarray(eris.get_ovvv())
        ovvv = ovvv - ovvv.transpose(0,3,2,1)
        v1a -= numpy.einsum('jabc,jc->ba', ovvv, t1a)
        v5a += einsum('kdbc,jkcd->bj', ovvv, t2aa) * .5
        woovo += einsum('idcb,kjbd->ijck', ovvv, tauaa) * .25
        wovvo += einsum('jbcd,kd->jbck', ovvv, t1a)
        wvvvo -= ovvv.conj().transpose(3,2,1,0) * .5
        wvvvo += einsum('jacd,kjbd->cabk', ovvv, t2aa)
        wvvVO += einsum('jacd,jKdB->caBK', ovvv, t2ab)
        ovvv = tmp = None

    if nvirb > 0 and noccb > 0:
        OVVV = numpy.asarray(eris.get_OVVV())
        OVVV = OVVV - OVVV.transpose(0,3,2,1)
        v1b -= numpy.einsum('jabc,jc->ba', OVVV, t1b)
        v5b += einsum('KDBC,JKCD->BJ', OVVV, t2bb) * .5
        wOOVO += einsum('idcb,kjbd->ijck', OVVV, taubb) * .25
        wOVVO += einsum('jbcd,kd->jbck', OVVV, t1b)
        wVVVO -= OVVV.conj().transpose(3,2,1,0) * .5
        wVVVO += einsum('jacd,kjbd->cabk', OVVV, t2bb)
        wVVvo += einsum('JACD,kJbD->CAbk', OVVV, t2ab)
        OVVV = tmp = None

    if nvirb > 0 and nocca > 0:
        OVvv = numpy.asarray(eris.get_OVvv())
        v1a += numpy.einsum('JCba,JC->ba', OVvv, t1b)
        v5a += einsum('KDbc,jKcD->bj', OVvv, t2ab)
        wOOvo += einsum('IDcb,kJbD->IJck', OVvv, tauab)
        wOVvo += einsum('JBcd,kd->JBck', OVvv, t1a)
        wOvvO -= einsum('JDcb,KD->JbcK', OVvv, t1b)
        wvvVO -= OVvv.conj().transpose(3,2,1,0)
        wvvvo -= einsum('KDca,jKbD->cabj', OVvv, t2ab)
        wvvVO -= einsum('KDca,JKBD->caBJ', OVvv, t2bb)
        wVVvo += einsum('KAcd,jKdB->BAcj', OVvv, t2ab)
        OVvv = tmp = None

    if nvira > 0 and noccb > 0:
        ovVV = numpy.asarray(eris.get_ovVV())
        v1b += numpy.einsum('jcBA,jc->BA', ovVV, t1a)
        v5b += einsum('kdBC,kJdC->BJ', ovVV, t2ab)
        wooVO += einsum('idCB,jKdB->ijCK', ovVV, tauab)
        wovVO += einsum('jbCD,KD->jbCK', ovVV, t1b)
        woVVo -= einsum('jdCB,kd->jBCk', ovVV, t1a)
        wVVvo -= ovVV.conj().transpose(3,2,1,0)
        wVVVO -= einsum('kdCA,kJdB->CABJ', ovVV, t2ab)
        wVVvo -= einsum('kdCA,jkbd->CAbj', ovVV, t2aa)
        wvvVO += einsum('kaCD,kJbD->baCJ', ovVV, t2ab)
        ovVV = tmp = None

    w3a += v5a
    w3b += v5b
    w3a += lib.einsum('cb,jb->cj', v1a, t1a)
    w3b += lib.einsum('cb,jb->cj', v1b, t1b)
    w3a -= lib.einsum('jk,jb->bk', v2a, t1a)
    w3b -= lib.einsum('jk,jb->bk', v2b, t1b)

    class _IMDS: pass
    imds = _IMDS()
    imds.ftmp = lib.H5TmpFile()
    dtype = numpy.result_type(t2ab, eris.vvvv).char
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocca,nocca,nocca,nocca), dtype)
    imds.wooOO = imds.ftmp.create_dataset('wooOO', (nocca,nocca,noccb,noccb), dtype)
    imds.wOOOO = imds.ftmp.create_dataset('wOOOO', (noccb,noccb,noccb,noccb), dtype)
    imds.wovvo = imds.ftmp.create_dataset('wovvo', (nocca,nvira,nvira,nocca), dtype)
    imds.wOVVO = imds.ftmp.create_dataset('wOVVO', (noccb,nvirb,nvirb,noccb), dtype)
    imds.wovVO = imds.ftmp.create_dataset('wovVO', (nocca,nvira,nvirb,noccb), dtype)
    imds.wOVvo = imds.ftmp.create_dataset('wOVvo', (noccb,nvirb,nvira,nocca), dtype)
    imds.woVVo = imds.ftmp.create_dataset('woVVo', (nocca,nvirb,nvirb,nocca), dtype)
    imds.wOvvO = imds.ftmp.create_dataset('wOvvO', (noccb,nvira,nvira,noccb), dtype)
    imds.woovo = imds.ftmp.create_dataset('woovo', (nocca,nocca,nvira,nocca), dtype)
    imds.wOOVO = imds.ftmp.create_dataset('wOOVO', (noccb,noccb,nvirb,noccb), dtype)
    imds.wOOvo = imds.ftmp.create_dataset('wOOvo', (noccb,noccb,nvira,nocca), dtype)
    imds.wooVO = imds.ftmp.create_dataset('wooVO', (nocca,nocca,nvirb,noccb), dtype)
    imds.wvvvo = imds.ftmp.create_dataset('wvvvo', (nvira,nvira,nvira,nocca), dtype)
    imds.wVVVO = imds.ftmp.create_dataset('wVVVO', (nvirb,nvirb,nvirb,noccb), dtype)
    imds.wVVvo = imds.ftmp.create_dataset('wVVvo', (nvirb,nvirb,nvira,nocca), dtype)
    imds.wvvVO = imds.ftmp.create_dataset('wvvVO', (nvira,nvira,nvirb,noccb), dtype)

    imds.woooo[:] = woooo
    imds.wOOOO[:] = wOOOO
    imds.wooOO[:] = wooOO
    imds.wovvo[:] = wovvo
    imds.wOVVO[:] = wOVVO
    imds.wovVO[:] = wovVO
    imds.wOVvo[:] = wOVvo
    imds.woVVo[:] = woVVo
    imds.wOvvO[:] = wOvvO
    imds.woovo[:] = woovo
    imds.wOOVO[:] = wOOVO
    imds.wOOvo[:] = wOOvo
    imds.wooVO[:] = wooVO
    imds.wvvvo[:] = wvvvo