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

from pyscf.grad import lagrange
from pyscf.mcscf.addons import StateAverageMCSCFSolver, StateAverageFCISolver, StateAverageMixFCISolver, state_average_mix_
from pyscf.grad.mp2 import _shell_prange
from pyscf.mcscf import mc1step, mc1step_symm, newton_casscf
from pyscf.grad import casscf as casscf_grad
from pyscf.grad import rhf as rhf_grad
from pyscf.fci.direct_spin1 import _unpack_nelec
from pyscf.fci.spin_op import spin_square
from pyscf.fci import cistring
from pyscf import lib, ao2mo, mcscf
import numpy as np
import copy, time, gc
from functools import reduce
from itertools import product
from scipy import linalg

def Lorb_dot_dgorb_dx (Lorb, mc, mo_coeff=None, ci=None, atmlst=None, mf_grad=None, eris=None, verbose=None):
    ''' Modification of pyscf.grad.casscf.kernel to compute instead the orbital
    Lagrange term nuclear gradient (sum_pq Lorb_pq d2_Ecas/d_lambda d_kpq)
    This involves removing nuclear-nuclear terms and making the substitution
    (D_[p]q + D_p[q]) -> D_pq
    (d_[p]qrs + d_pq[r]s + d_p[q]rs + d_pqr[s]) -> d_pqrs
    Where [] around an index implies contraction with Lorb from the left, so that the external index
    (regardless of whether the index on the rdm is bra or ket) is always the first index of Lorb. '''

    # dmo = smoT.dao.smo
    # dao = mo.dmo.moT
    t0 = (time.clock (), time.time ())

    if mo_coeff is None: mo_coeff = mc.mo_coeff
    if ci is None: ci = mc.ci
    if mf_grad is None: mf_grad = mc._scf.nuc_grad_method()
    if mc.frozen is not None:
        raise NotImplementedError

    mol = mc.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_occ = mo_coeff[:,:nocc]
    mo_core = mo_coeff[:,:ncore]
    mo_cas = mo_coeff[:,ncore:nocc]

    # MRH: new 'effective' MO coefficients including contraction from the Lagrange multipliers
    moL_coeff = np.dot (mo_coeff, Lorb)
    s0_inv = np.dot (mo_coeff,  mo_coeff.T)
    moL_core = moL_coeff[:,:ncore]
    moL_cas = moL_coeff[:,ncore:nocc]

    # MRH: these SHOULD be state-averaged! Use the actual sacasscf object!
    casdm1, casdm2 = mc.fcisolver.make_rdm12(ci, ncas, nelecas)

    # gfock = Generalized Fock, Adv. Chem. Phys., 69, 63
    # MRH: each index exactly once!
    dm_core = np.dot(mo_core, mo_core.T) * 2
    dm_cas = reduce(np.dot, (mo_cas, casdm1, mo_cas.T))
    # MRH: new density matrix terms
    dmL_core = np.dot(moL_core, mo_core.T) * 2
    dmL_cas = reduce(np.dot, (moL_cas, casdm1, mo_cas.T))
    dmL_core += dmL_core.T
    dmL_cas += dmL_cas.T
    dm1 = dm_core + dm_cas
    dm1L = dmL_core + dmL_cas
    # MRH: end new density matrix terms
    # MRH: wrap the integral instead of the density matrix. I THINK the sign is the same!
    # mo sets 0 and 2 should be transposed, 1 and 3 should be not transposed; this will lead to correct sign
    # Except I can't do this for the external index, because the external index is contracted to ovlp matrix,
    # not the 2RDM
    aapa = np.zeros ((ncas, ncas, nmo, ncas), dtype=dm_cas.dtype)
    aapaL = np.zeros ((ncas, ncas, nmo, ncas), dtype=dm_cas.dtype)
    for i in range (nmo):
        jbuf = eris.ppaa[i]
        kbuf = eris.papa[i]
        aapa[:,:,i,:] = jbuf[ncore:nocc,:,:].transpose (1,2,0)
        aapaL[:,:,i,:] += np.tensordot (jbuf, Lorb[:,ncore:nocc], axes=((0),(0)))
        kbuf = np.tensordot (kbuf, Lorb[:,ncore:nocc], axes=((1),(0))).transpose (1,2,0)
        aapaL[:,:,i,:] += kbuf + kbuf.transpose (1,0,2)
    # MRH: new vhf terms
    vj, vk   = mc._scf.get_jk(mol, (dm_core,  dm_cas))
    vjL, vkL = mc._scf.get_jk(mol, (dmL_core, dmL_cas))
    h1 = mc.get_hcore()
    vhf_c = vj[0] - vk[0] * .5
    vhf_a = vj[1] - vk[1] * .5
    vhfL_c = vjL[0] - vkL[0] * .5
    vhfL_a = vjL[1] - vkL[1] * .5
    # MRH: I rewrote this Feff calculation completely, double-check it
    gfock  = np.dot (h1, dm1L) # h1e
    gfock += np.dot ((vhf_c + vhf_a), dmL_core) # core-core and active-core, 2nd 1RDM linked
    gfock += np.dot ((vhfL_c + vhfL_a), dm_core) # core-core and active-core, 1st 1RDM linked
    gfock += np.dot (vhfL_c, dm_cas) # core-active, 1st 1RDM linked
    gfock += np.dot (vhf_c, dmL_cas) # core-active, 2nd 1RDM linked
    gfock  = np.dot (s0_inv, gfock) # Definition of quantity is in MO's; going (AO->MO->AO) incurs an inverse ovlp
    gfock += reduce (np.dot, (mo_coeff, np.einsum('uviw,uvtw->it', aapaL, casdm2), mo_cas.T)) # active-active
    # MRH: I have to contract this external 2RDM index explicitly on the 2RDM but fortunately I can do so here
    gfock += reduce (np.dot, (mo_coeff, np.einsum('uviw,vuwt->it', aapa, casdm2), moL_cas.T))
    # MRH: As of 04/18/2019, the two-body part of this is including aapaL is definitely, unambiguously correct
    dme0 = (gfock+gfock.T)/2 # This transpose is for the overlap matrix later on
    aapa = vj = vk = vhf_c = vhf_a = None

    vj, vk = mf_grad.get_jk (mol, (dm_core, dm_cas, dmL_core, dmL_cas))
    vhf1c, vhf1a, vhf1cL, vhf1aL = vj - vk * 0.5
    #vhf1c, vhf1a, vhf1cL, vhf1aL = mf_grad.get_veff(mol, (dm_core, dm_cas, dmL_core, dmL_cas))
    hcore_deriv = mf_grad.hcore_generator(mol)
    s1 = mf_grad.get_ovlp(mol)

    diag_idx = np.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)
    # MRH: contract the final two indices of the active-active 2RDM with L as you change to AOs
    # note tensordot always puts indices in the order of the arguments.
    dm2Lbuf = np.zeros ((ncas**2,nmo,nmo))
    # MRH: The second line below transposes the L; the third line transposes the derivative later on
    # Both the L and the derivative have to explore all indices
    dm2Lbuf[:,:,ncore:nocc]  = np.tensordot (Lorb[:,ncore:nocc], casdm2, axes=(1,2)).transpose (1,2,0,3).reshape (ncas**2,nmo,ncas)
    dm2Lbuf[:,ncore:nocc,:] += np.tensordot (Lorb[:,ncore:nocc], casdm2, axes=(1,3)).transpose (1,2,3,0).reshape (ncas**2,ncas,nmo)
    dm2Lbuf += dm2Lbuf.transpose (0,2,1)
    dm2Lbuf = np.ascontiguousarray (dm2Lbuf)
    dm2Lbuf = ao2mo._ao2mo.nr_e2(dm2Lbuf.reshape (ncas**2,nmo**2), mo_coeff.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)
    dm2Lbuf = lib.pack_tril(dm2Lbuf)
    dm2Lbuf[:,diag_idx] *= .5
    dm2Lbuf = dm2Lbuf.reshape(ncas,ncas,nao_pair)

    if atmlst is None:
        atmlst = list (range(mol.natm))
    aoslices = mol.aoslice_by_atom()
    de_hcore = np.zeros((len(atmlst),3))
    de_renorm = np.zeros((len(atmlst),3))
    de_eri = np.zeros((len(atmlst),3))
    de = np.zeros((len(atmlst),3))

    max_memory = mc.max_memory - lib.current_memory()[0]
    blksize = int(max_memory*.9e6/8 / (4*(aoslices[:,3]-aoslices[:,2]).max()*nao_pair))
    # MRH: 3 components of eri array and 1 density matrix array: FOUR arrays of this size are required!
    blksize = min(nao, max(2, blksize))
    lib.logger.info (mc, 'SA-CASSCF Lorb_dot_dgorb memory remaining for eri manipulation: {} MB; using blocksize = {}'.format (max_memory, blksize))
    t0 = lib.logger.timer (mc, 'SA-CASSCF Lorb_dot_dgorb 1-electron part', *t0)

    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        h1ao = hcore_deriv(ia)
        # MRH: h1e and Feff terms
        de_hcore[k] += np.einsum('xij,ij->x', h1ao, dm1L)
        de_renorm[k] -= np.einsum('xij,ij->x', s1[:,p0:p1], dme0[p0:p1]) * 2

        q1 = 0
        for b0, b1, nf in _shell_prange(mol, 0, mol.nbas, blksize):
            q0, q1 = q1, q1 + nf
            dm2_ao  = lib.einsum('ijw,pi,qj->pqw', dm2Lbuf, mo_cas[p0:p1], mo_cas[q0:q1])
            # MRH: now contract the first two indices of the active-active 2RDM with L as you go from MOs to AOs
            dm2_ao += lib.einsum('ijw,pi,qj->pqw', dm2buf, moL_cas[p0:p1], mo_cas[q0:q1])
            dm2_ao += lib.einsum('ijw,pi,qj->pqw', dm2buf, mo_cas[p0:p1], moL_cas[q0:q1])
            shls_slice = (shl0,shl1,b0,b1,0,mol.nbas,0,mol.nbas)
            gc.collect ()
            eri1 = mol.intor('int2e_ip1', comp=3, aosym='s2kl',
                             shls_slice=shls_slice).reshape(3,p1-p0,nf,nao_pair)
            # MRH: I still don't understand why there is a minus here!
            de_eri[k] -= np.einsum('xijw,ijw->x', eri1, dm2_ao) * 2
            eri1 = dm2_ao = None
            gc.collect ()
            t0 = lib.logger.timer (mc, 'SA-CASSCF Lorb_dot_dgorb atom {} ({},{}|{})'.format (ia, p1-p0, nf, nao_pair), *t0)
        # MRH: core-core and core-active 2RDM terms
        de_eri[k] += np.einsum('xij,ij->x', vhf1c[:,p0:p1], dm1L[p0:p1]) * 2
        de_eri[k] += np.einsum('xij,ij->x', vhf1cL[:,p0:p1], dm1[p0:p1]) * 2
        # MRH: active-core 2RDM terms
        de_eri[k] += np.einsum('xij,ij->x', vhf1a[:,p0:p1], dmL_core[p0:p1]) * 2
        de_eri[k] += np.einsum('xij,ij->x', vhf1aL[:,p0:p1], dm_core[p0:p1]) * 2

    # MRH: deleted the nuclear-nuclear part to avoid double-counting
    # lesson learned from debugging - mol.intor computes -1 * the derivative and only
    # for one index
    # on the other hand, mf_grad.hcore_generator computes the actual derivative of
    # h1 for both indices and with the correct sign

    lib.logger.debug (mc, "Orb lagrange hcore component:\n{}".format (de_hcore))
    lib.logger.debug (mc, "Orb lagrange renorm component:\n{}".format (de_renorm))
    lib.logger.debug (mc, "Orb lagrange eri component:\n{}".format (de_eri))
    de = de_hcore + de_renorm + de_eri

    return de

