#!/usr/bin/env python
# Copyright 2014-2019 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>
#

'''
Non-relativistic unrestricted Kohn-Sham electron spin-rotation coupling
(In testing)

Refs:
    J. Phys. Chem. A. 114, 9246, 2010
    Mole. Phys. 9, 6, 585, 1964
'''

from functools import reduce
import numpy, sys
from pyscf import lib
from pyscf.lib import logger
from pyscf.dft import numint
from pyscf.prop.nmr import uks as uks_nmr
from pyscf.prop.esr import uhf as uhf_esr
from pyscf.prop.esr.uhf import _write, align
from pyscf.data import nist
from pyscf.grad import rks as rks_grad


# Note mo10 is the imaginary part of MO^1
def para(obj, mo10, mo_coeff, mo_occ, qed_fac=1):
    mol = obj.mol
    effspin = mol.spin * .5
    muB = .5  # Bohr magneton
    #qed_fac = (nist.G_ELECTRON - 1)

    orboa = mo_coeff[0][:,mo_occ[0]>0]
    orbob = mo_coeff[1][:,mo_occ[1]>0]
    dm0a = numpy.dot(orboa, orboa.T)
    dm0b = numpy.dot(orbob, orbob.T)
    dm10a = [reduce(numpy.dot, (mo_coeff[0], x, orboa.T)) for x in mo10[0]]
    dm10b = [reduce(numpy.dot, (mo_coeff[1], x, orbob.T)) for x in mo10[1]]
    dm10a = numpy.asarray([x-x.T for x in dm10a])
    dm10b = numpy.asarray([x-x.T for x in dm10b])

    hso1e = uhf_esr.make_h01_soc1e(obj, mo_coeff, mo_occ, qed_fac)
    para1e =-numpy.einsum('xji,yij->xy', dm10a, hso1e)
    para1e+= numpy.einsum('xji,yij->xy', dm10b, hso1e)
    para1e *= 1./effspin / muB
    #_write(obj, align(para1e)[0], 'SOC(1e)/OZ')

    if obj.para_soc2e:
        raise NotImplementedError('dia_soc2e = %s' % obj.dia_soc2e)

    para = para1e
    return para


# Treat Vxc as one-particle operator Vnuc
def get_vxc_soc(ni, mol, grids, xc_code, dms, max_memory=2000, verbose=None):
    xctype = ni._xc_type(xc_code)
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi=1)
    ngrids = len(grids.weights)
    BLKSIZE = numint.BLKSIZE
    blksize = min(int(max_memory/12*1e6/8/nao/BLKSIZE)*BLKSIZE, ngrids)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    vmat = numpy.zeros((2,3,nao,nao))
    if xctype == 'LDA':
        buf = numpy.empty((4,blksize,nao))
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
                                 blksize=blksize, buf=buf):
            rho_a = make_rho(0, ao[0], mask, 'LDA')
            rho_b = make_rho(1, ao[0], mask, 'LDA')
            vxc = ni.eval_xc(xc_code, (rho_a, rho_b), 1, deriv=1)[1]
            vrho = vxc[0]
            aow = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,0])
            _cross3x3_(vmat[0], mol, aow, ao[1:], mask, shls_slice, ao_loc)
            aow = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,1])
            _cross3x3_(vmat[1], mol, aow, ao[1:], mask, shls_slice, ao_loc)
            rho = vxc = vrho = aow = None

    elif xctype == 'GGA':
        buf = numpy.empty((10,blksize,nao))
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
                                 blksize=blksize, buf=buf):
            rho_a = make_rho(0, ao, mask, 'GGA')
            rho_b = make_rho(1, ao, mask, 'GGA')
            vxc = ni.eval_xc(xc_code, (rho_a,rho_b), 1, deriv=1)[1]
            wva, wvb = numint._uks_gga_wv0((rho_a, rho_b), vxc, weight)

            ip_ao = ao[1:4]
            ipip_ao = ao[4:]
            aow = rks_grad._make_dR_dao_w(ao, wva)
            _cross3x3_(vmat[0], mol, aow, ip_ao, mask, shls_slice, ao_loc)
            aow = rks_grad._make_dR_dao_w(ao, wvb)
            _cross3x3_(vmat[1], mol, aow, ip_ao, mask, shls_slice, ao_loc)
            rho = vxc = vrho = vsigma = wv = aow = None
        vmat = vmat - vmat.transpose(0,1,3,2)

    else:
        raise NotImplementedError('meta-GGA')

    return vmat


