r"""A VQE simulation of H4 in STO-3G using UPS Ansatzes."""

# imports
from pytket.extensions.qiskit import AerStateBackend

from inquanto.algorithms import AlgorithmVQE
from inquanto.ansatzes import (
    FermionSpaceAnsatzsUPS,
    FermionSpaceAnsatztUPS,
    FermionSpaceStateExp,
    PerfectPairing,
)
from inquanto.computables import ExpectationValue, ExpectationValueDerivative
from inquanto.express import load_h5
from inquanto.mappings import QubitMappingJordanWigner
from inquanto.minimizers import MinimizerBasinHopping
from inquanto.protocols import SparseStatevectorProtocol
from inquanto.spaces import FermionSpace
from inquanto.states import FermionState

import random

# In this example we will find the ground state energy of the H4 molecule in the STO-3G basis using VQE. First, we load
# in the system Hamiltonian from the express module.
h4 = load_h5("h4_linear_sto-3g_r90pm.h5", as_tuple=True)
chemical_hamiltonian_operator = h4.hamiltonian_operator

# We now instantiate the space and state objects to construct the UPS ansatze. We have an 8-dimensional Fock space
# and 4 electrons, spin-paired occupying the 2 lowest spatial MOs.
space = FermionSpace(8)
state = FermionState([1, 1, 1, 1, 0, 0, 0, 0])

# Now we can instantiate the sUPS and tUPS ansatzes.
# For the sUPS we can define a specific operator ordering indexed by spatial MOs and type of excitation (single/double)
# This is the "optimal" H4 ordering as defined in https://doi.org/10.1038/s41534-023-00744-2
excitations = [
    ("s", 1, 2),
    ("d", 1, 2),
    ("d", 0, 3),
    ("d", 2, 3),
    ("s", 0, 3),
    ("d", 0, 3),
    ("s", 0, 2),
    ("d", 0, 3),
    ("s", 1, 3),
    ("d", 1, 3),
    ("d", 0, 3),
    ("s", 1, 3),
    ("d", 0, 1),
]

sups_ansatz = FermionSpaceAnsatzsUPS(space, state, excitations)

# For the tUPS as defined in https://doi.org/10.1103/PhysRevResearch.6.023300, we can decide on
# the number of layers, whether to append orbital optimisation and whether to use perfect pairing

# Vanilla tUPS
tups_ansatz = FermionSpaceAnsatztUPS(
    space,
    state,
    layers=1,
    orbital_optimisation=False,
    perfect_pairing=PerfectPairing.NONE,
)

# Orbital Optimised tUPS
tups_ansatz_oo = FermionSpaceAnsatztUPS(
    space,
    state,
    layers=1,
    orbital_optimisation=True,
    perfect_pairing=PerfectPairing.NONE,
)

# Perfect Paired tUPS, usually done with orbital optimisation
tups_ansatz_pp = FermionSpaceAnsatztUPS(
    space,
    state,
    layers=1,
    orbital_optimisation=True,
    perfect_pairing=PerfectPairing.ORBITAL_ENERGY,
)

# In order to run a quantum experiment we must express our Hamiltonian as a qubit operator, to do that we use the
# Jordan--Wigner mapping.
jw = QubitMappingJordanWigner()

ansatzes: list[FermionSpaceStateExp] = [sups_ansatz, tups_ansatz, tups_ansatz_oo, tups_ansatz_pp]
tags = ["N/A", "vanilla", "oo", "pp"]

for ansatz, pp_tag in zip(ansatzes, tags):
    ansatz_name = ansatz.__class__.__name__[18:]
    print(f"{ansatz_name=}\t{pp_tag=}")

    # We find how many optimisable parameters there are in the ansatz, and an estimate of the
    # quantum resources necessary for the ansatz state prep
    print(f"{ansatz.generate_report()['n_parameters']=}")
    print(f"{ansatz.circuit_resources()=}")

    

    # Using perfect pairing requires the Hamiltonian also to be permuted
    if pp_tag == "pp":
        spinorb_perfect_pairing = (
            chemical_hamiltonian_operator.generate_perfect_pairing_spinorb_index_ordering(8)
        )
        spinorb_to_qubit_indices = {s: q for q, s in enumerate(spinorb_perfect_pairing)}
        chemical_hamiltonian_operator = (
            chemical_hamiltonian_operator.to_FermionOperator().permuted_operator(
                spinorb_to_qubit_indices
            )
        )

    qubit_hamiltonian_operator = jw.operator_map(chemical_hamiltonian_operator)

    # The expectation value of the hamiltonian is the computable, which can be sped up by using its derivative
    expectation_value = ExpectationValue(ansatz, qubit_hamiltonian_operator)
    gradient = ExpectationValueDerivative(ansatz, qubit_hamiltonian_operator, ansatz.free_symbols())

    # Since we are doing a state-vector simulation we need to choose a state-vector protocol.
    protocol = SparseStatevectorProtocol(AerStateBackend())

    # Now we can run our VQE experiment after instantiation and calling the build method.
    # We are using the BasinHopping minimiser as defined in https://pubs.acs.org/doi/10.1021/jp970984n
    # As a python wrapper, it can be slow but is at least as good as the base SciPy minimizer, if not better.
    # niter=10 takes approx 5mins to run.
    rng = random.randint(1, 10)
    print(f"{rng=}")

    vqe = (
        AlgorithmVQE(
            objective_expression=expectation_value,
            gradient_expression=gradient,
            minimizer=MinimizerBasinHopping(niter=10, T=1e-3),
            initial_parameters=ansatz.state_symbols.construct_random(rng),
        )
        .build(protocol, protocol)
        .run()
    )

    energy = vqe.generate_report()["final_value"]
    fci_energy = -2.180316614
    hf_energy = h4.energy_hf

    print(
        f"HF energy: {hf_energy:.6f}, Minimum Energy: {energy:.6f}, FCI energy: {fci_energy:.6f}, Difference = {energy - fci_energy:.6f}\n"
    )