def Lci_dot_dgci_dx (Lci, weights, mc, mo_coeff=None, ci=None, atmlst=None, mf_grad=None, eris=None, verbose=None):
    ''' Modification of pyscf.grad.casscf.kernel to compute instead the CI
    Lagrange term nuclear gradient (sum_IJ Lci_IJ d2_Ecas/d_lambda d_PIJ)
    This involves removing all core-core and nuclear-nuclear terms and making the substitution
    sum_I w_I<L_I|p'q|I> + c.c. -> <0|p'q|0>
    sum_I w_I<L_I|p'r'sq|I> + c.c. -> <0|p'r'sq|0>
    The active-core terms (sum_I w_I<L_I|x'iyi|I>, sum_I w_I <L_I|x'iiy|I>, c.c.) must be retained.'''
    if mo_coeff is None: mo_coeff = mc.mo_coeff
    if ci is None: ci = mc.ci
    if mf_grad is None: mf_grad = mc._scf.nuc_grad_method()
    if mc.frozen is not None:
        raise NotImplementedError

    t0 = (time.clock (), time.time ())
    mol = mc.mol
    ncore = mc.ncore
    ncas = mc.ncas
    nocc = ncore + ncas
    nelecas = mc.nelecas
    nao, nmo = mo_coeff.shape
    nao_pair = nao * (nao+1) // 2
    nroots = len (ci)

    mo_occ = mo_coeff[:,:nocc]
    mo_core = mo_coeff[:,:ncore]
    mo_cas = mo_coeff[:,ncore:nocc]

    # MRH: TDMs + c.c. instead of RDMs; 06/30/2020: new interface in mcscf.addons makes this much more transparent
    casdm1, casdm2 = mc.fcisolver.trans_rdm12 (Lci, ci, ncas, nelecas)
    casdm1 += casdm1.transpose (1,0)
    casdm2 += casdm2.transpose (1,0,3,2)

# gfock = Generalized Fock, Adv. Chem. Phys., 69, 63
    dm_core = np.dot(mo_core, mo_core.T) * 2
    dm_cas = reduce(np.dot, (mo_cas, casdm1, mo_cas.T))
    aapa = np.zeros ((ncas, ncas, nmo, ncas), dtype=dm_cas.dtype)
    for i in range (nmo):
        aapa[:,:,i,:] = eris.ppaa[i][ncore:nocc,:,:].transpose (1,2,0)
    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
    # MRH: delete h1 + vhf_c from the first line below (core and core-core stuff)
    # Also extend gfock to span the whole space
    gfock = np.zeros_like (dm_cas)
    gfock[:,:nocc]   = reduce(np.dot, (mo_coeff.T, vhf_a, mo_occ)) * 2
    gfock[:,ncore:nocc]  = reduce(np.dot, (mo_coeff.T, h1 + vhf_c, mo_cas, casdm1))
    gfock[:,ncore:nocc] += np.einsum('uvpw,vuwt->pt', aapa, casdm2)
    dme0 = reduce(np.dot, (mo_coeff, (gfock+gfock.T)*.5, mo_coeff.T))
    aapa = vj = vk = vhf_c = vhf_a = h1 = gfock = None

    vj, vk = mf_grad.get_jk (mol, (dm_core, dm_cas))
    vhf1c, vhf1a = vj - vk * 0.5
    #vhf1c, vhf1a = mf_grad.get_veff(mol, (dm_core, dm_cas))
    hcore_deriv = mf_grad.hcore_generator(mol)
    s1 = mf_grad.get_ovlp(mol)

    diag_idx = np.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()
    de_hcore = np.zeros((len(atmlst),3))
    de_renorm = np.zeros((len(atmlst),3))
    de_eri = np.zeros((len(atmlst),3))
    de = np.zeros((len(atmlst),3))

    max_memory = mc.max_memory - lib.current_memory()[0]
    blksize = int(max_memory*.9e6/8 / (4*(aoslices[:,3]-aoslices[:,2]).max()*nao_pair))
    # MRH: 3 components of eri array and 1 density matrix array: FOUR arrays of this size are required!
    blksize = min(nao, max(2, blksize))
    lib.logger.info (mc, 'SA-CASSCF Lci_dot_dgci memory remaining for eri manipulation: {} MB; using blocksize = {}'.format (max_memory, blksize))
    t0 = lib.logger.timer (mc, 'SA-CASSCF Lci_dot_dgci 1-electron part', *t0)

    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        h1ao = hcore_deriv(ia)
        # MRH: dm1 -> dm_cas in the line below
        de_hcore[k] += np.einsum('xij,ij->x', h1ao, dm_cas)
        de_renorm[k] -= np.einsum('xij,ij->x', s1[:,p0:p1], dme0[p0:p1]) * 2

        q1 = 0
        for b0, b1, nf in _shell_prange(mol, 0, mol.nbas, blksize):
            q0, q1 = q1, q1 + nf
            dm2_ao = lib.einsum('ijw,pi,qj->pqw', dm2buf, mo_cas[p0:p1], mo_cas[q0:q1])
            shls_slice = (shl0,shl1,b0,b1,0,mol.nbas,0,mol.nbas)
            gc.collect ()
            eri1 = mol.intor('int2e_ip1', comp=3, aosym='s2kl',
                             shls_slice=shls_slice).reshape(3,p1-p0,nf,nao_pair)
            de_eri[k] -= np.einsum('xijw,ijw->x', eri1, dm2_ao) * 2
            eri1 = dm2_ao = None
            gc.collect ()
            t0 = lib.logger.timer (mc, 'SA-CASSCF Lci_dot_dgci atom {} ({},{}|{})'.format (ia, p1-p0, nf, nao_pair), *t0)
        # MRH: dm1 -> dm_cas in the line below. Also eliminate core-core terms
        de_eri[k] += np.einsum('xij,ij->x', vhf1c[:,p0:p1], dm_cas[p0:p1]) * 2
        de_eri[k] += np.einsum('xij,ij->x', vhf1a[:,p0:p1], dm_core[p0:p1]) * 2

    lib.logger.debug (mc, "CI lagrange hcore component:\n{}".format (de_hcore))
    lib.logger.debug (mc, "CI lagrange renorm component:\n{}".format (de_renorm))
    lib.logger.debug (mc, "CI lagrange eri component:\n{}".format (de_eri))
    de = de_hcore + de_renorm + de_eri
    return de

def as_scanner(mcscf_grad, state=None):
    '''Generating a nuclear gradients scanner/solver (for geometry optimizer).

    The returned solver is a function. This function requires one argument
    "mol" as input and returns energy and first order nuclear derivatives.

    The solver will automatically use the results of last calculation as the
    initial guess of the new calculation.  All parameters assigned in the
    nuc-grad object and SCF object (DIIS, conv_tol, max_memory etc) are
    automatically applied in the solver.

    Note scanner has side effects.  It may change many underlying objects
    (_scf, with_df, with_x2c, ...) during calculation.

    Examples:

    >>> from pyscf import gto, scf, mcscf
    >>> mol = gto.M(atom='N 0 0 0; N 0 0 1.1', verbose=0)
    >>> mc_grad_scanner = mcscf.CASSCF(scf.RHF(mol), 4, 4).nuc_grad_method().as_scanner()
    >>> etot, grad = mc_grad_scanner(gto.M(atom='N 0 0 0; N 0 0 1.1'))
    >>> etot, grad = mc_grad_scanner(gto.M(atom='N 0 0 0; N 0 0 1.5'))
    '''
    from pyscf import gto
    if isinstance(mcscf_grad, lib.GradScanner):
        return mcscf_grad

    #if state is None and (not hasattr (mcscf_grad, 'state') or (mcscf_grad.state is None)):
    #    return casscf_grad.as_scanner (mcscf_grad)

    lib.logger.info(mcscf_grad, 'Create scanner for %s', mcscf_grad.__class__)

    class CASSCF_GradScanner(mcscf_grad.__class__, lib.GradScanner):
        def __init__(self, g):
            lib.GradScanner.__init__(self, g)
            if state is None:
                self.state = g.state
            else:
                self.state = state
        def __call__(self, mol_or_geom, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)
            if 'state' in kwargs: self.state = kwargs['state']
            mc_scanner = self.base
            e_tot = mc_scanner(mol)
            if hasattr (mc_scanner, 'e_mcscf'): self.e_mcscf = mc_scanner.e_mcscf
            if hasattr (mc_scanner, 'e_states') and self.state is not None:
                e_tot = mc_scanner.e_states[self.state]
            self.mol = mol
            if not ('state' in kwargs):
                kwargs['state'] = self.state
            de = self.kernel(**kwargs)
            return e_tot, de

    return CASSCF_GradScanner(mcscf_grad)

