r"""Prepare a multiconfiguration circuit from the results of PySCF's large_ci method."""

import pyscf
from pyscf import mcscf
from inquanto.extensions.pyscf import (
    ChemistryDriverPySCFMolecularRHF,
    FromActiveSpace,
    pyscf_to_inquanto_parity_update,
)

from inquanto.states import FermionState, FermionStateString
from inquanto.mappings import QubitMappingJordanWigner
from inquanto.ansatzes import MultiConfigurationStateSparse

from inquanto.computables import ExpectationValue

# Define active space (6e in 6 spatial orbitals)
ncas = 6
nelecas = 6

# perform pyscf mean field
mol = pyscf.M(
    atom="""
 N  0 0 0.7
 N  0 0 -0.7 
""",
    basis="3-21g",
    symmetry=True,
)
myhf = mol.RHF().run()

# perform classical CASCI to yield eigenvalues and determinants and coeffs
mc = mcscf.CASCI(myhf, ncas, nelecas)
mc.symmetry = True
mc.verbose = 0
e_casci = mc.kernel()[0]
print("CAS-CI energy", e_casci)

# load pyscf mean-field into inquanto, choose active space
driver = ChemistryDriverPySCFMolecularRHF.from_mf(
    myhf, frozen=FromActiveSpace(ncas=ncas, nelecas=nelecas)
)

# obtain active space fermionic quantities from mf
ham, space, state = driver.get_system()
print("HF energy", driver.mf_energy)

# define mapping
mapping = QubitMappingJordanWigner()


# analyze CI wave function, obtain determinants and their coefficients
# we save FermionStateString: parity corrected coefficient pairs to
# create the total FermionState
ci_tol = 1e-5

fock_states = {}
coeffs = []

print("   det-alpha,    det-beta,    CI coefficients,   Parity flip")
for c, ia, ib in mc.fcisolver.large_ci(
    mc.ci, ncas, nelecas, tol=ci_tol, return_strs=False
):
    print(
        "   %s       %s      %.15f     %.1f"
        % (ia, ib, c, pyscf_to_inquanto_parity_update(ia, ib))
    )  # ia and ib are spatial alpha/beta
    det_occ_string = [0] * ncas * 2
    for spinorb in range(ncas * 2):
        if spinorb in ia * 2 or spinorb in ib * 2 + 1:
            det_occ_string[spinorb] = 1
    determinant = FermionStateString(det_occ_string)
    # save determinant and parity correct coeff
    fss_ref, coeff_fss = (determinant, c * pyscf_to_inquanto_parity_update(ia, ib))

    coeffs.append(coeff_fss)
    fock_states[fss_ref] = coeff_fss

# create the total Fermion State
fock_states = FermionState(fock_states)

# map to a qubit state.
# If ci_tol is > 0 then the state may need normalizing
qubit_states = fock_states.qubit_encode(mapping).normalized()

# use sparse method to prepare the CI state.
ci_ansatz = MultiConfigurationStateSparse(qubit_states)
print(ci_ansatz.circuit_resources())

# evaluate energy
e_mcsa = ExpectationValue(ci_ansatz, ham.qubit_encode(mapping))
evaluated_energy = e_mcsa.default_evaluate({})
print("Ansatz energy", evaluated_energy)

# compare to FCI
print("Error to CAS-CI", abs(evaluated_energy - e_casci))