"""
A PySCF-based spin-unrestricted sigma-SCF-SVO
"""


import h5py
import numpy as np

from pyscf import scf, ao2mo

from frankenstein import scf as fscf
from frankenstein.tools.scf_utils import (get_uscf_fock, get_uscf_energy,
    get_uscf_S2)
from frankenstein.tools.sscf_utils import get_usscf_fock_faster
from frankenstein.tools.tensor_utils import get_symm_mat_pow


def get_Vact(mol, mo_act, Vact_file, Vact_name):
    nmo = mo_act.shape[1]
    if h5py.is_hdf5(Vact_file):
        with h5py.File(Vact_file, "r") as fVact:
            if fVact[Vact_name].size == nmo**4:
                return fVact[Vact_name][()].reshape(*([nmo]*4))

    ao2mo.outcore.full(mol, mo_act, Vact_file, dataname=Vact_name,
        compact=False)
    with h5py.File(Vact_file, "r") as fVact:
        Vact = fVact[Vact_name][()].reshape(*([nmo]*4))

    return Vact


class pyUSSCFSVO(fscf.USSCF):

    def __init__(self, pymf, ncore, nact, **kwargs):
        """
        """
        if not isinstance(pymf, scf.hf.RHF):
            raise ValueError("Arg1 must be pyscf's RHF instance.")
        elif not pymf.converged:
            raise ValueError("Input RHF instance is not converged.")

        self.pymf = pymf
        self.h0 = pymf.get_hcore()

        self.mo_core = self.pymf.mo_coeff[:,:ncore] if ncore > 0 else None
        self.mo_act = self.pymf.mo_coeff[:,ncore:ncore+nact]
        X = get_symm_mat_pow(self.pymf.get_ovlp(), 0.5)
        self.mo_act_soao = X @ self.mo_act

        if self.mo_core is None:
            self.Ebath = 0.
            self.Gcore = 0.
        else:
            Pcore = self.mo_core@self.mo_core.T
            self.Gcore0 = pymf.get_fock(dm=Pcore*2.) - self.h0
            self.Ebath = np.einsum("ij,ji->", Pcore, 2*self.h0+self.Gcore0)
            self.Gcore = self.mo_act.T @ self.Gcore0 @ self.mo_act

        # the following attributes can be set via initialization
        self.conv = 7
        self.opt = "diis"
        self.max_iter = 100
        self.guess = "core"
        self.max_diis = 20
        self.dm_method = "bfgs"
        self.verbose = self.pymf.mol.verbose
        self.no_roothaan = False
        self.mom = False

        self.guess_mix = 0

        self.omega = None
        self.Vact_file = "__eris.h5py"
        self.Vact_name = "eri"

        self.__dict__.update(kwargs)

        # do not touch the following argument unless you know what you are doing
        self.unrestricted = True

        self.initialization()

    # properties
    @property
    def ncore(self):
        return 0 if self.mo_core is None else self.mo_core.shape[1]

    @property
    def nmo(self):
        return self.mo_act.shape[1]

    @property
    def noccs(self):
        nocc = self.pymf.mol.nelectron//2
        return [nocc-self.ncore] * 2

    @property
    def nvirs(self):
        return [self.nmo-self.noccs[i] for i in [0,1]]

    @property
    def e_nuc(self):
        return self.pymf.energy_nuc()

    @property
    def e_tot(self):
        return self.e_scf + self.e_nuc

    def initialization(self):
        self.h = self.mo_act.T @ self.h0 @ self.mo_act + self.Gcore
        self.V = get_Vact(self.pymf.mol, self.mo_act, self.Vact_file,
            self.Vact_name)

        def __rdm1_builder_ez(mo_coeff):
            noccs = self.noccs
            mo_occ = [mo_coeff[i][:,:noccs[i]] for i in [0,1]]
            return [mo_occ[i]@mo_occ[i].T for i in [0,1]]

        def __roothaan_step_ez(fock):
            es = [None]*2
            us = [None]*2
            for i in [0,1]:
                es[i], us[i] = np.linalg.eigh(fock[i])

            return es, us

        def __uhf_fock_builder_ez(rdm1):
            return get_uscf_fock(self.h, self.V, rdm1)

        def __fock_builder_ez(Cs, omega, uhf_fock, e_uscf, ret_var=False):
            return get_usscf_fock_faster(self.h, self.V, Cs, self.noccs,
                omega, uhf_fock, e_uscf, ret_var)

        def __e_scf_builder_ez(uhf_fock, rdm1):
            return get_uscf_energy(self.h, uhf_fock, rdm1) + self.Ebath

        def __S2_builder_ez(rdm1):
            rdm10 = [self.mo_act_soao @ rdm1[i] @ self.mo_act_soao.T
                for i in [0,1]]
            return get_uscf_S2(rdm10, self.noccs)

        self.rdm1_builder_ez = __rdm1_builder_ez
        self.roothaan_step_ez = __roothaan_step_ez
        self.uhf_fock_builder_ez = __uhf_fock_builder_ez
        self.fock_builder_ez = __fock_builder_ez
        self.e_scf_builder_ez = __e_scf_builder_ez
        self.S2_builder_ez = __S2_builder_ez

    def update(self, mo_coeff=None):
        if mo_coeff is None: mo_coeff = self.mo_coeff
        self.rdm1 = self.rdm1_builder_ez(mo_coeff)
        self.uhf_fock = self.uhf_fock_builder_ez(self.rdm1)
        self.e_scf = self.e_scf_builder_ez(self.uhf_fock, self.rdm1)
        self.fock, self.var = self.fock_builder_ez(mo_coeff, self.omega,
            self.uhf_fock, self.e_scf, ret_var=True)
        self.S2 = self.S2_builder_ez(self.rdm1)

    def get_init_guess(self, C0=None, D0=None):
        if C0 is None:
            self.mo_coeff = [np.eye(self.nmo) for i in [0,1]]
        else:
            self.mo_coeff = C0

        if self.guess_mix > 0:
            from frankenstein.tools.scf_utils import homo_lumo_mix
            homo_lumo_mix(self.mo_coeff[1], self.noccs[1], 0.1*self.guess_mix)


