r"""A VQE simulation of H2 in STO-3G using a HVA Ansatz."""

from pytket.extensions.qiskit import AerStateBackend
from inquanto.algorithms import AlgorithmVQE
from inquanto.ansatzes import HamiltonianVariationalAnsatzChemistry
from inquanto.computables import ExpectationValue
from inquanto.express import load_h5
from inquanto.mappings import QubitMappingJordanWigner
from inquanto.minimizers import MinimizerScipy
from inquanto.protocols import SparseStatevectorProtocol
from inquanto.spaces import FermionSpace


# Load in the system from the express module.
h2_sto3g = load_h5("h2_sto3g.h5", as_tuple=True)
ham = h2_sto3g.hamiltonian_operator
# In order to construct an ansatz we need to create a space and a state object.
space = FermionSpace(4)
state = space.generate_occupation_state_from_list([1, 1, 0, 0])

# Initiate the HV ansatz for chemistry with one layer
ansatz = HamiltonianVariationalAnsatzChemistry(reference=state, hamiltonian_operator=ham,
                                               qubit_mapping=QubitMappingJordanWigner(), step_number=1)

# Create the computable for the quantity of interest, then build and run the algorithm.
ev = ExpectationValue(ansatz, QubitMappingJordanWigner().operator_map(ham))
# Initialize the algorithm, supplying a set of initial parameters for the HV ansatz
# here our optimization process only uses the total energy, not the gradients
vqe = AlgorithmVQE(
    objective_expression=ev,
    minimizer=MinimizerScipy(),
    initial_parameters=ansatz.state_symbols.construct_random(mu=0.0, sigma=0.2),
)
vqe.build(
    protocol_objective=SparseStatevectorProtocol(AerStateBackend()),
)
vqe.run()

print("Minimum Energy: {}".format(vqe.final_value))