# Copyright 2014-2018,2021 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.

import numpy as np

import pyscf.pbc.tools as tools
from mpi4py import MPI
from pyscf.lib import logger
from pyscf.pbc.mpitools import mpi_load_balancer
from pyscf import lib
from pyscf.pbc import lib as pbclib
from pyscf.pbc.tools.tril import tril_index, unpack_tril

comm = MPI.COMM_WORLD

#einsum = np.einsum
einsum = lib.einsum
dot = np.dot

# This is restricted (R)CCSD
# Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004)

### Eqs. (37)-(39) "kappa"

def cc_tau1(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    tmp_oovv_shape = BLKSIZE + (nocc,nocc,nvir,nvir)
    tmp_oovv  = np.empty(tmp_oovv_shape,dtype=t1.dtype)
    tmp2_oovv = np.zeros_like(tmp_oovv)

    tau1_ooVv = feri2['tau1_ooVv']
    tau1_oOvv = feri2['tau1_oOvv']
    tau1_oovv_rev = feri2['tau1_oovv_rev']
    tau2_Oovv = feri2['tau2_Oovv']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterki, ki in enumerate(ranges0):
            for iterkj, kj in enumerate(ranges1):
                for iterka, ka in enumerate(ranges2):
                    kb = kconserv[ki,ka,kj]
                    tmp_oovv[iterki,iterkj,iterka] = t2[ki,kj,ka].copy()
                    tmp2_oovv[iterki,iterkj,iterka] *= 0.0
                    if ki == ka and kj == kb:
                        tmp2_oovv[iterki,iterkj,iterka] = einsum('ia,jb->ijab',t1[ki],t1[kj])

                    tau1_oovv_rev[kj,ka,kb] = (tmp_oovv[iterki,iterkj,iterka] + tmp2_oovv[iterki,iterkj,iterka])

        tau1_ooVv[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,nvir*min(ranges2):nvir*(max(ranges2)+1)] = \
                (tmp_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] +
                 tmp2_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] ).transpose(0,1,2,5,3,4,6).reshape(len(ranges0),len(ranges1),len(ranges2)*nvir,nocc,nocc,nvir)
        tau1_oOvv[min(ranges0):max(ranges0)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges1):nocc*(max(ranges1)+1)] = \
                (tmp_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] +
                 tmp2_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] ).transpose(0,2,1,4,3,5,6).reshape(len(ranges0),len(ranges2),len(ranges1)*nocc,nocc,nvir,nvir)
        tau2_Oovv[min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges0):nocc*(max(ranges0)+1)] = \
                (tmp_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] +
                 2*tmp2_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] ).transpose(1,2,0,3,4,5,6).reshape(len(ranges1),len(ranges2),len(ranges0)*nocc,nocc,nvir,nvir)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Foo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    Fki = np.empty((nkpts,nocc,nocc),dtype=t2.dtype)

    for ki in range(nkpts):
        kk = ki
        Fki[ki] = eris.fock[ki,:nocc,:nocc].copy()
        for kl in range(nkpts):
            for kc in range(nkpts):
                #Fki[ki] += einsum('lkcd,licd->ki',eris.SoOvv[kk,kc],tau1_oOvv[ki,kc])
                kd = kconserv[kk,kc,kl]
                Soovv = 2*eris.oovv[kk,kl,kc] - eris.oovv[kk,kl,kd].transpose(0,1,3,2)
                #Fki[ki] += einsum('klcd,ilcd->ki',Soovv,t2[ki,kl,kc])
                Fki[ki] += einsum('klcd,ilcd->ki',Soovv,unpack_tril(t2,nkpts,ki,kl,kc,kconserv[ki,kc,kl]))
            #if ki == kc:
            kd = kconserv[kk,ki,kl]
            Soovv = 2*eris.oovv[kk,kl,ki] - eris.oovv[kk,kl,kd].transpose(0,1,3,2)
            Fki[ki] += einsum('klcd,ic,ld->ki',Soovv,t1[ki],t1[kl])
    return Fki