if __name__ == "__main__":
    import os
    from pyscf import gto, scf
    from frankenstein.tools.pyscf_utils import get_pymol

    geom = "O; H 1 0.9614; H 1 0.9614 2 104.4"
    basis = "cc-pVDZ"

    pymol = get_pymol(geom, basis, verbose=0)
    pymf = scf.RHF(pymol)
    pymf.kernel()
    print(pymf.e_tot)
    print(pymf.e_tot - pymf.energy_nuc())

    # pyusvo = pyUSSCFSVO(pymf, 1, 8, verbose=1, omega=-87, opt="gdm")
    # pyusvo.kernel()

    de_list = []
    e0 = pymf.e_tot - pymf.energy_nuc()
    # for dw in np.arange(8, 16, 0.05):
    # for dw in np.arange(16, 18, 0.05):
    for dw in [16]:
    # for dw in np.arange(0,8,0.1):
        omega = e0 + dw/27.211399
        pyusvo = pyUSSCFSVO(pymf, 1, 6, verbose=1, omega=omega, opt="gdm")
        pyusvo.guess_mix = 3
        pyusvo.kernel()

        pyusvo.omega = None
        pyusvo.guess_mix = 0
        pyusvo.kernel(C0=pyusvo.mo_coeff)

        de = (pyusvo.e_scf - e0) * 27.211399
        de_str = "{:.4f}".format(de)
        if pyusvo.is_converged and not de_str in de_list:
            de_list.append(de_str)
            print("{: .2f}  {: .10f}  {: .10f}  {: .8f}".format(dw, de,
                pyusvo.var, pyusvo.S2))

    os.system("rm -f {:s}".format(pyusvo.Vact_file))
