#!/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>
#

'''
CASCI analytical nuclear gradients

Ref.
J. Comput. Chem., 5, 589
'''

import sys
import time
from functools import reduce
import numpy
from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.grad import rhf as rhf_grad
from pyscf.grad.mp2 import _shell_prange
from pyscf.scf import cphf

if sys.version_info < (3,):
    RANGE_TYPE = list
else:
    RANGE_TYPE = range


def grad_elec(mc_grad, mo_coeff=None, ci=None, atmlst=None, verbose=None):
    mc = mc_grad.base
    if mo_coeff is None: mo_coeff = mc._scf.mo_coeff
    if ci is None: ci = mc.ci

    time0 = time.clock(), time.time()
    log = logger.new_logger(mc_grad, verbose)
    mol = mc_grad.mol
    ncore = mc.ncore
    ncas = mc.ncas
    nocc = ncore + ncas
    nelecas = mc.nelecas
    nao, nmo = mo_coeff.shape
    nao_pair = nao * (nao+1) // 2
    mo_energy = mc._scf.mo_energy

    mo_occ = mo_coeff[:,:nocc]
    mo_core = mo_coeff[:,:ncore]
    mo_cas = mo_coeff[:,ncore:nocc]
    neleca, nelecb = mol.nelec
    assert(neleca == nelecb)
    orbo = mo_coeff[:,:neleca]
    orbv = mo_coeff[:,neleca:]

    casdm1, casdm2 = mc.fcisolver.make_rdm12(ci, ncas, nelecas)
    dm_core = numpy.dot(mo_core, mo_core.T) * 2
    dm_cas = reduce(numpy.dot, (mo_cas, casdm1, mo_cas.T))
    aapa = ao2mo.kernel(mol, (mo_cas, mo_cas, mo_coeff, mo_cas), compact=False)
    aapa = aapa.reshape(ncas,ncas,nmo,ncas)
    vj, vk = mc._scf.get_jk(mol, (dm_core, dm_cas))
    h1 = mc.get_hcore()
    vhf_c = vj[0] - vk[0] * .5
    vhf_a = vj[1] - vk[1] * .5
    # Imat = h1_{pi} gamma1_{iq} + h2_{pijk} gamma_{iqkj}
    Imat = numpy.zeros((nmo,nmo))
    Imat[:,:nocc] = reduce(numpy.dot, (mo_coeff.T, h1 + vhf_c + vhf_a, mo_occ)) * 2
    Imat[:,ncore:nocc] = reduce(numpy.dot, (mo_coeff.T, h1 + vhf_c, mo_cas, casdm1))
    Imat[:,ncore:nocc] += lib.einsum('uviw,vuwt->it', aapa, casdm2)
    aapa = vj = vk = vhf_c = vhf_a = h1 = None

    ee = mo_energy[:,None] - mo_energy
    zvec = numpy.zeros_like(Imat)
    zvec[:ncore,ncore:neleca] = Imat[:ncore,ncore:neleca] / -ee[:ncore,ncore:neleca]
    zvec[ncore:neleca,:ncore] = Imat[ncore:neleca,:ncore] / -ee[ncore:neleca,:ncore]
    zvec[nocc:,neleca:nocc] = Imat[nocc:,neleca:nocc] / -ee[nocc:,neleca:nocc]
    zvec[neleca:nocc,nocc:] = Imat[neleca:nocc,nocc:] / -ee[neleca:nocc,nocc:]

    zvec_ao = reduce(numpy.dot, (mo_coeff, zvec+zvec.T, mo_coeff.T))
    vhf = mc._scf.get_veff(mol, zvec_ao) * 2
    xvo = reduce(numpy.dot, (orbv.T, vhf, orbo))
    xvo += Imat[neleca:,:neleca] - Imat[:neleca,neleca:].T
    def fvind(x):
        x = x.reshape(xvo.shape)
        dm = reduce(numpy.dot, (orbv, x, orbo.T))
        v = mc._scf.get_veff(mol, dm + dm.T)
        v = reduce(numpy.dot, (orbv.T, v, orbo))
        return v * 2
    dm1resp = cphf.solve(fvind, mo_energy, mc._scf.mo_occ, xvo, max_cycle=30)[0]
    zvec[neleca:,:neleca] = dm1resp

    zeta = numpy.einsum('ij,j->ij', zvec, mo_energy)
    zeta = reduce(numpy.dot, (mo_coeff, zeta, mo_coeff.T))

    zvec_ao = reduce(numpy.dot, (mo_coeff, zvec+zvec.T, mo_coeff.T))
    p1 = numpy.dot(mo_coeff[:,:neleca], mo_coeff[:,:neleca].T)
    vhf_s1occ = reduce(numpy.dot, (p1, mc._scf.get_veff(mol, zvec_ao), p1))

    Imat[:ncore,ncore:neleca] = 0
    Imat[ncore:neleca,:ncore] = 0
    Imat[nocc:,neleca:nocc] = 0
    Imat[neleca:nocc,nocc:] = 0
    Imat[neleca:,:neleca] = Imat[:neleca,neleca:].T
    im1 = reduce(numpy.dot, (mo_coeff, Imat, mo_coeff.T))

    casci_dm1 = dm_core + dm_cas
    hf_dm1 = mc._scf.make_rdm1(mo_coeff, mc._scf.mo_occ)
    hcore_deriv = mc_grad.hcore_generator(mol)
    s1 = mc_grad.get_ovlp(mol)

    diag_idx = numpy.arange(nao)
    diag_idx = diag_idx * (diag_idx+1) // 2 + diag_idx
    casdm2_cc = casdm2 + casdm2.transpose(0,1,3,2)
    dm2buf = ao2mo._ao2mo.nr_e2(casdm2_cc.reshape(ncas**2,ncas**2), mo_cas.T,
                                (0, nao, 0, nao)).reshape(ncas**2,nao,nao)
    dm2buf = lib.pack_tril(dm2buf)
    dm2buf[:,diag_idx] *= .5
    dm2buf = dm2buf.reshape(ncas,ncas,nao_pair)
    casdm2 = casdm2_cc = None

    if atmlst is None:
        atmlst = range(mol.natm)
    aoslices = mol.aoslice_by_atom()