def _cross3x3_(out, mol, ao1, ao2, mask, shls_slice, ao_loc):
    out[0] += numint._dot_ao_ao(mol, ao1[1], ao2[2], mask, shls_slice, ao_loc)
    out[0] -= numint._dot_ao_ao(mol, ao1[2], ao2[1], mask, shls_slice, ao_loc)
    out[1] += numint._dot_ao_ao(mol, ao1[2], ao2[0], mask, shls_slice, ao_loc)
    out[1] -= numint._dot_ao_ao(mol, ao1[0], ao2[2], mask, shls_slice, ao_loc)
    out[2] += numint._dot_ao_ao(mol, ao1[0], ao2[1], mask, shls_slice, ao_loc)
    out[2] -= numint._dot_ao_ao(mol, ao1[1], ao2[0], mask, shls_slice, ao_loc)
    return out

# Jia, start to work here
class ESR(uhf_esr.ESR):
    '''dE = B dot gtensor dot s'''
    def __init__(self, scf_method):
        uhf_esr.ESR.__init__(self, scf_method)
        self.dia_soc2e = False
        self.para_soc2e = False

    def para(self, mo10=None, mo_coeff=None, mo_occ=None):
        if mo_coeff is None: mo_coeff = self._scf.mo_coeff
        if mo_occ is None:   mo_occ = self._scf.mo_occ
        if mo10 is None:
            self.mo10, self.mo_e10 = self.solve_mo1()
            mo10 = self.mo10
        return para(self, mo10, mo_coeff, mo_occ)

    #make_para_soc2e = make_para_soc2e
    get_fock = uks_nmr.get_fock


if __name__ == '__main__':
    from pyscf import gto, scf
    mol = gto.M(atom='H 0 0.1 0; H 0 0 1.',
                basis='ccpvdz', spin=1, charge=-1, verbose=3)
    mf = scf.UKS(mol).set(xc='bp86').run()
    esr_obj = ESR(mf)
    esr_obj.gauge_orig = (0,0,0)
    esr_obj.para_soc2e = False
    esr_obj.so_eff_charge = True
    print(esr_obj.kernel())

    mol = gto.M(atom='''
                H 0   0   1
                H 1.2 0   1
                H .1  1.1 0.3
                H .8  .7  .6
                ''',
                basis='ccpvdz', spin=1, charge=1, verbose=3)
    mf = scf.UKS(mol).set(xc='bp86').run()
    gobj = GTensor(mf)
    #print(gobj.kernel())
    gobj.para_soc2e = 'SSO'
    gobj.dia_soc2e = None
    gobj.so_eff_charge = False
    nao, nmo = mf.mo_coeff[0].shape
    nelec = mol.nelec
    numpy.random.seed(1)
    mo10 =[numpy.random.random((3,nmo,nelec[0])),
           numpy.random.random((3,nmo,nelec[1]))]
    print(lib.finger(para(gobj, mo10, mf.mo_coeff, mf.mo_occ)) - -2.1813250579863279e-05)
    numpy.random.seed(1)
    dm0 = numpy.random.random((2,nao,nao))
    dm0 = dm0 + dm0.transpose(0,2,1)
    dm10 = numpy.random.random((2,3,nao,nao))
    dm10 = dm10 - dm10.transpose(0,1,3,2)
    print(lib.finger(make_para_soc2e(gobj, dm0, dm10)) - 0.0036073897889263721)


