r"""A VQE simulation of a 2x4 Hubbard model using a HVA Ansatz."""

import numpy as np
from pytket.extensions.qiskit import AerStateBackend

from inquanto.algorithms import AlgorithmVQE
from inquanto.ansatzes import (
    RealUnrestrictedBasisRotationAnsatz,
    HamiltonianVariationalAnsatzHubbard,
)
from inquanto.computables import ExpectationValue, ExpectationValueDerivative
from inquanto.express import  DriverGeneralizedHubbard2D
from inquanto.ansatzes import SiteLabeling2D
from inquanto.mappings import QubitMappingJordanWigner
from inquanto.minimizers import MinimizerScipy
from inquanto.protocols import SparseStatevectorProtocol


# prepare an initial state for Hubbard model solution
# the initial state is a single Slater determinant diagonalising the non-interacting problem
def get_slater_state_RBR(hopping_matrix, n_particles):
    """Prepares a Slater determinant ground state of a non-interacting Hamiltonian.
    Uses a basis rotation method included in the RealBasisRotationAnsatz.

    Args:
        hopping_matrix: Non interacting hamiltonian matrix in the basis of Hubbard sites.
        n_particles: Number of spinless particles. Assumes nel_up = nel_dn.

    Returns: Circuit representing the ground state of the non-interacting Hamiltonian.

    """
    # diagonalise the non-interacting problem
    n_sites = hopping_matrix.shape[0]  # number of Hubbard sites
    eigvals, eigvecs = np.linalg.eigh(
        hopping_matrix
    )  # diagonalise the non-interacting problem
    reference = [1] * 2 * n_particles + [0] * (
        n_sites - n_particles
    ) * 2  # prepare a Slater determinant

    # rotate to the basis diagonalising the non-interacting problem
    rbr = RealUnrestrictedBasisRotationAnsatz(
        reference
    )  # initialise the basis rotation
    _params = rbr.ansatz_parameters_from_unitary(
        eigvecs, eigvecs
    )  # extract basis rotation parameters
    rbr_circuit = rbr.get_circuit(_params)  # generate basis rotation circuit
    return rbr_circuit

# This example shows how to find the ground state of a 2x4 square lattice Hubbard model using VQE.
# Hamiltonian Variational ansatz (HVA) is used.
# It splits the Hubbard Hamiltonian into various terms, e.g. vertical or horizontal hopping,
# Hubbard U term, and applies these terms in an exponentiated form to a given reference state.
# The HVA is used in each iteration of the minimizer searching for the ground state.

# define the parameters of a 2D Hubbard model calculation
e, t, u, v, nx, ny = (0.0, -1.0, 8.0, 0.0, 2, 4)
nel_up, nel_dn = 3, 3

# initialise the Hubbard model driver
labels = SiteLabeling2D.snake
dgh = DriverGeneralizedHubbard2D(
    e=e, u=u, t=t, v=v, n_x=nx, n_y=ny, site_labeling=labels, pbc=False
)

# extract the Hubbard Hamiltonian, and its one-body part separately
hamiltonian, _, _ = dgh.get_system()
one_body, _ = dgh.get_one_two_body_terms()

# define fermion to qubit mapping and map the Hamiltonian
mapping = QubitMappingJordanWigner()
qubit_op = mapping.operator_map(hamiltonian)

# prepare an initial state for Hubbard model solution
# the initial state is a single Slater determinant diagonalising the non-interacting problem
slater_circuit = get_slater_state_RBR(
    one_body.todense(), nel_up
)  # spinless electron number, nel_u=nel_d

# define the ansatz for a Hubbard model VQE calculation
init_params = [
    0.9829,
    0.6825,
    -0.5875,
    1.2791,
    1.9934,
    -0.004,
    -0.8857,
    0.4147,
    0.9944,
    -0.8973,
    -1.0177,
    0.0789,
]  # initial guess for VQE
backend = AerStateBackend()
compiled_circuit = backend.get_compiled_circuit(
    slater_circuit
)  # compile the slater state-prep circuit
n_layer = 3  # number of layers in the ansatz used for accuracy adjustment
ansatz = HamiltonianVariationalAnsatzHubbard(compiled_circuit, lat_dims=[nx, ny], hamiltonian_operator=hamiltonian,
                                             step_number=n_layer, site_labeling=labels)

# build the VQE algorithm
minimizer = MinimizerScipy(method="L-BFGS-B")
expectation_value = ExpectationValue(
    ansatz, qubit_op
)  # define the expectation value computable
derivative = ExpectationValueDerivative(
    ansatz, qubit_op, ansatz.state_symbols.as_set()
)  # define gradient
vqe = AlgorithmVQE(
    objective_expression=expectation_value,
    minimizer=minimizer,
    gradient_expression=derivative,
    initial_parameters=ansatz.state_symbols.construct_from_array(init_params),
)
protocol_objective = SparseStatevectorProtocol(
    backend
)  # define statevector protocols for the VQE experiment
protocol_gradient = SparseStatevectorProtocol(backend)
vqe.build(
    protocol_objective=protocol_objective, protocol_gradient=protocol_gradient
)
vqe.run()

# print final VQE energy
print("VQE Energy:    ", vqe.final_value)
print("VQE Parameters:", vqe.final_parameters)