active_space_uhf

import os

from pyscf import fci, ao2mo, mcscf
from pyscf.fci.cistring import make_strings
from pyscf import gto, scf
from pyscf.fci.cistring import make_strings
import pyscf
from pyscf.cc.addons import spatial2spin
from functools import reduce

import numpy as np

import openfermion as of


def get_molecule():
    from openfermionpyscf import run_pyscf

    dim = 2.55
    geometry = [['O', [0, 0, 0]], ['O', [0, dim, 0]]]
    basis = 'sto-3g'
    charge = 0
    multiplicity = 3
    molecule = of.MolecularData(geometry=geometry,
                                charge=charge,
                                multiplicity=multiplicity,
                                basis=basis)
    molecule = run_pyscf(molecule, run_scf=True)# , run_fci=True)
    return molecule


def active_space_uhf(mol, mf):
    cas = mcscf.CASCI(mf, ncas=8, nelecas=sum(mol.nelec)-4)
    h1, ecore = cas.get_h1eff()
    eri = cas.get_h2cas()
    h1 = np.ascontiguousarray(h1)
    eri = np.ascontiguousarray(eri)
    eri = ao2mo.restore('s8', eri, 8)
    ecore = np.float(ecore)

    active_mol = gto.M()
    active_mol.nelectron = sum(cas.nelecas)

    mf_uhf = scf.UHF(active_mol)
    mf_uhf.nelec = cas.nelecas
    mf_uhf.get_hcore = lambda *args: np.asarray(h1)
    mf_uhf.get_ovlp = lambda *args: np.eye(8)
    mf_uhf.energy_nuc = lambda *args: ecore
    mf_uhf._eri = eri # ao2mo.restore('8', np.zeros((8, 8, 8, 8)), 8)
    mf_uhf.init_guess = '1e'
    mf_uhf.kernel()


def main():
    molecule = get_molecule()
    molecule.filename = os.path.join(os.getcwd(), molecule.name)
    molecule.save()

    mf = molecule._pyscf_data['scf']
    mol = molecule._pyscf_data['mol']

    # do joonho's trick of defining the ROHF core
    active_space_uhf(mol, mf)


if __name__ == "__main__":
    main()