r"""Example running VQE for unitary cluster Jastrow ansatz."""

import warnings

from pytket.extensions.qiskit import AerStateBackend

from inquanto.algorithms import AlgorithmVQE
from inquanto.ansatzes import (
    FermionSpaceAnsatzUnitaryClusterJastrow,
    GeneralAnsatz,
    UnitaryClusterJastrowAnsatz,
    UnitaryClusterJastrowModel,
)
from inquanto.computables import ExpectationValue, ExpectationValueDerivative
from inquanto.express import load_h5
from inquanto.minimizers import MinimizerScipy
from inquanto.operators import QubitOperator
from inquanto.protocols import SparseStatevectorProtocol
from inquanto.spaces import FermionSpace
from inquanto.states import FermionState

warnings.filterwarnings("ignore")

# Load hamiltonians for H2 and H4, and their FermionSpace and FermionStates
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


# Loop through all 3 flavours of Jastrow: REAL_K, IMAGINARY_K, GENERAL_K
for m in UnitaryClusterJastrowModel:
    print(f"{m=}")
    if m == UnitaryClusterJastrowModel.GENERAL_K:
        warnings.warn(
            f"Doing VQE with {UnitaryClusterJastrowModel.GENERAL_K} is extremely sensitive to initialised parameters;"
            "it appears to be most stable when parameters initialised near 0, setting variance to 0.1."
            "The Fermionic version appears to be more stable, and may lead to different parameters."
        )

    ucj_ansatz = UnitaryClusterJastrowAnsatz(state, m)
    fermion_ansatz = FermionSpaceAnsatzUnitaryClusterJastrow(space, state, mode=m)

    ansatze: list[GeneralAnsatz] = [fermion_ansatz, ucj_ansatz]
    tags = ["Fermionic", "UCJ"]

    # Loop over fermionic and exact implementations
    # Fermionic extremely quick, exact is slow due to exponentiate/log of matrix and various decompositions

    for ansatz, tag in zip(ansatze, tags[:]):
        print(f"{tag=}")
        exp_val = ExpectationValue(ansatz, h2_hamiltonian)
        grad = ExpectationValueDerivative(ansatz, h2_hamiltonian, ansatz.free_symbols_ordered())
        protocol = SparseStatevectorProtocol(AerStateBackend())

        params = ansatz.state_symbols.construct_random(seed=123456789, sigma=0.11)

        vqe_exp = (
            AlgorithmVQE(
                MinimizerScipy("SLSQP"),
                exp_val,
                gradient_expression=grad,
                initial_parameters=params,
            )
            .build(protocol, protocol)
            .run()
        )

        vqe = vqe_exp.generate_report()
        print(f"{vqe['final_value']=}\t{fci_energy=}\n\n")
        final_params = {k.name: v for (k, v) in vqe_exp.final_parameters.items()}