r"""A basic Hubbard model example using the spin correlation computable."""

import numpy as np
from matplotlib import pyplot as plt
from pytket.extensions.qiskit import AerStateBackend

from inquanto.ansatzes import (
    RealUnrestrictedBasisRotationAnsatz,
    HamiltonianVariationalAnsatzHubbard,
    SiteLabeling2D,
)
from inquanto.computables.composite import SpinSzComputable, \
    SpinCorrelationComputable
from inquanto.express import DriverGeneralizedHubbard2D
from inquanto.mappings import QubitMappingJordanWigner
from inquanto.protocols import SparseStatevectorProtocol

# This example aims to reproduce properties of a 2D Hubbard model reported in
# https://arxiv.org/abs/2112.02025


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


if __name__ == "__main__":
    # define the parameters of a 2D Hubbard model calculation
    e, t, u, v, nx, ny = (0.0, -1.0, 4.0, 0.0, 2, 4)
    nel_up, nel_dn = 4, 4

    # 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, f_space, f_state = 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

    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

    # build the ground state HVA ansatz
    ansatz = HamiltonianVariationalAnsatzHubbard(compiled_circuit, lat_dims=[nx, ny],
                                                 hamiltonian_operator=hamiltonian,
                                                 step_number=n_layer, site_labeling=labels)
    # print ansatz properties
    print('ANSATZ: ', ansatz.generate_report(), "\n")
    # using parameters from a VQE optimisation (skipped here)
    params_vqe = {"sym0_l0": 0.9481931804543265, "sym0_l1": 1.0382884839454627, "sym0_l2": -0.924132119830008,
                  "sym1_l0": 1.1082735281121958, "sym1_l1": 2.022564270774505, "sym1_l2": -0.08425721713315654,
                  "sym2_l0": -0.6123566351342871, "sym2_l1": 0.7032504509643882, "sym2_l2": 0.9633159435158412,
                  "sym3_l0": -0.6556278859648624, "sym3_l1": -1.010560960893456, "sym3_l2": -0.12415371369806133} #U=4,nel=8
    # substitute the parameters into the ansatz
    ansatz.symbol_substitution(params_vqe)

    # define the protocol to measure observables
    sv = SparseStatevectorProtocol(AerStateBackend())
    # extract the evaluator for given parameters
    sv_evaluator = sv.get_evaluator(params_vqe)

    # define spin Sz (one-body) and spin correlation (two-body)
    spin_sz_comp = SpinSzComputable(f_space, ansatz, mapping)
    spin_corr_comp = SpinCorrelationComputable(f_space, ansatz, mapping)
    # evaluate spin Sz and spin correlation for given parameters
    spin_sz_val = spin_sz_comp.evaluate(evaluator=sv_evaluator)
    spin_corr_val = spin_corr_comp.evaluate(evaluator=sv_evaluator)
    # print spin Sz for all sites:
    print(f"spin Sz on sites = {spin_sz_val}")

    # extract the spin correlation
    spin_corr = np.real(spin_corr_val.data)
    print(f"spin Correlation on sites = {spin_corr}")
    # plot spin correlation
    plt.figure()
    plt.imshow(spin_corr)
    plt.colorbar()
    plt.title(f"Spin correlation")
    plt.show()