def cc_Fvv(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    Fac = np.empty((nkpts,nvir,nvir),dtype=t2.dtype)

    for ka in range(nkpts):
        kc = ka
        Fac[ka] = eris.fock[ka,nocc:,nocc:].copy()
        #for kk in range(nkpts):
        #    Fac[ka] += -einsum('lkcd,lkad->ac',eris.SoOvv[kk,kc],tau1_oOvv[kk,ka])
        for kl in range(nkpts):
            for kk in range(nkpts):
                kd = kconserv[kk,kc,kl]
                Soovv = 2*eris.oovv[kk,kl,kc] - eris.oovv[kk,kl,kd].transpose(0,1,3,2)
                #Fac[ka] += -einsum('klcd,klad->ac',Soovv,t2[kk,kl,ka])
                Fac[ka] += -einsum('klcd,klad->ac',Soovv,unpack_tril(t2,nkpts,kk,kl,ka,kconserv[kk,ka,kl]))
            #if kk == ka
            kd = kconserv[ka,kc,kl]
            Soovv = 2*eris.oovv[ka,kl,kc] - eris.oovv[ka,kl,kd].transpose(0,1,3,2)
            Fac[ka] += -einsum('klcd,ka,ld->ac',Soovv,t1[ka],t1[kl])
    return Fac

def cc_Fov(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    Fkc = np.empty((nkpts,nocc,nvir),dtype=t2.dtype)
    Fkc[:] = eris.fock[:,:nocc,nocc:].copy()
    for kk in range(nkpts):
        for kl in range(nkpts):
            Soovv = 2.*eris.oovv[kk,kl,kk] - eris.oovv[kk,kl,kl].transpose(0,1,3,2)
            Fkc[kk] += einsum('klcd,ld->kc',Soovv,t1[kl])
    return Fkc

### Eqs. (40)-(41) "lambda"

def Loo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    fov = eris.fock[:,:nocc,nocc:]
    Lki = cc_Foo(cc,t1,t2,eris,feri2)
    for ki in range(nkpts):
        Lki[ki] += einsum('kc,ic->ki',fov[ki],t1[ki])
        SoOov = (2*eris.ooov[ki,:,ki] -
                 eris.ooov[:,ki,ki].transpose(0,2,1,3,4)).transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nocc,nvir)
        Lki[ki] += einsum('lkic,lc->ki',SoOov,t1.reshape(nkpts*nocc,nvir))
    return Lki

def Lvv(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    fov = eris.fock[:,:nocc,nocc:]
    Lac = cc_Fvv(cc,t1,t2,eris,feri2)
    for ka in range(nkpts):
        Lac[ka] += -einsum('kc,ka->ac',fov[ka],t1[ka])
        for kk in range(nkpts):
            Svovv = 2*eris.ovvv[kk,ka,kk].transpose(1,0,3,2) - eris.ovvv[kk,ka,ka].transpose(1,0,2,3)
            Lac[ka] += einsum('akcd,kd->ac',Svovv,t1[kk])
    return Lac

### Eqs. (42)-(45) "chi"

def cc_Woooo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    #Wklij = np.array(eris.oooo, copy=True)
    #for pqr in range(nUnique_klist):
    #    kk, kl, ki = unique_klist[pqr]
    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    oooo_tmp_shape = BLKSIZE + (nocc,nocc,nocc,nocc)
    oooo_tmp = np.empty(shape=oooo_tmp_shape,dtype=t1.dtype)

    tau1_ooVv = feri2['tau1_ooVv']
    #Woooo     = feri2['Woooo']
    Woooo_rev = feri2['Woooo_rev']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterkk, kk in enumerate(ranges0):
            for iterkl, kl in enumerate(ranges1):
                for iterki, ki in enumerate(ranges2):
                    kj = kconserv[kk,ki,kl]
                    oooo_tmp[iterkk,iterkl,iterki] = np.array(eris.oooo[kk,kl,ki],copy=True)
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('klic,jc->klij',eris.ooov[kk,kl,ki],t1[kj])
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('klcj,ic->klij',eris.ooov[kl,kk,kj].transpose(1,0,3,2),t1[ki])

                    # Note the indices and way the tau1 is stored : instead of a loop over kpt='kc' and
                    # loop over mo='c', the (kc,k,l,c,d) index is changed instead to (nkpts*nvir,k,l,d)
                    # so that we only have to loop over the first index, saving read operations.
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('ckld,cijd->klij',eris.ooVv[kk,kl],tau1_ooVv[ki,kj])

                    #for kc in range(nkpts):
                    #    oooo_tmp[iterkk,iterkl,iterki] += einsum('klcd,ijcd->klij',eris.oovv[kk,kl,kc],t2[ki,kj,kc])
                    #oooo_tmp[iterkk,iterkl,iterki] += einsum('klcd,ic,jd->klij',eris.oovv[kk,kl,ki],t1[ki],t1[kj])
                    #Woooo[kk,kl,ki] = oooo_tmp[iterkk,iterkl,iterki]
                    #Woooo[kl,kk,kj] = oooo_tmp[iterkk,iterkl,iterki].transpose(1,0,3,2)
                    Woooo_rev[kl,ki,kj] = oooo_tmp[iterkk,iterkl,iterki]

        #Woooo[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
        #                oooo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]

        # for if you want to take into account symmetry of Woooo integral
        #feri2.Woooo[min(ranges1):max(ranges1)+1,min(ranges0):max(ranges0)+1,min(ranges2):max(ranges2)+1] = \
        #                oooo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(0,1,2,4,3,6,5)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Wvvvv(cc,t1,t2,eris,feri2=None):
    ## Slow:
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    vvvv_tmp = np.empty((nvir,nvir,nvir,nvir),dtype=t1.dtype)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(1,nkpts,))
    loader.set_ranges((range(nkpts),range(nkpts),))

    Wvvvv = feri2['Wvvvv']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1 = loader.get_blocks_from_data(data)
        for ka in ranges0:
            for kc in ranges1:
                for kb in range(ka+1):
                    kd = kconserv[ka,kc,kb]
                    vvvv_tmp = np.array(eris.vvvv[ka,kb,kc],copy=True)
                    vvvv_tmp += einsum('akcd,kb->abcd',eris.ovvv[kb,ka,kd].transpose(1,0,3,2),-t1[kb])
                    vvvv_tmp += einsum('kbcd,ka->abcd',eris.ovvv[ka,kb,kc],-t1[ka])
                    Wvvvv[ka,kb,kc] = vvvv_tmp
                    Wvvvv[kb,ka,kd] = vvvv_tmp.transpose(1,0,3,2)
        loader.slave_finished()

    ## Fast
    #nocc,nvir = t1.shape
    #Wabcd = np.empty((nvir,)*4)
    #for a in range(nvir):
    #    Wabcd[a,:] = einsum('kcd,kb->bcd',eris.vovv[a],-t1)
    ##Wabcd += einsum('kbcd,ka->abcd',eris.ovvv,-t1)
    #Wabcd += lib.dot(-t1.T,eris.ovvv.reshape(nocc,-1)).reshape((nvir,)*4)
    #Wabcd += np.asarray(eris.vvvv)

    comm.Barrier()
    return

def cc_Wvoov(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    #Wakic = np.empty((nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir),dtype=t1.dtype)

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    voov_tmp_shape = BLKSIZE + (nvir,nocc,nocc,nvir)
    voov_tmp = np.empty(voov_tmp_shape,dtype=t1.dtype)

    tau2_Oovv = feri2['tau2_Oovv']
    #Wvoov     = feri2['Wvoov']
    WvOov     = feri2['WvOov']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        #ix = sum([[min(x),max(x)+1] for x in (ranges0,ranges1,ranges2)], [])
        #eris_ooov = eris.ooov[ix[0]:ix[1], ix[2]:ix[3], ix[4]:ix[5]]
        for iterka, ka in enumerate(ranges0):
            for iterkk, kk in enumerate(ranges1):
                for iterki, ki in enumerate(ranges2):
                    kc = kconserv[ka,ki,kk]
                    voov_tmp[iterka,iterkk,iterki] = np.array(eris.ovvo[kk,ka,kc]).transpose(1,0,3,2)
                    voov_tmp[iterka,iterkk,iterki] -= einsum('lkic,la->akic',eris.ooov[ka,kk,ki],t1[ka])
                    voov_tmp[iterka,iterkk,iterki] += einsum('akdc,id->akic',eris.ovvv[kk,ka,kc].transpose(1,0,3,2),t1[ki])
                    # Beginning of change
                    #for kl in range(nkpts):
                    #    # kl - kd + kk = kc
                    #    # => kd = kl - kc + kk
                    #    kd = kconserv[kl,kc,kk]
                    #    Soovv = 2*np.array(eris.oovv[kl,kk,kd]) - np.array(eris.oovv[kl,kk,kc]).transpose(0,1,3,2)
                    #    voov_tmp[iterka,iterkk,iterki] += 0.5*einsum('lkdc,ilad->akic',Soovv,t2[ki,kl,ka])
                    #    voov_tmp[iterka,iterkk,iterki] -= 0.5*einsum('lkdc,ilda->akic',eris.oovv[kl,kk,kd],t2[ki,kl,kd])
                    #voov_tmp[iterka,iterkk,iterki] -= einsum('lkdc,id,la->akic',eris.oovv[ka,kk,ki],t1[ki],t1[ka])
                    #Wvoov[ka,kk,ki] = voov_tmp[iterka,iterkk,iterki]

                    # Making various intermediates...
                    #t2_oOvv = t2[ki,:,ka].transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nvir,nvir)
                    t2_oOvv = unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)]).transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nvir,nvir)
                    #eris_oOvv = eris.oovv[kk,:,kc].transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nvir,nvir)

                    voov_tmp[iterka,iterkk,iterki] += 0.5*einsum('lkcd,liad->akic',eris.SoOvv[kk,kc],t2_oOvv)
                    voov_tmp[iterka,iterkk,iterki] -= 0.5*einsum('lkcd,liad->akic',eris.oOvv[kk,kc],tau2_Oovv[ki,ka])

                    # End of change
        #Wvoov[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
        #        voov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        WvOov[min(ranges0):max(ranges0)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges1):nocc*(max(ranges1)+1)] = \
                voov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(0,2,1,4,3,5,6).reshape(
                    len(ranges1),len(ranges2),len(ranges0)*nocc,nvir,nocc,nvir)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Wvovo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    vovo_tmp_shape = BLKSIZE + (nvir,nocc,nvir,nocc)
    vovo_tmp = np.empty(shape=vovo_tmp_shape,dtype=t1.dtype)

    Wvovo = feri2['Wvovo']
    WvOVo = feri2['WvOVo']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterka, ka in enumerate(ranges0):
            for iterkk, kk in enumerate(ranges1):
                for iterkc, kc in enumerate(ranges2):
                    ki = kconserv[ka,kc,kk]
                    vovo_tmp[iterka,iterkk,iterkc] = np.array(eris.ovov[kk,ka,ki]).transpose(1,0,3,2)
                    vovo_tmp[iterka,iterkk,iterkc] -= einsum('lkci,la->akci',eris.ooov[kk,ka,ki].transpose(1,0,3,2),t1[ka])
                    vovo_tmp[iterka,iterkk,iterkc] += einsum('akcd,id->akci',eris.ovvv[kk,ka,ki].transpose(1,0,3,2),t1[ki])
                    # Beginning of change
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    vovo_tmp[iterka,iterkk,iterkc] -= 0.5*einsum('lkcd,ilda->akci',eris.oovv[kl,kk,kc],t2[ki,kl,kd])
                    #vovo_tmp[iterka,iterkk,iterkc] -= einsum('lkcd,id,la->akci',eris.oovv[ka,kk,kc],t1[ki],t1[ka])
                    #Wvovo[ka,kk,kc] = vovo_tmp[iterka,iterkk,iterkc]

                    oovvf = eris.oovv[:,kk,kc].reshape(nkpts*nocc,nocc,nvir,nvir)
                    #t2f = t2[:,ki,ka].copy() #This is a tau like term
                    t2f = unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).copy() #This is a tau like term
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    if ki == kd and kl == ka:
                    #        t2f[kl] += 2*einsum('id,la->liad',t1[ki],t1[ka])
                    kd = kconserv[ka,kc,kk]
                    t2f[ka] += 2*einsum('id,la->liad',t1[kd],t1[ka])
                    t2f = t2f.reshape(nkpts*nocc,nocc,nvir,nvir)
                    vovo_tmp[iterka,iterkk,iterkc] -= 0.5*einsum('lkcd,liad->akci',oovvf,t2f)

        Wvovo[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                vovo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        WvOVo[min(ranges0):max(ranges0)+1,nocc*min(ranges1):nocc*(max(ranges1)+1),nvir*min(ranges2):nvir*(max(ranges2)+1)] = \
                vovo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(0,1,4,2,5,3,6).reshape(
                    len(ranges0),len(ranges1)*nocc,len(ranges2)*nvir,nvir,nocc)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Wovov(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    ovov_tmp_shape = BLKSIZE + (nocc,nvir,nocc,nvir)
    ovov_tmp = np.empty(shape=ovov_tmp_shape,dtype=t1.dtype)

    #Wovov = feri2['Wovov']
    WOvov = feri2['WOvov']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterkk, kk in enumerate(ranges0):
            for iterka, ka in enumerate(ranges1):
                for iterki, ki in enumerate(ranges2):
                    kc = kconserv[kk,ki,ka]
                    ovov_tmp[iterkk,iterka,iterki] = np.array(eris.ovov[kk,ka,ki],copy=True)
                    ovov_tmp[iterkk,iterka,iterki] -= einsum('lkci,la->kaic',eris.ooov[kk,ka,ki].transpose(1,0,3,2),t1[ka])
                    ovov_tmp[iterkk,iterka,iterki] += einsum('akcd,id->kaic',eris.ovvv[kk,ka,ki].transpose(1,0,3,2),t1[ki])
                    # Beginning of change
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    ovov_tmp[iterka,iterkk,iterkc] -= 0.5*einsum('lkcd,ilda->akci',eris.oovv[kl,kk,kc],t2[ki,kl,kd])
                    #ovov_tmp[iterka,iterkk,iterkc] -= einsum('lkcd,id,la->akci',eris.oovv[ka,kk,kc],t1[ki],t1[ka])
                    #Wvovo[ka,kk,kc] = ovov_tmp[iterka,iterkk,iterkc]

                    oovvf = eris.oovv[:,kk,kc].reshape(nkpts*nocc,nocc,nvir,nvir)
                    #t2f = t2[:,ki,ka].copy() #This is a tau like term
                    t2f = unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).copy() #This is a tau like term
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    if ki == kd and kl == ka:
                    #        t2f[kl] += 2*einsum('id,la->liad',t1[ki],t1[ka])
                    kd = kconserv[ka,kc,kk]
                    t2f[ka] += 2*einsum('id,la->liad',t1[kd],t1[ka])
                    t2f = t2f.reshape(nkpts*nocc,nocc,nvir,nvir)
                    ovov_tmp[iterkk,iterka,iterki] -= 0.5*einsum('lkcd,liad->kaic',oovvf,t2f)

        #Wovov[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
        #        ovov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        WOvov[min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges0):nocc*(max(ranges0)+1)] = \
                ovov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(1,2,0,3,4,5,6).reshape(
                    len(ranges1),len(ranges2),len(ranges0)*nocc,nvir,nocc,nvir)
        loader.slave_finished()
    comm.Barrier()
    return

# EOM Intermediates w/ k-points

def Wooov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wklid = np.zeros((nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir),dtype=t2.dtype)
    else:
        Wklid = fint['Wooov']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ooov_tmp_size = BLKSIZE + (nocc,nocc,nocc,nvir)
    ooov_tmp = np.empty(ooov_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_ooov_kli = _cp(eris.ooov[s0,s1,s2])
        eris_oovv_kli = _cp(eris.oovv[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkl,kl in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    ooov_tmp[iterkk,iterkl,iterki] = eris_ooov_kli[iterkk,iterkl,iterki].copy()
                    ooov_tmp[iterkk,iterkl,iterki] += einsum('ic,klcd->klid', t1[ki],
                                                             eris_oovv_kli[iterkk,iterkl,iterki])
        Wklid[s0,s1,s2] = ooov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wklid, op=MPI.SUM)

    return Wklid


def Wvovv(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Walcd = np.zeros((nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir),dtype=t2.dtype)
    else:
        Walcd = fint['Wvovv']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nvir*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    vovv_tmp_size = BLKSIZE + (nvir,nocc,nvir,nvir)
    vovv_tmp = np.empty(vovv_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_vovv_alc = _cp(eris.vovv[s0,s1,s2])
        eris_oovv_alc = _cp(eris.oovv[s0,s1,s2])

        for iterka,ka in enumerate(ranges0):
            for iterkl,kl in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    # vovv[ka,kl,kc,kd] <= ovvv[kl,ka,kd,kc].transpose(1,0,3,2)
                    vovv_tmp[iterka,iterkl,iterkc] = eris_vovv_alc[iterka,iterkl,iterkc] #np.array(eris.ovvv[kl,ka,kd]).transpose(1,0,3,2)
                    vovv_tmp[iterka,iterkl,iterkc] += -einsum('ka,klcd->alcd',t1[ka],eris_oovv_alc[iterka,iterkl,iterkc])
        Walcd[s0,s1,s2] = vovv_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Walcd, op=MPI.SUM)

    return Walcd

def W1ovvo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc),dtype=t1.dtype)
    else:
        Wkaci  = fint['W1ovvo']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovvo_tmp_size = BLKSIZE + (nocc,nvir,nvir,nocc)
    ovvo_tmp = np.empty(ovvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        eris_ovvo_kac = _cp(eris.ovvo[s0,s1,s2])
        eris_oovv_kXc = _cp(eris.oovv[s0,:,s2])
        eris_oovv_Xkc = _cp(eris.oovv[:,s0,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterka,ka in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    ki = kconserv[kk,kc,ka]
                    ovvo_tmp[iterkk,iterka,iterkc] = _cp(eris_ovvo_kac[iterkk,iterka,iterkc])
                    #St2 = 2.*t2[ki,:,ka]
                    St2 = 2.*unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)])
                    #St2 -= t2[:,ki,ka].transpose(0,2,1,3,4)
                    St2 -= unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).transpose(0,2,1,3,4)
                    ovvo_tmp[iterkk,iterka,iterkc] += einsum('klcd,ilad->kaci',
                                                             eris_oovv_kXc[iterkk,:,iterkc].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir),
                                                             St2.transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    ovvo_tmp[iterkk,iterka,iterkc] += -einsum('lkcd,ilad->kaci',
                                                              eris_oovv_Xkc[:,iterkk,iterkc].reshape(nocc*nkpts,nocc,nvir,nvir),
                                                              unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)]).transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
#                                                                t2[ki,:,ka].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
        Wkaci[s0,s1,s2] = ovvo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci

def W2ovvo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc),dtype=t1.dtype)
        WWooov = Wooov(cc,t1,t2,eris)
    else:
        Wkaci  = fint['W2ovvo']
        WWooov = fint['Wooov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovvo_tmp_size = BLKSIZE + (nocc,nvir,nvir,nocc)
    ovvo_tmp = np.empty(ovvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        Wooov_akX     = _cp(WWooov[s1,s0])
        eris_ovvv_kac = _cp(eris.ovvv[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterka,ka in enumerate(ranges1):