class Gradients (lagrange.Gradients):

    def __init__(self, mc, state=None):
        self.__dict__.update (mc.__dict__)
        nmo = mc.mo_coeff.shape[-1]
        self.ngorb = np.count_nonzero (mc.uniq_var_indices (nmo, mc.ncore, mc.ncas, mc.frozen))
        self.nroots = mc.fcisolver.nroots
        neleca, nelecb = _unpack_nelec (mc.nelecas)
        self.spin_states = [neleca - nelecb,] * self.nroots
        self.na_states = [cistring.num_strings (mc.ncas, neleca),] * self.nroots
        self.nb_states = [cistring.num_strings (mc.ncas, nelecb),] * self.nroots
        if isinstance (mc.fcisolver, StateAverageMixFCISolver):
            self.nroots = p0 = 0
            for solver in mc.fcisolver.fcisolvers:
                self.nroots += solver.nroots
                nea, neb = mc.fcisolver._get_nelec (solver, (neleca, nelecb))
                self.spin_states[p0:self.nroots] = (nea - neb for x in range (solver.nroots))
                self.na_states[p0:self.nroots] = (cistring.num_strings (mc.ncas, nea) for x in range (solver.nroots))
                self.nb_states[p0:self.nroots] = (cistring.num_strings (mc.ncas, neb) for x in range (solver.nroots))
                p0 = self.nroots
        self.nci = sum ([na * nb for na, nb in zip (self.na_states, self.nb_states)])
        if state is not None:
            self.state = state
        elif hasattr (mc, 'nuc_grad_state'):
            self.state = mc.nuc_grad_state
        else:
            self.state = None
        self.eris = None
        self.weights = np.array ([1])
        try:
            self.e_states = np.asarray (mc.e_states)
        except AttributeError as e:
            self.e_states = np.asarray (mc.e_tot)
        if isinstance (mc, StateAverageMCSCFSolver):
            self.weights = np.asarray (mc.weights)
        assert (len (self.weights) == self.nroots), '{} {} {}'.format (mc.fcisolver.__class__, self.weights, self.nroots)
        lagrange.Gradients.__init__(self, mc, self.ngorb+self.nci)
        self.max_cycle = mc.max_cycle_macro

    def pack_uniq_var (self, xorb, xci):
        # TODO: point-group symmetry of the xci components? CSFs?
        xorb = self.base.pack_uniq_var (xorb)
        xci = np.concatenate ([x.ravel () for x in xci])
        return np.append (xorb, xci)

    def unpack_uniq_var (self, x):
        # TODO: point-group symmetry of the xci components? CSFs?
        xorb, x = self.base.unpack_uniq_var (x[:self.ngorb]), x[self.ngorb:]
        xci = []
        for na, nb in zip (self.na_states, self.nb_states):
            xci.append (x[:na*nb].reshape (na, nb))
            x = x[na*nb:]
        return xorb, xci

    def make_fcasscf (self, state=None, casscf_attr={}, fcisolver_attr={}):
        ''' Make a fake CASSCF object for ostensible single-state calculations '''
        fcasscf = mcscf.CASSCF (self.base._scf, self.base.ncas, self.base.nelecas)
        fcasscf.__dict__.update (self.base.__dict__)

        if isinstance (fcasscf.fcisolver, StateAverageFCISolver):
            if isinstance (fcasscf.fcisolver, StateAverageMixFCISolver):
                p0 = 0
                for solver in fcasscf.fcisolver.fcisolvers:
                    p1 = p0 + solver.nroots
                    if p0 <= state < p1:
                        solver_class = solver.__class__
                        solver_obj = solver
                        break
                    p0 = p1
            else:
                solver_class = self.base.fcisolver._base_class
                solver_obj = self.base.fcisolver
            fcasscf.fcisolver = solver_class (self.base.mol)
            fcasscf.fcisolver.__dict__.update (solver_obj.__dict__)
            fcasscf.fcisolver.nroots = 1
        fcasscf.__dict__.update (casscf_attr)
        fcasscf.fcisolver.__dict__.update (fcisolver_attr)
        fcasscf.verbose, fcasscf.stdout = self.verbose, self.stdout
        fcasscf._tag_gfock_ov_nonzero = True
        return fcasscf

    def make_fcasscf_sa (self, casscf_attr={}, fcisolver_attr={}):
        ''' Make a fake SA-CASSCF object to get around weird inheritance conflicts '''
        fcasscf = self.make_fcasscf (state=0, casscf_attr={}, fcisolver_attr={})
        fcasscf.__dict__.update (self.base.__dict__)
        if isinstance (self.base, StateAverageMCSCFSolver):
            if isinstance (self.base.fcisolver, StateAverageMixFCISolver):
                fcasscf = state_average_mix_(fcasscf, self.base.fcisolver.fcisolvers, self.base.weights)
            else:
                fcasscf.state_average_(self.base.weights)
        fcasscf.__dict__.update (casscf_attr)
        fcasscf.fcisolver.__dict__.update (fcisolver_attr)
        return fcasscf

    def kernel (self, state=None, atmlst=None, verbose=None, mo=None, ci=None, eris=None, mf_grad=None, e_states=None, level_shift=None, **kwargs):
        if state is None: state = self.state
        if atmlst is None: atmlst = self.atmlst
        if verbose is None: verbose = self.verbose
        if mo is None: mo = self.base.mo_coeff
        if ci is None: ci = self.base.ci
        if eris is None:
            eris = self.eris = self.base.ao2mo (mo)
        if mf_grad is None: mf_grad = self.base._scf.nuc_grad_method ()
        if state is None:
            return casscf_grad.Gradients (self.base).kernel (mo_coeff=mo, ci=ci, atmlst=atmlst, verbose=verbose)
        if e_states is None:
            try:
                e_states = self.e_states = np.asarray (self.base.e_states)
            except AttributeError as e:
                e_states = self.e_states = np.asarray (self.base.e_tot)
        if level_shift is None: level_shift=self.level_shift
        return lagrange.Gradients.kernel (self, state=state, atmlst=atmlst, verbose=verbose, mo=mo, ci=ci, eris=eris, mf_grad=mf_grad, e_states=e_states, level_shift=level_shift, **kwargs)

    def get_wfn_response (self, atmlst=None, state=None, verbose=None, mo=None, ci=None, **kwargs):
        if state is None: state = self.state
        if atmlst is None: atmlst = self.atmlst
        if verbose is None: verbose = self.verbose
        if mo is None: mo = self.base.mo_coeff
        if ci is None: ci = self.base.ci
        ndet = self.na_states[state] * self.nb_states[state]
        fcasscf = self.make_fcasscf (state)
        fcasscf.mo_coeff = mo
        fcasscf.ci = ci[state]
        eris = fcasscf.ao2mo (mo)
        g_all_state = newton_casscf.gen_g_hop (fcasscf, mo, ci[state], eris, verbose)[0]
        g_all = np.zeros (self.nlag)
        g_all[:self.ngorb] = g_all_state[:self.ngorb]
        # No need to reshape or anything, just use the magic of repeated slicing
        offs = sum ([na * nb for na, nb in zip (self.na_states[:state],
            self.nb_states[:state])]) if state > 0 else 0
        g_all[self.ngorb:][offs:][:ndet] = g_all_state[self.ngorb:]
        return g_all

    def get_Aop_Adiag (self, atmlst=None, state=None, verbose=None, mo=None, ci=None, eris=None, level_shift=None, **kwargs):
        if state is None: state = self.state
        if atmlst is None: atmlst = self.atmlst
        if verbose is None: verbose = self.verbose
        if mo is None: mo = self.base.mo_coeff
        if ci is None: ci = self.base.ci
        if eris is None and self.eris is None:
            eris = self.eris = self.base.ao2mo (mo)
        elif eris is None:
            eris = self.eris
        if not isinstance (self.base, StateAverageMCSCFSolver) and isinstance (ci, list): ci = ci[0]
        fcasscf = self.make_fcasscf_sa ()
        Aop, Adiag = newton_casscf.gen_g_hop (fcasscf, mo, ci, eris, verbose)[2:]
        # Eliminate the component of Aop (x) which is parallel to the state-average space
        # The Lagrange multiplier equations are not defined there
        return self.project_Aop (Aop, ci, state), Adiag

    def get_ham_response (self, state=None, atmlst=None, verbose=None, mo=None, ci=None, eris=None, mf_grad=None, **kwargs):
        if state is None: state = self.state
        if atmlst is None: atmlst = self.atmlst
        if verbose is None: verbose = self.verbose
        if mo is None: mo = self.base.mo_coeff
        if ci is None: ci = self.base.ci
        if eris is None and self.eris is None:
            eris = self.eris = self.base.ao2mo (mo)
        elif eris is None:
            eris = self.eris
        fcasscf_grad = casscf_grad.Gradients (self.make_fcasscf (state))
        return fcasscf_grad.kernel (mo_coeff=mo, ci=ci[state], atmlst=atmlst, verbose=verbose)

    def get_LdotJnuc (self, Lvec, state=None, atmlst=None, verbose=None, mo=None, ci=None, eris=None, mf_grad=None, **kwargs):
        if state is None: state = self.state
        if atmlst is None: atmlst = self.atmlst
        if verbose is None: verbose = self.verbose
        if mo is None: mo = self.base.mo_coeff
        if ci is None: ci = self.base.ci[state]
        if eris is None and self.eris is None:
            eris = self.eris = self.base.ao2mo (mo)
        elif eris is None:
            eris = self.eris
        ncas = self.base.ncas
        nelecas = self.base.nelecas
        if getattr(self.base.fcisolver, 'gen_linkstr', None):
            linkstr  = self.base.fcisolver.gen_linkstr(ncas, nelecas, False)
        else:
            linkstr  = None

        # Just sum the weights now... Lorb can be implicitly summed
        # Lci may be in the csf basis
        Lorb, Lci = self.unpack_uniq_var (Lvec)
        #Lorb = self.base.unpack_uniq_var (Lvec[:self.ngorb])
        #Lci = Lvec[self.ngorb:].reshape (self.nroots, -1)
        #ci = np.ravel (ci).reshape (self.nroots, -1)

        # CI part
        t0 = (time.clock (), time.time ())
        de_Lci = Lci_dot_dgci_dx (Lci, self.weights, self.base, mo_coeff=mo, ci=ci, atmlst=atmlst, mf_grad=mf_grad, eris=eris, verbose=verbose)
        lib.logger.info (self, '--------------- %s gradient Lagrange CI response ---------------',
                    self.base.__class__.__name__)
        if verbose >= lib.logger.INFO: rhf_grad._write(self, self.mol, de_Lci, atmlst)
        lib.logger.info (self, '----------------------------------------------------------------')
        t0 = lib.logger.timer (self, '{} gradient Lagrange CI response'.format (self.base.__class__.__name__), *t0)

        # Orb part
        de_Lorb = Lorb_dot_dgorb_dx (Lorb, self.base, mo_coeff=mo, ci=ci, atmlst=atmlst, mf_grad=mf_grad, eris=eris, verbose=verbose)
        lib.logger.info (self, '--------------- %s gradient Lagrange orbital response ---------------',
                    self.base.__class__.__name__)
        if verbose >= lib.logger.INFO: rhf_grad._write(self, self.mol, de_Lorb, atmlst)
        lib.logger.info (self, '----------------------------------------------------------------------')
        t0 = lib.logger.timer (self, '{} gradient Lagrange orbital response'.format (self.base.__class__.__name__), *t0)

        return de_Lci + de_Lorb

    def debug_lagrange (self, Lvec, bvec, Aop, Adiag, state=None, mo=None, ci=None, **kwargs):
        # This needs to be rewritten substantially to work properly with state_average_mix
        if state is None: state = self.state
        if mo is None: mo = self.base.mo_coeff
        if ci is None: ci = self.base.ci
        def _debug_cispace (xci, label):
            xci_norm = [np.dot (c.ravel (), c.ravel ()) for c in xci]
            try:
                xci_ss = self.base.fcisolver.states_spin_square (xci, self.base.ncas, self.base.nelecas)[0]
            except AttributeError:
                nelec = sum (_unpack_nelec (self.base.nelecas))
                xci_ss = [spin_square (x, self.base.ncas, ((nelec+m)//2,(nelec-m)//2))[0]
                    for x, m in zip (xci, self.spin_states)]
            xci_ss = [x / max (y, 1e-8) for x, y in zip (xci_ss, xci_norm)]
            xci_multip = [np.sqrt (x+.25) - .5 for x in xci_ss]
            for ix, (norm, ss, multip) in enumerate (zip (xci_norm, xci_ss, xci_multip)):
                lib.logger.debug (self,
                    ' State {} {} norm = {:.7e} ; <S^2> = {:.7f} ; 2S+1 = {:.7f}'.format
                    (ix, label, norm, ss, multip))
        borb, bci = self.unpack_uniq_var (bvec)
        lib.logger.debug (self, 'Orbital rotation gradient norm = {:.7e}'.format (linalg.norm (borb)))
        _debug_cispace (bci, 'CI gradient')
        Aorb, Aci = self.unpack_uniq_var (Adiag)
        lib.logger.debug (self, 'Orbital rotation Hamiltonian diagonal norm = {:.7e}'.format (linalg.norm (Aorb)))
        _debug_cispace (Aci, 'Hamiltonian diagonal')
        Lorb, Lci = self.unpack_uniq_var (Lvec)
        lib.logger.debug (self, 'Orbital rotation Lagrange vector norm = {:.7e}'.format (linalg.norm (Lorb)))
        _debug_cispace (Lci, 'Lagrange vector')
        #lib.logger.info (self, '{} gradient: state = {}'.format (self.base.__class__.__name__, state))
        #ngorb = self.ngorb
        #nci = self.nci
        #nroots = self.nroots
        #ndet = nci // nroots
        #ncore = self.base.ncore
        #ncas = self.base.ncas
        #nelecas = self.base.nelecas
        #nocc = ncore + ncas
        #nlag = self.nlag
        #ci = np.asarray (self.base.ci).reshape (nroots, -1)
        #err = Aop (Lvec) + bvec
        #eorb = self.base.unpack_uniq_var (err[:ngorb])
        #eci = err[ngorb:].reshape (nroots, -1)
        #borb = self.base.unpack_uniq_var (bvec[:ngorb])
        #bci = bvec[ngorb:].reshape (nroots, -1)
        #Lorb = self.base.unpack_uniq_var (Lvec[:ngorb])
        #Lci = Lvec[ngorb:].reshape (nroots, ndet)
        #Aci = Adiag[ngorb:].reshape (nroots, ndet)
        #Lci_ci_ovlp = (np.asarray (ci).reshape (nroots,-1).conjugate () @ Lci.T).T
        #Lci_Lci_ovlp = (Lci.conjugate () @ Lci.T).T
        #eci_ci_ovlp = (np.asarray (ci).reshape (nroots,-1).conjugate () @ eci.T).T
        #bci_ci_ovlp = (np.asarray (ci).reshape (nroots,-1).conjugate () @ bci.T).T
        #ci_ci_ovlp = ci.conjugate () @ ci.T
        #lib.logger.debug (self, "{} gradient RHS, inactive-active orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, borb[:ncore,ncore:nocc]))
        #lib.logger.debug (self, "{} gradient RHS, inactive-external orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, borb[:ncore,nocc:]))
        #lib.logger.debug (self, "{} gradient RHS, active-external orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, borb[ncore:nocc,nocc:]))
        #lib.logger.debug (self, "{} gradient residual, inactive-active orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, eorb[:ncore,ncore:nocc]))
        #lib.logger.debug (self, "{} gradient residual, inactive-external orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, eorb[:ncore,nocc:]))
        #lib.logger.debug (self, "{} gradient residual, active-external orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, eorb[ncore:nocc,nocc:]))
        #lib.logger.debug (self, "{} gradient Lagrange factor, inactive-active orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, Lorb[:ncore,ncore:nocc]))
        #lib.logger.debug (self, "{} gradient Lagrange factor, inactive-external orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, Lorb[:ncore,nocc:]))
        #lib.logger.debug (self, "{} gradient Lagrange factor, active-external orbital rotations:\n{}".format (
        #    self.base.__class__.__name__, Lorb[ncore:nocc,nocc:]))
        #'''
        #lib.logger.debug (self, "{} gradient RHS, inactive-inactive orbital rotations (redundant!):\n{}".format (
        #    self.base.__class__.__name__, borb[:ncore,:ncore]))
        #lib.logger.debug (self, "{} gradient RHS, active-active orbital rotations (redundant!):\n{}".format (
        #    self.base.__class__.__name__, borb[ncore:nocc,ncore:nocc]))
        #lib.logger.debug (self, "{} gradient RHS, external-external orbital rotations (redundant!):\n{}".format (
        #    self.base.__class__.__name__, borb[nocc:,nocc:]))
        #lib.logger.debug (self, "{} gradient Lagrange factor, inactive-inactive orbital rotations (redundant!):\n{}".format (
        #    self.base.__class__.__name__, Lorb[:ncore,:ncore]))
        #lib.logger.debug (self, "{} gradient Lagrange factor, active-active orbital rotations (redundant!):\n{}".format (
        #    self.base.__class__.__name__, Lorb[ncore:nocc,ncore:nocc]))
        #lib.logger.debug (self, "{} gradient Lagrange factor, external-external orbital rotations (redundant!):\n{}".format (
        #    self.base.__class__.__name__, Lorb[nocc:,nocc:]))
        #'''
        #lib.logger.debug (self, "{} gradient Lagrange factor, CI part overlap with true CI SA space:\n{}".format (
        #    self.base.__class__.__name__, Lci_ci_ovlp))
        #lib.logger.debug (self, "{} gradient Lagrange factor, CI part self overlap matrix:\n{}".format (
        #    self.base.__class__.__name__, Lci_Lci_ovlp))
        #lib.logger.debug (self, "{} gradient Lagrange factor, CI vector self overlap matrix:\n{}".format (
        #    self.base.__class__.__name__, ci_ci_ovlp))
        #lib.logger.debug (self, "{} gradient Lagrange factor, CI part response overlap with SA space:\n{}".format (
        #    self.base.__class__.__name__, bci_ci_ovlp))
        #lib.logger.debug (self, "{} gradient Lagrange factor, CI part residual overlap with SA space:\n{}".format (
        #    self.base.__class__.__name__, eci_ci_ovlp))
        #neleca, nelecb = _unpack_nelec (nelecas)
        #spin = neleca - nelecb + 1
        #csf = CSFTransformer (ncas, neleca, nelecb, spin)
        #ecsf = csf.vec_det2csf (eci, normalize=False, order='C')
        #err_norm_det = linalg.norm (err)
        #err_norm_csf = linalg.norm (np.append (eorb, ecsf.ravel ()))
        #lib.logger.debug (self, "{} gradient: determinant residual = {}, CSF residual = {}".format (
        #    self.base.__class__.__name__, err_norm_det, err_norm_csf))
        #ci_lbls, ci_csf   = csf.printable_largest_csf (ci,  10, isdet=True, normalize=True,  order='C')
        #bci_lbls, bci_csf = csf.printable_largest_csf (bci, 10, isdet=True, normalize=False, order='C')
        #eci_lbls, eci_csf = csf.printable_largest_csf (eci, 10, isdet=True, normalize=False, order='C')
        #Lci_lbls, Lci_csf = csf.printable_largest_csf (Lci, 10, isdet=True, normalize=False, order='C')
        #Aci_lbls, Aci_csf = csf.printable_largest_csf (Aci, 10, isdet=True, normalize=False, order='C')
        #ncsf = bci_csf.shape[1]
        #for iroot in range (self.nroots):
        #    lib.logger.debug (self, "{} gradient Lagrange factor, CI part root {} spin square: {}".format (
        #        self.base.__class__.__name__, iroot, spin_square (Lci[iroot], ncas, nelecas)))
        #    lib.logger.debug (self, "Base CI vector")
        #    for icsf in range (ncsf):
        #        lib.logger.debug (self, '{} {}'.format (ci_lbls[iroot,icsf], ci_csf[iroot,icsf]))
        #    lib.logger.debug (self, "CI gradient:")
        #    for icsf in range (ncsf):
        #        lib.logger.debug (self, '{} {}'.format (bci_lbls[iroot,icsf], bci_csf[iroot,icsf]))
        #    lib.logger.debug (self, "CI residual:")
        #    for icsf in range (ncsf):
        #        lib.logger.debug (self, '{} {}'.format (eci_lbls[iroot,icsf], eci_csf[iroot,icsf]))
        #    lib.logger.debug (self, "CI Lagrange vector:")
        #    for icsf in range (ncsf):
        #        lib.logger.debug (self, '{} {}'.format (Lci_lbls[iroot,icsf], Lci_csf[iroot,icsf]))
        #    lib.logger.debug (self, "Diagonal of Hessian matrix CI part:")
        #    for icsf in range (ncsf):
        #        lib.logger.debug (self, '{} {}'.format (Aci_lbls[iroot,icsf], Aci_csf[iroot,icsf]))
        #'''
        #Afull = np.zeros ((nlag, nlag))
        #dum = np.zeros ((nlag))
        #for ix in range (nlag):
        #    dum[ix] = 1
        #    Afull[ix,:] = Aop (dum)
        #    dum[ix] = 0
        #Afull_orborb = Afull[:ngorb,:ngorb]
        #Afull_orbci = Afull[:ngorb,ngorb:].reshape (ngorb, nroots, ndet)
        #Afull_ciorb = Afull[ngorb:,:ngorb].reshape (nroots, ndet, ngorb)
        #Afull_cici = Afull[ngorb:,ngorb:].reshape (nroots, ndet, nroots, ndet).transpose (0, 2, 1, 3)
        #lib.logger.debug (self, "Orb-orb Hessian:\n{}".format (Afull_orborb))
        #for iroot in range (nroots):
        #    lib.logger.debug (self, "Orb-ci Hessian root {}:\n{}".format (iroot, Afull_orbci[:,iroot,:]))
        #    lib.logger.debug (self, "Ci-orb Hessian root {}:\n{}".format (iroot, Afull_ciorb[iroot,:,:]))
        #    for jroot in range (nroots):
        #        lib.logger.debug (self, "Ci-ci Hessian roots {},{}:\n{}".format (iroot, jroot, Afull_cici[iroot,jroot,:,:]))
        #'''

    def get_lagrange_precond (self, Adiag, level_shift=None, ci=None, **kwargs):
        if level_shift is None: level_shift = self.level_shift
        if ci is None: ci = self.base.ci
        return SACASLagPrec (Adiag=Adiag, level_shift=level_shift, ci=ci, grad_method=self)

    def get_lagrange_callback (self, Lvec_last, itvec, geff_op):
        def my_call (x):
            itvec[0] += 1
            geff = geff_op (x)
            deltax = x - Lvec_last
            gorb, gci = self.unpack_uniq_var (geff)
            deltaorb, deltaci = self.unpack_uniq_var (deltax)
            gci = np.concatenate ([g.ravel () for g in gci])
            deltaci = np.concatenate ([d.ravel () for d in deltaci])
            lib.logger.info (self, ('Lagrange optimization iteration {}, |gorb| = {}, |gci| = {}, '
                '|dLorb| = {}, |dLci| = {}').format (itvec[0], linalg.norm (gorb), linalg.norm (gci),
                linalg.norm (deltaorb), linalg.norm (deltaci)))
            Lvec_last[:] = x[:]
        return my_call

    def project_Aop (self, Aop, ci, state):
        ''' Wrap the Aop function to project out redundant degrees of freedom for the CI part.  What's redundant
            changes between SA-CASSCF and MC-PDFT so modify this part in child classes. '''
        def my_Aop (x):
            Ax = Aop (x)
            Ax_orb, Ax_ci = self.unpack_uniq_var (Ax)
            for i, j in product (range (self.nroots), repeat=2):
                # I'm assuming the only symmetry here that's actually built into the data structure is solver.spin
                # This will be the case as long as the various solvers are determinants with a common total charge
                # occupying a common set of orbitals
                if self.spin_states[i] != self.spin_states[j]: continue
                Ax_ci[i] -= np.dot (Ax_ci[i].ravel (), ci[j].ravel ()) * ci[j]
            #Ax_ci = Ax[self.ngorb:].reshape (self.nroots, -1)
            #ci_arr = np.asarray (ci).reshape (self.nroots, -1)
            #ovlp = np.dot (ci_arr.conjugate (), Ax_ci.T)
            #Ax_ci -= np.dot (ovlp.T, ci_arr)
            #Ax[self.ngorb:] = Ax_ci.ravel ()
            return self.pack_uniq_var (Ax_orb, Ax_ci)
        return my_Aop

    as_scanner = as_scanner

class SACASLagPrec (lagrange.LagPrec):
    ''' A callable preconditioner for solving the Lagrange equations. Based on Mol. Phys. 99, 103 (2001).
    Attributes:

    nroots : integer
        Number of roots in the SA space
    nlag : integer
        Number of Lagrange degrees of freedom
    ngorb : integer
        Number of Lagrange degrees of freedom which are orbital rotations
    level_shift : float
        numerical shift applied to CI rotation Hessian
    ci : ndarray of shape (nroots, ndet or ncscf)
        Ci vectors of the SA space
    Rorb : ndarray of shape (ngorb)
        Diagonal inverse Hessian matrix for orbital rotations
    Rci : ndarray of shape (nroots, ndet or ncsf)
        Diagonal inverse Hessian matrix for CI rotations including a level shift
    Rci_sa : ndarray of shape (nroots (I), ndet or ncsf, nroots (K))
        First two factors of the inverse diagonal CI Hessian projected into SA space:
        Rci(I)|J> <J|Rci(I)|K>^{-1} <K|Rci(I)
        note: right-hand bra and R_I factor not included due to storage considerations
        Make the operand's matrix element with <K|Rci(I) before taking the dot product!
'''

    # TODO: fix me (subclass me? wrap me?) for state_average_mix
    def __init__(self, Adiag=None, level_shift=None, ci=None, grad_method=None):
        self.level_shift = level_shift
        self.nroots = grad_method.nroots
        self.nlag = grad_method.nlag
        self.ngorb = grad_method.ngorb
        self.spin_states = grad_method.spin_states
        self.na_states = grad_method.na_states
        self.nb_states = grad_method.nb_states
        self.grad_method = grad_method
        Aorb, Aci = self.unpack_uniq_var (Adiag)
        self._init_orb (Aorb)
        self._init_ci (Aci, ci)

    def unpack_uniq_var (self, x):
        return self.grad_method.unpack_uniq_var (x)

    def pack_uniq_var (self, xorb, xci):
        return self.grad_method.pack_uniq_var (xorb, xci)

    def _init_orb (self, Aorb):
        self.Rorb = Aorb
        self.Rorb[abs(self.Rorb)<1e-8] = 1e-8
        self.Rorb = 1./self.Rorb

    def _init_ci (self, Aci_spins, ci_spins):
        self.ci = []
        self.Rci = []
        self.Rci_sa = []
        for [Aci, ci] in self._iterate_ci (Aci_spins, ci_spins):
            nroots = Aci.shape[0]
            Rci = Aci + self.level_shift
            Rci[abs(Rci)<1e-8] = 1e-8
            Rci = 1./Rci
            # R_I|J>
            # Indices: I, det, J
            Rci_cross = Rci[:,:,None] * ci.T[None,:,:]
            # S(I)_JK = <J|R_I|K> (first index of CI contract with middle index of R_I|J> and reshape to put I first)
            Sci = np.tensordot (ci.conjugate (), Rci_cross, axes=(1,1)).transpose (1,0,2)
            # R_I|J> S(I)_JK^-1 (can only loop explicitly because of necessary call to linalg.inv)
            # Indices: I, det, K
            Rci_sa = np.zeros_like (Rci_cross)
            for iroot in range (nroots):
                Rci_sa[iroot] = np.dot (Rci_cross[iroot], linalg.inv (Sci[iroot]))
            self.ci.append (ci)
            self.Rci.append (Rci)
            self.Rci_sa.append (Rci_sa)

    def _iterate_ci (self, *args):
        # All args must be iterables over CI vectors in input order
        # Eventually, get rid of copying (np.asarray, etc.)
        # Don't assume args are ndarrays on input
        for my_spin in np.unique (self.spin_states):
            idx = np.where (self.spin_states == my_spin)[0]
            yield [np.asarray ([arg[i] for i in idx]).reshape (len (idx), -1) for arg in args]

    def __call__(self, x):
        xorb, xci = self.unpack_uniq_var (x)
        Mxorb = self.orb_prec (xorb)
        Mxci = self.ci_prec (xci)
        return self.pack_uniq_var (Mxorb, Mxci)

    def orb_prec (self, xorb):
        return self.Rorb * xorb

    def ci_prec (self, xci_spins):
        Mxci = [None,] * self.nroots
        for ix_spin, [xci, desort_spin] in enumerate (self._iterate_ci (xci_spins, list(range(self.nroots)))):
            desort_spin = np.atleast_1d (np.squeeze (desort_spin))
            nroots = xci.shape[0]
            ci = self.ci[ix_spin]
            Rci = self.Rci[ix_spin]
            Rci_sa = self.Rci_sa[ix_spin]
            # R_I|H I> (indices: I, det)
            Rx = Rci * xci
            # <J|R_I|H I> (indices: J, I)
            sa_ovlp = np.dot (ci.conjugate (), Rx.T)
            # R_I|J> S(I)_JK^-1 <K|R_I|H I> (indices: I, det)
            Rx_sub = np.zeros_like (Rx)
            for iroot in range (nroots):
                Rx_sub[iroot] = np.dot (Rci_sa[iroot], sa_ovlp[:,iroot])
            for i, j in enumerate (desort_spin):
                try:
                    Mxci[j] = Rx[i] - Rx_sub[i]
                except Exception as e:
                    print (i, j, desort_spin)
                    raise (e)
        assert (all ([i is not None for i in Mxci]))
        return Mxci

from pyscf import mcscf
mcscf.addons.StateAverageMCSCFSolver.Gradients = lib.class_as_method(Gradients)


#!/usr/bin/env python3
# Copyright 2018 Mitsubishi Electric Research Labs (Takaaki Hori)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
import numpy as np
import paddle
import six

class CTCPrefixScorePD():
    """Batch processing of CTCPrefixScore

    which is based on Algorithm 2 in WATANABE et al.
    "HYBRID CTC/ATTENTION ARCHITECTURE FOR END-TO-END SPEECH RECOGNITION,"
    but extended to efficiently compute the label probablities for multiple
    hypotheses simultaneously
    See also Seki et al. "Vectorized Beam Search for CTC-Attention-Based
    Speech Recognition," In INTERSPEECH (pp. 3825-3829), 2019.
    """

    def __init__(self, x, xlens, blank, eos, margin=0):
        """Construct CTC prefix scorer

        `margin` is M in eq.(22,23)

        :param paddle.Tensor x: input label posterior sequences (B, T, O)
        :param paddle.Tensor xlens: input lengths (B,)
        :param int blank: blank label id
        :param int eos: end-of-sequence id
        :param int margin: margin parameter for windowing (0 means no windowing)
        """
        # In the comment lines,
        # we assume T: input_length, B: batch size, W: beam width, O: output dim.
        self.logzero = -10000000000.0
        self.blank = blank
        self.eos = eos
        self.batch = x.size(0)
        self.input_length = x.size(1)
        self.odim = x.size(2)
        self.dtype = x.dtype

        # Pad the rest of posteriors in the batch
        # TODO(takaaki-hori): need a better way without for-loops
        for i, l in enumerate(xlens):
            if l < self.input_length:
                x[i, l:, :] = self.logzero
                x[i, l:, blank] = 0
        # Reshape input x
        xn = x.transpose([1, 0, 2])  # (B, T, O) -> (T, B, O)
        xb = xn[:, :, self.blank].unsqueeze(2).expand(-1, -1,
                                                      self.odim)  # (T,B,O)
        self.x = paddle.stack([xn, xb])  # (2, T, B, O)
        self.end_frames = paddle.to_tensor(xlens) - 1  # (B,)

        # Setup CTC windowing
        self.margin = margin
        if margin > 0:
            self.frame_ids = paddle.arange(self.input_length, dtype=self.dtype)
        # Base indices for index conversion
        # B idx, hyp idx. shape (B*W, 1)
        self.idx_bh = None
        # B idx. shape (B,)
        self.idx_b = paddle.arange(self.batch)
        # B idx, O idx. shape (B, 1)
        self.idx_bo = (self.idx_b * self.odim).unsqueeze(1)

    def __call__(self, y, state, scoring_ids=None, att_w=None):
        """Compute CTC prefix scores for next labels

        :param list y: prefix label sequences
        :param tuple state: previous CTC state
        :param paddle.Tensor scoring_ids: selected next ids to score (BW, O'), O' <= O
        :param paddle.Tensor att_w: attention weights to decide CTC window
        :return new_state, ctc_local_scores (BW, O)
        """
        output_length = len(y[0]) - 1  # ignore sos
        last_ids = [yi[-1] for yi in y]  # last output label ids
        n_bh = len(last_ids)  # batch * hyps
        n_hyps = n_bh // self.batch  # assuming each utterance has the same # of hyps
        self.scoring_num = scoring_ids.size(
            -1) if scoring_ids is not None else 0
        # prepare state info
        if state is None:
            r_prev = paddle.full(
                (self.input_length, 2, self.batch, n_hyps),
                self.logzero,
                dtype=self.dtype, )  # (T, 2, B, W)
            r_prev[:, 1] = paddle.cumsum(self.x[0, :, :, self.blank],
                                         0).unsqueeze(2)
            r_prev = r_prev.view(-1, 2, n_bh)  # (T, 2, BW)
            s_prev = 0.0  # score
            f_min_prev = 0  # eq. 22-23
            f_max_prev = 1  # eq. 22-23
        else:
            r_prev, s_prev, f_min_prev, f_max_prev = state

        # select input dimensions for scoring
        if self.scoring_num > 0:
            # (BW, O)
            scoring_idmap = paddle.full(
                (n_bh, self.odim), -1, dtype=paddle.long)
            snum = self.scoring_num
            if self.idx_bh is None or n_bh > len(self.idx_bh):
                self.idx_bh = paddle.arange(n_bh).view(-1, 1)  # (BW, 1)
            scoring_idmap[self.idx_bh[:n_bh], scoring_ids] = paddle.arange(snum)
            scoring_idx = (
                scoring_ids + self.idx_bo.repeat(1, n_hyps).view(-1,
                                                                 1)  # (BW,1)
            ).view(-1)  # (BWO)
            # x_ shape (2, T, B*W, O)
            x_ = paddle.index_select(
                self.x.view(2, -1, self.batch * self.odim), scoring_idx,
                2).view(2, -1, n_bh, snum)
        else:
            scoring_ids = None
            scoring_idmap = None
            snum = self.odim
            # x_ shape (2, T, B*W, O)
            x_ = self.x.unsqueeze(3).repeat(1, 1, 1, n_hyps, 1).view(2, -1,
                                                                     n_bh, snum)

        # new CTC forward probs are prepared as a (T x 2 x BW x S) tensor
        # that corresponds to r_t^n(h) and r_t^b(h) in a batch.
        r = paddle.full(
            (self.input_length, 2, n_bh, snum),
            self.logzero,
            dtype=self.dtype, )
        if output_length == 0:
            r[0, 0] = x_[0, 0]

        r_sum = paddle.logsumexp(r_prev, 1)  #(T,BW)
        log_phi = r_sum.unsqueeze(2).repeat(1, 1, snum)  # (T, BW, O)
        if scoring_ids is not None:
            for idx in range(n_bh):
                pos = scoring_idmap[idx, last_ids[idx]]
                if pos >= 0:
                    log_phi[:, idx, pos] = r_prev[:, 1, idx]
        else:
            for idx in range(n_bh):
                log_phi[:, idx, last_ids[idx]] = r_prev[:, 1, idx]

        # decide start and end frames based on attention weights
        if att_w is not None and self.margin > 0:
            f_arg = paddle.matmul(att_w, self.frame_ids)
            f_min = max(int(f_arg.min().cpu()), f_min_prev)
            f_max = max(int(f_arg.max().cpu()), f_max_prev)
            start = min(f_max_prev, max(f_min - self.margin, output_length, 1))
            end = min(f_max + self.margin, self.input_length)
        else:
            f_min = f_max = 0
            # if one frame one out, the output_length is the eating frame num now.
            start = max(output_length, 1)
            end = self.input_length

        # compute forward probabilities log(r_t^n(h)) and log(r_t^b(h))
        for t in range(start, end):
            rp = r[t - 1]  # (2 x BW x O')
            rr = paddle.stack([rp[0], log_phi[t - 1], rp[0], rp[1]]).view(
                2, 2, n_bh, snum)  # (2,2,BW,O')
            r[t] = paddle.logsumexp(rr, 1) + x_[:, t]

        # compute log prefix probabilities log(psi)
        log_phi_x = paddle.concat(
            (log_phi[0].unsqueeze(0), log_phi[:-1]), axis=0) + x_[0]
        if scoring_ids is not None:
            log_psi = paddle.full(
                (n_bh, self.odim), self.logzero, dtype=self.dtype)
            log_psi_ = paddle.logsumexp(
                paddle.concat(
                    (log_phi_x[start:end], r[start - 1, 0].unsqueeze(0)),
                    axis=0),
                axis=0, )
            for si in range(n_bh):
                log_psi[si, scoring_ids[si]] = log_psi_[si]
        else:
            log_psi = paddle.logsumexp(
                paddle.concat(
                    (log_phi_x[start:end], r[start - 1, 0].unsqueeze(0)),
                    axis=0),
                axis=0, )

        for si in range(n_bh):
            log_psi[si, self.eos] = r_sum[self.end_frames[si // n_hyps], si]

        # exclude blank probs
        log_psi[:, self.blank] = self.logzero

        return (log_psi - s_prev), (r, log_psi, f_min, f_max, scoring_idmap)

    def index_select_state(self, state, best_ids):
        """Select CTC states according to best ids

        :param state    : CTC state
        :param best_ids : index numbers selected by beam pruning (B, W)
        :return selected_state
        """
        r, s, f_min, f_max, scoring_idmap = state
        # convert ids to BHO space
        n_bh = len(s)
        n_hyps = n_bh // self.batch
        vidx = (best_ids + (self.idx_b *
                            (n_hyps * self.odim)).view(-1, 1)).view(-1)
        # select hypothesis scores
        s_new = paddle.index_select(s.view(-1), vidx, 0)
        s_new = s_new.view(-1, 1).repeat(1, self.odim).view(n_bh, self.odim)
        # convert ids to BHS space (S: scoring_num)
        if scoring_idmap is not None:
            snum = self.scoring_num
            hyp_idx = (best_ids // self.odim +
                       (self.idx_b * n_hyps).view(-1, 1)).view(-1)
            label_ids = paddle.fmod(best_ids, self.odim).view(-1)
            score_idx = scoring_idmap[hyp_idx, label_ids]
            score_idx[score_idx == -1] = 0
            vidx = score_idx + hyp_idx * snum
        else:
            snum = self.odim
        # select forward probabilities
        r_new = paddle.index_select(r.view(-1, 2, n_bh * snum), vidx, 2).view(
            -1, 2, n_bh)
        return r_new, s_new, f_min, f_max

    def extend_prob(self, x):
        """Extend CTC prob.

        :param paddle.Tensor x: input label posterior sequences (B, T, O)
        """

        if self.x.shape[1] < x.shape[1]:  # self.x (2,T,B,O); x (B,T,O)
            # Pad the rest of posteriors in the batch
            # TODO(takaaki-hori): need a better way without for-loops
            xlens = [x.size(1)]
            for i, l in enumerate(xlens):
                if l < self.input_length:
                    x[i, l:, :] = self.logzero
                    x[i, l:, self.blank] = 0
            tmp_x = self.x
            xn = x.transpose([1, 0, 2])  # (B, T, O) -> (T, B, O)
            xb = xn[:, :, self.blank].unsqueeze(2).expand(-1, -1, self.odim)
            self.x = paddle.stack([xn, xb])  # (2, T, B, O)
            self.x[:, :tmp_x.shape[1], :, :] = tmp_x
            self.input_length = x.size(1)
            self.end_frames = paddle.to_tensor(xlens) - 1

    def extend_state(self, state):
        """Compute CTC prefix state.

        :param state    : CTC state
        :return ctc_state
        """

        if state is None:
            # nothing to do
            return state
        else:
            r_prev, s_prev, f_min_prev, f_max_prev = state

            r_prev_new = paddle.full(
                (self.input_length, 2),
                self.logzero,
                dtype=self.dtype, )
            start = max(r_prev.shape[0], 1)
            r_prev_new[0:start] = r_prev
            for t in range(start, self.input_length):
                r_prev_new[t, 1] = r_prev_new[t - 1, 1] + self.x[0, t, :,
                                                                 self.blank]

            return (r_prev_new, s_prev, f_min_prev, f_max_prev)

class CTCPrefixScore():
    """Compute CTC label sequence scores

    which is based on Algorithm 2 in WATANABE et al.
    "HYBRID CTC/ATTENTION ARCHITECTURE FOR END-TO-END SPEECH RECOGNITION,"
    but extended to efficiently compute the probablities of multiple labels
    simultaneously
    """

    def __init__(self, x, blank, eos, xp):
        self.xp = xp
        self.logzero = -10000000000.0
        self.blank = blank
        self.eos = eos
        self.input_length = len(x)
        self.x = x  # (T, O)

    def initial_state(self):
        """Obtain an initial CTC state

        :return: CTC state
        """
        # initial CTC state is made of a frame x 2 tensor that corresponds to
        # r_t^n(<sos>) and r_t^b(<sos>), where 0 and 1 of axis=1 represent
        # superscripts n and b (non-blank and blank), respectively.
        # r shape (T, 2)
        r = self.xp.full((self.input_length, 2), self.logzero, dtype=np.float32)
        r[0, 1] = self.x[0, self.blank]
        for i in six.moves.range(1, self.input_length):
            r[i, 1] = r[i - 1, 1] + self.x[i, self.blank]
        return r

    def __call__(self, y, cs, r_prev):
        """Compute CTC prefix scores for next labels

        :param y     : prefix label sequence
        :param cs    : array of next labels
        :param r_prev: previous CTC state
        :return ctc_scores, ctc_states
        """
        # initialize CTC states
        output_length = len(y) - 1  # ignore sos
        # new CTC states are prepared as a frame x (n or b) x n_labels tensor
        # that corresponds to r_t^n(h) and r_t^b(h).
        # r shape (T, 2, n_labels)
        r = self.xp.ndarray((self.input_length, 2, len(cs)), dtype=np.float32)
        xs = self.x[:, cs]
        if output_length == 0:
            r[0, 0] = xs[0]
            r[0, 1] = self.logzero
        else:
            # Although the code does not exactly follow Algorithm 2,
            # we don't have to change it because we can assume
            # r_t(h)=0 for t < |h| in CTC forward computation
            # (Note: we assume here that index t starts with 0).
            # The purpose of this difference is to reduce the number of for-loops.
            # https://github.com/espnet/espnet/pull/3655
            # where we start to accumulate r_t(h) from t=|h|
            # and iterate r_t(h) = (r_{t-1}(h) + ...) to T-1,
            # avoiding accumulating zeros for t=1~|h|-1.
            # Thus, we need to set r_{|h|-1}(h) = 0,
            # i.e., r[output_length-1] = logzero, for initialization.
            # This is just for reducing the computation.
            r[output_length - 1] = self.logzero

        # prepare forward probabilities for the last label
        r_sum = self.xp.logaddexp(r_prev[:, 0],
                                  r_prev[:, 1])  # log(r_t^n(g) + r_t^b(g))
        last = y[-1]
        if output_length > 0 and last in cs:
            log_phi = self.xp.ndarray(
                (self.input_length, len(cs)), dtype=np.float32)
            for i in six.moves.range(len(cs)):
                log_phi[:, i] = r_sum if cs[i] != last else r_prev[:, 1]
        else:
            log_phi = r_sum

        # compute forward probabilities log(r_t^n(h)), log(r_t^b(h)),
        # and log prefix probabilities log(psi)
        start = max(output_length, 1)
        log_psi = r[start - 1, 0]
        for t in six.moves.range(start, self.input_length):
            r[t, 0] = self.xp.logaddexp(r[t - 1, 0], log_phi[t - 1]) + xs[t]
            r[t, 1] = (self.xp.logaddexp(r[t - 1, 0], r[t - 1, 1]) +
                       self.x[t, self.blank])
            log_psi = self.xp.logaddexp(log_psi, log_phi[t - 1] + xs[t])

        # get P(...eos|X) that ends with the prefix itself
        eos_pos = self.xp.where(cs == self.eos)[0]
        if len(eos_pos) > 0:
            log_psi[eos_pos] = r_sum[-1]  # log(r_T^n(g) + r_T^b(g))

        # exclude blank probs
        blank_pos = self.xp.where(cs == self.blank)[0]
        if len(blank_pos) > 0:
            log_psi[blank_pos] = self.logzero

        # return the log prefix probability and CTC states, where the label axis
        # of the CTC states is moved to the first axis to slice it easily
        # log_psi shape (n_labels,), state shape (n_labels, T, 2)
        return log_psi, self.xp.rollaxis(r, 2)


#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   WTM.py
@Time    :   2020/10/06 17:13:43
@Author  :   Leilan Zhang
@Version :   1.0
@Contact :   zhangleilan@gmail.com
@Desc    :   None
'''

import os
import re
import time
import pickle
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
from .wae import WAE
import sys
sys.path.append('..')
from utils import evaluate_topic_quality, smooth_curve

class WTM:
    def __init__(self, bow_dim=10000, n_topic=20, device=None, dist='gmm_std', taskname=None, dropout=0.0):
        self.bow_dim = bow_dim
        self.n_topic = n_topic
        self.wae = WAE(encode_dims=[bow_dim, 1024, 512, n_topic], decode_dims=[n_topic, 512, bow_dim], dropout=dropout, nonlin='relu')
        self.device = device
        self.id2token = None
        self.dist = dist
        self.taskname = taskname
        if device != None:
            self.wae = self.wae.to(device)

    def train(self, train_data, batch_size=256, learning_rate=1e-3, test_data=None, num_epochs=100, is_evaluate=False, log_every=5, beta=1.0):
        self.wae.train()
        self.id2token = {v: k for k,v in train_data.dictionary.token2id.items()}
        data_loader = DataLoader(train_data, batch_size=batch_size,shuffle=True, num_workers=4, collate_fn=train_data.collate_fn)

        optimizer = torch.optim.Adam(self.wae.parameters(), lr=learning_rate)
        #scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)
        trainloss_lst, valloss_lst = [], []
        c_v_lst, c_w2v_lst, c_uci_lst, c_npmi_lst, mimno_tc_lst, td_lst = [], [], [], [], [], []
        for epoch in range(num_epochs):
            epochloss_lst = []
            for iter, data in enumerate(data_loader):
                optimizer.zero_grad()

                txts, bows = data
                bows = bows.to(self.device)

                bows_recon, theta_q = self.wae(bows)

                theta_prior = self.wae.sample(dist=self.dist, batch_size=len(bows), ori_data=bows).to(self.device)

                logsoftmax = torch.log_softmax(bows_recon, dim=1)
                rec_loss = -1.0 * torch.sum(bows*logsoftmax)
                #rec_loss = F.binary_cross_entropy(torch.softmax(bows_recon,dim=1),bows,reduction='sum')
                #rec_loss = F.binary_cross_entropy(bows_recon,bows,reduction='sum')
                mmd = self.wae.mmd_loss(theta_q, theta_prior, device=self.device, t=0.1)
                #mmd = self.wae.mmd_loss(hid_vecs, theta_prior, device=self.device, t=0.1)
                s = torch.sum(bows)/len(bows)
                lamb = (5.0*s*torch.log(torch.tensor(1.0 *bows.shape[-1]))/torch.log(torch.tensor(2.0)))
                mmd = mmd * lamb

                loss = rec_loss + mmd * beta

                loss.backward()
                optimizer.step()

                trainloss_lst.append(loss.item()/len(bows))
                epochloss_lst.append(loss.item()/len(bows))
                if (iter+1) % 10 == 0:
                    print(f'Epoch {(epoch+1):>3d}\tIter {(iter+1):>4d}\tLoss:{loss.item()/len(bows):<.7f}\tRec Loss:{rec_loss.item()/len(bows):<.7f}\tMMD:{mmd.item()/len(bows):<.7f}')
            #scheduler.step()
            if (epoch+1) % log_every == 0:
                print(f'Epoch {(epoch+1):>3d}\tLoss:{sum(epochloss_lst)/len(epochloss_lst):<.7f}')
                print('\n'.join([str(lst) for lst in self.show_topic_words()]))
                print('='*30)
                smth_pts = smooth_curve(trainloss_lst)
                plt.plot(np.array(range(len(smth_pts)))*log_every, smth_pts)
                plt.xlabel('epochs')
                plt.title('Train Loss')
                plt.savefig('wlda_trainloss.png')
                if test_data!=None:
                    c_v,c_w2v,c_uci,c_npmi,mimno_tc, td = self.evaluate(test_data,calc4each=False)
                    c_v_lst.append(c_v), c_w2v_lst.append(c_w2v), c_uci_lst.append(c_uci),c_npmi_lst.append(c_npmi), mimno_tc_lst.append(mimno_tc), td_lst.append(td)
                save_name = f'./ckpt/WTM_{self.taskname}_tp{self.n_topic}_{self.dist}_{time.strftime("%Y-%m-%d-%H-%M", time.localtime())}.ckpt'
                torch.save(self.wae.state_dict(),save_name)
        scrs = {'c_v':c_v_lst,'c_w2v':c_w2v_lst,'c_uci':c_uci_lst,'c_npmi':c_npmi_lst,'mimno_tc':mimno_tc_lst,'td':td_lst}
        '''
        for scr_name,scr_lst in scrs.items():
            plt.cla()
            plt.plot(np.array(range(len(scr_lst)))*log_every,scr_lst)
            plt.savefig(f'wlda_{scr_name}.png')
        '''
        plt.cla()
        for scr_name,scr_lst in scrs.items():
            if scr_name in ['c_v','c_w2v','td']:
                plt.plot(np.array(range(len(scr_lst)))*log_every,scr_lst,label=scr_name)
        plt.title('Topic Coherence')
        plt.xlabel('epochs')
        plt.legend()
        plt.savefig(f'wlda_tc_scores.png')

    def evaluate(self, test_data, calc4each=False):
        topic_words = self.show_topic_words()
        return evaluate_topic_quality(topic_words, test_data, taskname=self.taskname, calc4each=calc4each)

    def inference(self, doc_bow):
        # doc_bow: torch.tensor [vocab_size]; optional: np.array [vocab_size]
        if isinstance(doc_bow,np.array):
            doc_bow = torch.from_numpy(doc_bow)
        doc_bow = doc_bow.reshape(1,self.bow_dim).to(self.device)
        with torch.no_grad():
            theta = F.softmax(self.wae.encode(doc_bow),dim=1)
            return theta.detach().cpu().squeeze(0).numpy()

    def inference(self, doc_tokenized, dictionary,normalize=True):
        doc_bow = torch.zeros(1,self.bow_dim)
        for token in doc_tokenized:
            try:
                idx = dictionary.token2id[token]
                doc_bow[0][idx] = 1.0
            except:
                print(f'{token} not in the vocabulary.')
        doc_bow = doc_bow.to(self.device)
        with torch.no_grad():
            theta = self.wae.encode(doc_bow)
            if normalize:
                theta = F.softmax(theta,dim=1)
            return theta.detach().cpu().squeeze(0).numpy()

    def get_embed(self,train_data, num=1000):
        self.wae.eval()
        data_loader = DataLoader(train_data, batch_size=512,shuffle=False, num_workers=4, collate_fn=train_data.collate_fn)
        embed_lst = []
        txt_lst = []
        for data_batch in data_loader:
            txts, bows = data_batch
            embed = self.inference(bows,train_data.dictionary)
            embed_lst.append(embed)
            txt_lst.append(txts)
            cnt += embed.shape[0]
            if cnt>=num:
                break
        embed_lst = torch.concat(embed_lst,dim=0)[:num]
        txt_lst = torch.concat(txt_lst,dim=0)[:num]
        return txt_lst, embed_lst

    def get_topic_word_dist(self,normalize=True):
        self.wae.eval()
        with torch.no_grad():
            idxes = torch.eye(self.n_topic).to(self.device)
            word_dist = self.wae.decode(idxes)  # word_dist: [n_topic, vocab.size]
            if normalize:
                word_dist = F.softmax(word_dist,dim=1)
            return word_dist.detach().cpu().numpy()

    def show_topic_words(self, topic_id=None, topK=15):
        self.wae.eval()
        topic_words = []
        idxes = torch.eye(self.n_topic).to(self.device)
        word_dist = self.wae.decode(idxes)
        word_dist = F.softmax(word_dist, dim=1)
        vals, indices = torch.topk(word_dist, topK, dim=1)
        vals = vals.cpu().tolist()
        indices = indices.cpu().tolist()
        if topic_id == None:
            for i in range(self.n_topic):
                topic_words.append([self.id2token[idx] for idx in indices[i]])
        else:
            topic_words.append([self.id2token[idx] for idx in indices[topic_id]])
        return topic_words

if __name__ == '__main__':
    model = WAE(encode_dims=[1024, 512, 256, 20],
                decode_dims=[20, 128, 768, 1024])
    model = model.cuda()
    inpt = torch.randn(234, 1024).cuda()
    out, mu, log_var = model(inpt)
    print(out.shape)
    print(mu.shape)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

################################################################################
#
#   CanTherm -
#
#   Copyright (c) 2010 by Joshua W. Allen (jwallen@mit.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

import os.path
import logging
import math
import numpy
import matplotlib
matplotlib.rc('mathtext', default='regular')

from rmgpy.quantity import constants
from rmgpy.statmech import HinderedRotor, HarmonicOscillator
from rmgpy.species import Species, TransitionState
from rmgpy.kinetics import Arrhenius
from rmgpy.reaction import Reaction

from gaussian import GaussianLog
from states import projectRotors, applyEnergyCorrections

################################################################################

# The model chemistry used
# The energies for each species and transition state will be automatically
# adjusted using standard reference energies for that method
modelChemistry = ''

# A dictionary associating species identifiers with species objects
speciesDict = {}
# A dictionary associating transition state identifiers with transition state objects
transitionStateDict = {}
# A dictionary associating reaction identifiers with reaction objects
reactionDict = {}
# A dictionary associated species and transition state identifiers with geometry objects
geometryDict = {}

# The file to save the output to
outputFile = ''

################################################################################

def setOutputFile(path):
    global outputFile
    outputFile = path
    f = open(path, 'w')
    f.close()

def setModelChemistry(method):
    """
    Set the model chemistry used in this quantum chemisty calculation to
    `method`.
    """
    global modelChemistry
    modelChemistry = method

################################################################################

def hinderedRotor(scanLog, pivots, top, symmetry):
    pivots = [p-1 for p in pivots]
    top = [t-1 for t in top]
    return [scanLog, pivots, top, symmetry]

################################################################################

def loadConfiguration(energyLog, geomLog, statesLog, extSymmetry, spinMultiplicity, freqScaleFactor, linear, rotors, atoms, bonds, E0=None, TS=False):

    logging.debug('    Reading optimized geometry...')
    log = GaussianLog(geomLog)
    geom = log.loadGeometry()

    logging.debug('    Reading energy...')
    if E0 is None:
        if energyLog is not None: log = GaussianLog(energyLog)
        E0 = log.loadEnergy()
    else:
        E0 *= 4.35974394e-18 * constants.Na     # Hartree/particle to J/mol
    E0 = applyEnergyCorrections(E0, modelChemistry, atoms, bonds)
    logging.debug('         E0 (0 K) = %g kcal/mol' % (E0 / 4184))

    logging.debug('    Reading molecular degrees of freedom...')
    log = GaussianLog(statesLog)
    states = log.loadStates(symmetry=extSymmetry)
    states.spinMultiplicity = spinMultiplicity

    F = log.loadForceConstantMatrix()

    if F is not None and len(geom.mass) > 1 and len(rotors) > 0:

        logging.debug('    Fitting %i hindered rotors...' % len(rotors))
        for scanLog, pivots, top, symmetry in rotors:
            log = GaussianLog(scanLog)

            Vlist, angle = log.loadScanEnergies()

            inertia = geom.getInternalReducedMomentOfInertia(pivots, top)

            barr, symm = log.fitCosinePotential()
            cosineRotor = HinderedRotor(inertia=(inertia*constants.Na*1e23,"amu*angstrom^2"), symmetry=symm, barrier=(barr/4184.,"kcal/mol"))
            fourier = log.fitFourierSeriesPotential()
            fourierRotor = HinderedRotor(inertia=(inertia*constants.Na*1e23,"amu*angstrom^2"), symmetry=symmetry, fourier=(fourier,"J/mol"))

            Vlist_cosine = cosineRotor.getPotential(angle)
            Vlist_fourier = fourierRotor.getPotential(angle)

            rms_cosine = numpy.sqrt(numpy.sum((Vlist_cosine - Vlist) * (Vlist_cosine - Vlist)) / (len(Vlist) - 1)) / 4184.
            rms_fourier = numpy.sqrt(numpy.sum((Vlist_fourier - Vlist) * (Vlist_fourier - Vlist))/ (len(Vlist) - 1)) / 4184.
            print rms_cosine, rms_fourier, symm, symmetry

            # Keep the rotor with the most accurate potential
            rotor = cosineRotor if rms_cosine < rms_fourier else fourierRotor
            # However, keep the cosine rotor if it is accurate enough, the
            # fourier rotor is not significantly more accurate, and the cosine
            # rotor has the correct symmetry
            if rms_cosine < 0.05 and rms_cosine / rms_fourier > 0.25 and rms_cosine / rms_fourier < 4.0 and symmetry == symm:
                rotor = cosineRotor

            states.modes.append(rotor)

            import pylab
            phi = numpy.arange(0, 6.3, 0.02, numpy.float64)
            fig = pylab.figure()
            pylab.plot(angle, Vlist / 4184, 'ok')
            linespec = '-r' if rotor is cosineRotor else '--r'
            pylab.plot(phi, cosineRotor.getPotential(phi) / 4184, linespec)
            linespec = '-b' if rotor is fourierRotor else '--b'
            pylab.plot(phi, fourierRotor.getPotential(phi) / 4184, linespec)
            pylab.legend(['scan', 'cosine', 'fourier'], loc=1)
            pylab.xlim(0, 2*math.pi)

            axes = fig.get_axes()[0]
            axes.set_xticks([float(j*math.pi/4) for j in range(0,9)])
            axes.set_xticks([float(j*math.pi/8) for j in range(0,17)], minor=True)
            axes.set_xticklabels(['$0$', '$\pi/4$', '$\pi/2$', '$3\pi/4$', '$\pi$', '$5\pi/4$', '$3\pi/2$', '$7\pi/4$', '$2\pi$'])

        pylab.show()

        logging.debug('    Determining frequencies from reduced force constant matrix...')
        frequencies = list(projectRotors(geom, F, rotors, linear, TS))

    elif len(states.modes) > 2:
        frequencies = states.modes[2].frequencies.values
        rotors = []
    else:
        frequencies = []
        rotors = []

    for mode in states.modes:
        if isinstance(mode, HarmonicOscillator):
            mode.frequencies.values = numpy.array(frequencies, numpy.float) * freqScaleFactor

    return E0, geom, states

def loadSpecies(label, geomLog, statesLog, extSymmetry, spinMultiplicity, freqScaleFactor, linear, rotors, atoms, bonds, directory=None, E0=None, energyLog=None):
    global modelChemistry
    logging.info('Loading species %s...' % label)
    if directory:
        geomLog = os.path.join(directory, geomLog)
        statesLog = os.path.join(directory, statesLog)
        if energyLog: energyLog = os.path.join(directory, energyLog)
        for rotor in rotors:
            rotor[0] = os.path.join(directory, rotor[0])
    E0, geom, states = loadConfiguration(energyLog, geomLog, statesLog, extSymmetry, spinMultiplicity, freqScaleFactor, linear, rotors, atoms, bonds, E0, TS=False)
    speciesDict[label] = Species(label=label, thermo=None, states=states, E0=(E0/1000.,"kJ/mol"))
    geometryDict[label] = geom

def loadTransitionState(label, geomLog, statesLog, extSymmetry, spinMultiplicity, freqScaleFactor, linear, rotors, atoms, bonds, directory=None, E0=None, energyLog=None):
    global modelChemistry
    logging.info('Loading transition state %s...' % label)
    if directory:
        geomLog = os.path.join(directory, geomLog)
        statesLog = os.path.join(directory, statesLog)
        if energyLog: energyLog = os.path.join(directory, energyLog)
        for rotor in rotors:
            rotor[0] = os.path.join(directory, rotor[0])
    E0, geom, states = loadConfiguration(energyLog, geomLog, statesLog, extSymmetry, spinMultiplicity, freqScaleFactor, linear, rotors, atoms, bonds, E0, TS=True)
    log = GaussianLog(statesLog)
    frequency = log.loadNegativeFrequency() * freqScaleFactor
    transitionStateDict[label] = TransitionState(label=label, states=states, frequency=(frequency,"cm^-1"), E0=(E0/1000.,"kJ/mol"))
    geometryDict[label] = geom

################################################################################

def loadReaction(label, reactants, products, transitionState, degeneracy=1):
    global speciesDict, transitionStateDict, reactionDict
    logging.info('Loading reaction %s...' % label)
    rxn = Reaction(
        reactants=[speciesDict[s] for s in reactants],
        products=[speciesDict[s] for s in products],
        transitionState=transitionStateDict[transitionState],
    )
    rxn.degeneracy = degeneracy
    reactionDict[label] = rxn

################################################################################

def generateStates(label):
    global outputFile, speciesDict, transitionStateDict
    from states import saveStates
    if label in speciesDict:
        saveStates(speciesDict[label], geometryDict[label], label, outputFile)
    elif label in transitionStateDict:
        saveStates(transitionStateDict[label], geometryDict[label], label, outputFile)

def generateThermo(label, model, plot=False):
    global outputFile, speciesDict
    from thermo import generateThermoModel, saveThermo
    generateThermoModel(speciesDict[label], model, plot)
    saveThermo(speciesDict[label], label, outputFile)

def generateKinetics(label, tunneling='', plot=False):
    global outputFile, reactionDict
    from kinetics import generateKineticsModel, saveKinetics
    generateKineticsModel(reactionDict[label], tunneling, plot)
    saveKinetics(reactionDict[label], tunneling, label, outputFile)