[((49, 11, 49, 36), 'numpy.dot', 'numpy.dot', ({(49, 21, 49, 26): 'orboa', (49, 28, 49, 35): 'orboa.T'}, {}), '(orboa, orboa.T)', False, 'import numpy, sys\n'), ((50, 11, 50, 36), 'numpy.dot', 'numpy.dot', ({(50, 21, 50, 26): 'orbob', (50, 28, 50, 35): 'orbob.T'}, {}), '(orbob, orbob.T)', False, 'import numpy, sys\n'), ((53, 12, 53, 49), 'numpy.asarray', 'numpy.asarray', ({(53, 26, 53, 48): '[(x - x.T) for x in dm10a]'}, {}), '([(x - x.T) for x in dm10a])', False, 'import numpy, sys\n'), ((54, 12, 54, 49), 'numpy.asarray', 'numpy.asarray', ({(54, 26, 54, 48): '[(x - x.T) for x in dm10b]'}, {}), '([(x - x.T) for x in dm10b])', False, 'import numpy, sys\n'), ((56, 12, 56, 66), 'pyscf.prop.esr.uhf.make_h01_soc1e', 'uhf_esr.make_h01_soc1e', ({(56, 35, 56, 38): 'obj', (56, 40, 56, 48): 'mo_coeff', (56, 50, 56, 56): 'mo_occ', (56, 58, 56, 65): 'qed_fac'}, {}), '(obj, mo_coeff, mo_occ, qed_fac)', True, 'from pyscf.prop.esr import uhf as uhf_esr\n'), ((58, 13, 58, 54), 'numpy.einsum', 'numpy.einsum', ({(58, 26, 58, 39): '"""xji,yij->xy"""', (58, 41, 58, 46): 'dm10b', (58, 48, 58, 53): 'hso1e'}, {}), "('xji,yij->xy', dm10b, hso1e)", False, 'import numpy, sys\n'), ((79, 11, 79, 37), 'numpy.zeros', 'numpy.zeros', ({(79, 23, 79, 36): '(2, 3, nao, nao)'}, {}), '((2, 3, nao, nao))', False, 'import numpy, sys\n'), ((123, 14, 123, 78), 'pyscf.dft.numint._dot_ao_ao', 'numint._dot_ao_ao', ({(123, 32, 123, 35): 'mol', (123, 37, 123, 43): 'ao1[1]', (123, 45, 123, 51): 'ao2[2]', (123, 53, 123, 57): 'mask', (123, 59, 123, 69): 'shls_slice', (123, 71, 123, 77): 'ao_loc'}, {}), '(mol, ao1[1], ao2[2], mask, shls_slice, ao_loc)', False, 'from pyscf.dft import numint\n'), ((124, 14, 124, 78), 'pyscf.dft.numint._dot_ao_ao', 'numint._dot_ao_ao', ({(124, 32, 124, 35): 'mol', (124, 37, 124, 43): 'ao1[2]', (124, 45, 124, 51): 'ao2[1]', (124, 53, 124, 57): 'mask', (124, 59, 124, 69): 'shls_slice', (124, 71, 124, 77): 'ao_loc'}, {}), '(mol, ao1[2], ao2[1], mask, shls_slice, ao_loc)', False, 'from pyscf.dft import numint\n'), ((125, 14, 125, 78), 'pyscf.dft.numint._dot_ao_ao', 'numint._dot_ao_ao', ({(125, 32, 125, 35): 'mol', (125, 37, 125, 43): 'ao1[2]', (125, 45, 125, 51): 'ao2[0]', (125, 53, 125, 57): 'mask', (125, 59, 125, 69): 'shls_slice', (125, 71, 125, 77): 'ao_loc'}, {}), '(mol, ao1[2], ao2[0], mask, shls_slice, ao_loc)', False, 'from pyscf.dft import numint\n'), ((126, 14, 126, 78), 'pyscf.dft.numint._dot_ao_ao', 'numint._dot_ao_ao', ({(126, 32, 126, 35): 'mol', (126, 37, 126, 43): 'ao1[0]', (126, 45, 126, 51): 'ao2[2]', (126, 53, 126, 57): 'mask', (126, 59, 126, 69): 'shls_slice', (126, 71, 126, 77): 'ao_loc'}, {}), '(mol, ao1[0], ao2[2], mask, shls_slice, ao_loc)', False, 'from pyscf.dft import numint\n'), ((127, 14, 127, 78), 'pyscf.dft.numint._dot_ao_ao', 'numint._dot_ao_ao', ({(127, 32, 127, 35): 'mol', (127, 37, 127, 43): 'ao1[0]', (127, 45, 127, 51): 'ao2[1]', (127, 53, 127, 57): 'mask', (127, 59, 127, 69): 'shls_slice', (127, 71, 127, 77): 'ao_loc'}, {}), '(mol, ao1[0], ao2[1], mask, shls_slice, ao_loc)', False, 'from pyscf.dft import numint\n'), ((128, 14, 128, 78), 'pyscf.dft.numint._dot_ao_ao', 'numint._dot_ao_ao', ({(128, 32, 128, 35): 'mol', (128, 37, 128, 43): 'ao1[1]', (128, 45, 128, 51): 'ao2[0]', (128, 53, 128, 57): 'mask', (128, 59, 128, 69): 'shls_slice', (128, 71, 128, 77): 'ao_loc'}, {}), '(mol, ao1[1], ao2[0], mask, shls_slice, ao_loc)', False, 'from pyscf.dft import numint\n'), ((153, 10, 154, 61), 'pyscf.gto.M', 'gto.M', (), '', False, 'from pyscf import gto, scf\n'), ((162, 10, 168, 60), 'pyscf.gto.M', 'gto.M', (), '', False, 'from pyscf import gto, scf\n'), ((177, 4, 177, 24), 'numpy.random.seed', 'numpy.random.seed', ({(177, 22, 177, 23): '(1)'}, {}), '(1)', False, 'import numpy, sys\n'), ((181, 4, 181, 24), 'numpy.random.seed', 'numpy.random.seed', ({(181, 22, 181, 23): '(1)'}, {}), '(1)', False, 'import numpy, sys\n'), ((182, 10, 182, 42), 'numpy.random.random', 'numpy.random.random', ({(182, 30, 182, 41): '(2, nao, nao)'}, {}), '((2, nao, nao))', False, 'import numpy, sys\n'), ((184, 11, 184, 45), 'numpy.random.random', 'numpy.random.random', ({(184, 31, 184, 44): '(2, 3, nao, nao)'}, {}), '((2, 3, nao, nao))', False, 'import numpy, sys\n'), ((51, 13, 51, 57), 'functools.reduce', 'reduce', ({(51, 20, 51, 29): 'numpy.dot', (51, 31, 51, 56): '(mo_coeff[0], x, orboa.T)'}, {}), '(numpy.dot, (mo_coeff[0], x, orboa.T))', False, 'from functools import reduce\n'), ((52, 13, 52, 57), 'functools.reduce', 'reduce', ({(52, 20, 52, 29): 'numpy.dot', (52, 31, 52, 56): '(mo_coeff[1], x, orbob.T)'}, {}), '(numpy.dot, (mo_coeff[1], x, orbob.T))', False, 'from functools import reduce\n'), ((57, 13, 57, 54), 'numpy.einsum', 'numpy.einsum', ({(57, 26, 57, 39): '"""xji,yij->xy"""', (57, 41, 57, 46): 'dm10a', (57, 48, 57, 53): 'hso1e'}, {}), "('xji,yij->xy', dm10a, hso1e)", False, 'import numpy, sys\n'), ((81, 14, 81, 42), 'numpy.empty', 'numpy.empty', ({(81, 26, 81, 41): '(4, blksize, nao)'}, {}), '((4, blksize, nao))', False, 'import numpy, sys\n'), ((135, 8, 135, 46), 'pyscf.prop.esr.uhf.ESR.__init__', 'uhf_esr.ESR.__init__', ({(135, 29, 135, 33): 'self', (135, 35, 135, 45): 'scf_method'}, {}), '(self, scf_method)', True, 'from pyscf.prop.esr import uhf as uhf_esr\n'), ((178, 11, 178, 48), 'numpy.random.random', 'numpy.random.random', ({(178, 31, 178, 47): '(3, nmo, nelec[0])'}, {}), '((3, nmo, nelec[0]))', False, 'import numpy, sys\n'), ((179, 11, 179, 48), 'numpy.random.random', 'numpy.random.random', ({(179, 31, 179, 47): '(3, nmo, nelec[1])'}, {}), '((3, nmo, nelec[1]))', False, 'import numpy, sys\n'), ((90, 18, 90, 70), 'numpy.einsum', 'numpy.einsum', ({(90, 31, 90, 43): '"""xpi,p->xpi"""', (90, 45, 90, 51): 'ao[1:]', (90, 53, 90, 69): 'weight * vrho[:, (0)]'}, {}), "('xpi,p->xpi', ao[1:], weight * vrho[:, (0)])", False, 'import numpy, sys\n'), ((92, 18, 92, 70), 'numpy.einsum', 'numpy.einsum', ({(92, 31, 92, 43): '"""xpi,p->xpi"""', (92, 45, 92, 51): 'ao[1:]', (92, 53, 92, 69): 'weight * vrho[:, (1)]'}, {}), "('xpi,p->xpi', ao[1:], weight * vrho[:, (1)])", False, 'import numpy, sys\n'), ((97, 14, 97, 43), 'numpy.empty', 'numpy.empty', ({(97, 26, 97, 42): '(10, blksize, nao)'}, {}), '((10, blksize, nao))', False, 'import numpy, sys\n'), ((105, 23, 105, 71), 'pyscf.dft.numint._uks_gga_wv0', 'numint._uks_gga_wv0', ({(105, 43, 105, 57): '(rho_a, rho_b)', (105, 59, 105, 62): 'vxc', (105, 64, 105, 70): 'weight'}, {}), '((rho_a, rho_b), vxc, weight)', False, 'from pyscf.dft import numint\n'), ((109, 18, 109, 50), 'pyscf.grad.rks._make_dR_dao_w', 'rks_grad._make_dR_dao_w', ({(109, 42, 109, 44): 'ao', (109, 46, 109, 49): 'wva'}, {}), '(ao, wva)', True, 'from pyscf.grad import rks as rks_grad\n'), ((111, 18, 111, 50), 'pyscf.grad.rks._make_dR_dao_w', 'rks_grad._make_dR_dao_w', ({(111, 42, 111, 44): 'ao', (111, 46, 111, 49): 'wvb'}, {}), '(ao, wvb)', True, 'from pyscf.grad import rks as rks_grad\n'), ((155, 9, 155, 21), 'pyscf.scf.UKS', 'scf.UKS', ({(155, 17, 155, 20): 'mol'}, {}), '(mol)', False, 'from pyscf import gto, scf\n'), ((169, 9, 169, 21), 'pyscf.scf.UKS', 'scf.UKS', ({(169, 17, 169, 20): 'mol'}, {}), '(mol)', False, 'from pyscf import gto, scf\n')]