r"""Example demonstrating the use of unitary cluster Jastrow ansatz, both exact and fermionic implementations."""
from inquanto.states import FermionState
from inquanto.spaces import FermionSpace
from inquanto.ansatzes import UnitaryClusterJastrowAnsatz, UnitaryClusterJastrowModel, FermionSpaceAnsatzUnitaryClusterJastrow
from inquanto.computables import ExpectationValue
from inquanto.express import load_h5
from inquanto.operators import QubitOperator

from sympy import Symbol


h2_data = load_h5("h2_sto3g.h5")
h2_hamiltonian: QubitOperator = h2_data["hamiltonian_operator"].to_FermionOperator().qubit_encode()

space, state = FermionSpace(4), FermionState([1, 1, 0, 0])
fci_energy = -1.136


# These are VQE parameters found by running the example script in the algorithms/vqe folder
h2_optimal_ucj_vqe_params = {
    UnitaryClusterJastrowModel.REAL_K:{
        Symbol("Im-(J0)_{0}^{1}"): -0.3882322739560017,
        Symbol("Im-(J0)_{0}^{2}"): -0.10671614600771638,
        Symbol("Im-(J0)_{0}^{3}"): 1.1809444840062906,
        Symbol("Im-(J0)_{1}^{2}"): -0.38957106509030526,
        Symbol("Im-(J0)_{1}^{3}"): 0.041859626611391074,
        Symbol("Im-(J0)_{2}^{3}"): -0.38980959900100637,
        Symbol("Re-(K0)_{0}^{2}"): 1.3439355651187306,
        Symbol("Re-(K0)_{1}^{3}"): -0.23051964890038554
    },
    UnitaryClusterJastrowModel.IMAGINARY_K:{
        Symbol("Im-(J0)_{0}^{1}"): -1.1737763239366452,
        Symbol("Im-(J0)_{0}^{2}"): -0.10671609482591318,
        Symbol("Im-(J0)_{0}^{3}"): 1.9673637788895841,
        Symbol("Im-(J0)_{1}^{2}"): -1.1748427328351632,
        Symbol("Im-(J0)_{1}^{3}"): 0.04185961623443913,
        Symbol("Im-(J0)_{2}^{3}"): 0.39458653057272636,
        Symbol("Im-(K0)_{0}^{2}"): 2.9117078609172005,
        Symbol("Im-(K0)_{1}^{3}"): 0.23511607353377678
    },
    UnitaryClusterJastrowModel.GENERAL_K:{
        Symbol("Im-(J0)_{0}^{1}"): -0.03621567711405116,
        Symbol("Im-(J0)_{0}^{2}"): -0.10671614029671447,
        Symbol("Im-(J0)_{0}^{3}"): 0.04849712208638738,
        Symbol("Im-(J0)_{1}^{2}"): 0.06811138114087331,
        Symbol("Im-(J0)_{1}^{3}"): 0.041859639288645736,
        Symbol("Im-(J0)_{2}^{3}"): -0.06706134522398013,
        Symbol("Im-(K0)_{0}^{2}"): 0.5805185634243769,
        Symbol("Im-(K0)_{1}^{3}"): -0.3805352397324376,
        Symbol("Re-(K0)_{0}^{2}"): -0.33434555052086185,
        Symbol("Re-(K0)_{1}^{3}"): 0.657194688556897
    }
}

# These are VQE energies found by running the example script in the algorithms/vqe folder
h2_ucj_vqe_energies = {
    UnitaryClusterJastrowModel.REAL_K:-1.1356619897094613,
    UnitaryClusterJastrowModel.IMAGINARY_K:-1.136048120079276,
    UnitaryClusterJastrowModel.GENERAL_K:-1.1368405637503423
}

# Loop through REAL_K, IMAGINARY_K, GENERAL_K versions
for m in UnitaryClusterJastrowModel:

    print("Fermionic") # Suffers from Trotter error
    fermion_ansatz = FermionSpaceAnsatzUnitaryClusterJastrow(space, state, k=1, mode=m)
    print(f"{m=}\n{fermion_ansatz.generate_report()}\n{len(fermion_ansatz.free_symbols())=}")
    
    print("Exact") # No Trotter due to the generalized decomposition of the orbital rotation matrix
    exact_ansatz = UnitaryClusterJastrowAnsatz(state, mode=m, k=1)
    print(f"{m=}\n{exact_ansatz.generate_report()}\n{len(exact_ansatz.free_symbols())=}")
    ansatz_params = h2_optimal_ucj_vqe_params[m]
    # Fermionic parameters converted to Givens/generalized Givens rotation angles
    circuit_params = exact_ansatz.map_ansatz_parameters_to_circuit_parameters(ansatz_params)
    # print(f"{ansatz_params=}\n{circuit_params=}\n\n")


    # Expectation Value computable requires the ansatz parameters, with the energy matching what was found during VQE
    # This should print 0
    exp_val = ExpectationValue(exact_ansatz, h2_hamiltonian).default_evaluate(ansatz_params)
    check_energy = h2_ucj_vqe_energies[m]
    print(f"{exp_val-check_energy=}\n\n")