r"""A simulation of H2 using the Information Theory QPE algorithm and QPDE protocol."""

from inquanto.mappings import QubitMappingJordanWigner
from inquanto.states import QubitState
from inquanto.computables import ExpectationValue
from inquanto.ansatzes import CircuitAnsatz
from pytket.circuit import Circuit
from pytket.extensions.qiskit import AerBackend, AerStateBackend
from inquanto.protocols import IterativePhaseDiffEstimation
from inquanto.protocols import UStrat
from inquanto.algorithms.phase_estimation._iterative._info_theory import (
    AlgorithmInfoTheoryQPE,
)
from inquanto.operators import QubitOperator, QubitOperatorList
from pytket.circuit import StatePreparationBox
from pytket._tket.circuit import QControlBox
from inquanto.ansatzes import TrotterAnsatz
from inquanto.express import run_vqe
from inquanto.minimizers import MinimizerScipy
from inquanto.express import load_h5

# This example calculates the ground tate energy of H2 using the IQPE algorithm and QPDE protocol.
# The reference state for QPDE is taken as the VQE H2 ground state.

def get_ctrlstate(qstate: QubitState, n_qubits: int = 5):
    """
    Constructs a control state circuit for QPDE based on the given qubit state.

    Assumes the H2 QPDE problem with 5 qubits.

    Args:
        qstate (QubitState): The input qubit state used to construct the
            control state circuit.
        n_qubits: number of qubits in the circuit.


    Returns:
        CircuitAnsatz: The resulting control state circuit wrapped in a
        circuit ansatz.
    """
    # Wrap the state in a StatePreparationBox
    qstate_stv = qstate.get_numeric_representation()
    multi_state_box = StatePreparationBox(qstate_stv)
    # Add control action for QPDE
    ctrlstate = Circuit(n_qubits)
    ctrlstate.add_qcontrolbox(QControlBox(multi_state_box, 1, 1), list(range(n_qubits)))
    ctrlstate_ans = CircuitAnsatz(ctrlstate)
    return ctrlstate_ans

def get_vqe_state(vqe_ansatz: CircuitAnsatz, qubit_hamiltonian: QubitOperator):
    """Constructs and returns the VQE ground state for the provided quantum ansatz and Hamiltonian.

    Args:
        vqe_ansatz: CircuitAnsatz
            The quantum circuit ansatz to be optimized for variational quantum computation.
        qubit_hamiltonian: QubitOperator
            The Hamiltonian of the quantum system to compute the ground state for.

    Returns:
        QubitState
            The quantum state representing the computed ground state of the given
            Hamiltonian.
    """

    # Prepare minimizer and backend
    backend = AerStateBackend()
    minimizer = MinimizerScipy(method="L-BFGS-B")
    # Define and run the VQE algorithm
    vqe = run_vqe(
        vqe_ansatz,
        qubit_hamiltonian,
        backend=backend,
        with_gradient=True,
        minimizer=minimizer,
    )
    # Return the VQE ground state
    parameters = vqe_ansatz.state_symbols.construct_from_dict(vqe.final_parameters)
    vqe_ansatz.symbol_substitution(parameters)
    return vqe_ansatz.to_QubitState()


if __name__ == '__main__':

    # Set up the H2 problem
    h2 = load_h5("h2_sto3g.h5")
    fop = h2["hamiltonian_operator"]
    qop = QubitMappingJordanWigner().operator_map(fop)
    ene_exact = h2["energy_casci"]

    # Prepare the ansatz
    ansatz = TrotterAnsatz(
        QubitOperatorList.from_list([QubitOperator("Y0 X1 X2 X3", 1j)]),
        QubitState([1, 1, 0, 0]),
    )

    # Get the approximate H2 ground state from a VQE calculation
    s0_qstate = get_vqe_state(ansatz, qop)
    # Use the VQE ground state to build the Ctrl-Excitation circuit needed for QPDE
    s0_ctrlstate = get_ctrlstate(s0_qstate, n_qubits=5)
    # QPDE will return the energy difference between H2 and vaccum.
    # So vaccum energy needs to be added
    energy_core = 0.7430177069924179  # The 1e part of the energy.

    # Reference energy of s0_qstate to check QPDE against
    s0_energy = ExpectationValue(s0_qstate, qop).default_evaluate({})

    # Build QPDE protocol and use it for the IQPE algorithm
    time = 0.25 / ene_exact  # 0.25 is the decimal representation of 0.01 (3 bits)
    n_trotter = 1
    evo_ope_exp = qop.trotterize(trotter_number=n_trotter) * time

    # Define Vacuum state as the initial state for QPDE.
    state = CircuitAnsatz(Circuit())

    # Build the QPDE protocol using Aer backend
    backend = AerBackend()
    protocol = IterativePhaseDiffEstimation(
        backend=backend,
        optimisation_level=0,
        n_shots=100,
    ).build(
        state=state,
        ctrl_excit=s0_ctrlstate,
        evolution_operator_exponents=evo_ope_exp,
        u_strat=UStrat.PAULI_EXP_BOX,
    )

    # Set up the Information Theory IQPE algorithm
    n_bits = 3
    k_max = 2 ** n_bits
    resolution = 2 ** (n_bits + 5)
    n_samples = 200

    algorithm = AlgorithmInfoTheoryQPE(
            resolution=resolution,
            k_max=k_max,
            n_samples=n_samples,
            verbose=True
    )

    # Build the algorithm using the QPDE protocol
    algorithm.build(protocol=protocol)
    handles_mapping = algorithm.run_async()
    algorithm.join(handles_mapping)

    # Obtain the final QPDE energy
    mu, sigma = algorithm.final_value()
    energy_mu = -mu / time

    # QPDE returns the energy difference between s0_qstate and vaccum state.
    # Vaccum state includes the core energy.
    # So we need to add the core energy to match the VQE result.
    print(f"FINAL H2 energy from QPDE: {energy_mu + energy_core}")
    print(f"VQE energy to check against: {s0_energy}")