r"""Use of HVA Ansatz for the Hubbard model. """
from pytket.extensions.qiskit import AerStateBackend
from inquanto.ansatzes import (
    HamiltonianVariationalAnsatzHubbard,
)
from inquanto.ansatzes import SplitFermionTermListHvaHubbard
from inquanto.express import  DriverGeneralizedHubbard2D
from inquanto.ansatzes import SiteLabeling2D
from inquanto.mappings import QubitMappingJordanWigner

# This example shows how to use the Hamiltonian Variational ansatz for Hubbard model.
# 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.
if __name__ == "__main__":
    # define the parameters of a 2D Hubbard model problem
    e, t, u, v, nx, ny = (0.0, -1.0, 8.0, 0.0, 2, 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, as well as the fermion state
    hamiltonian, _, f_state = dgh.get_system()
    print(f"Hamiltonian: {hamiltonian}\n")
    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)

    n_layer = 3  # number of layers in the ansatz used for accuracy adjustment
    # create the HVA
    ansatz = HamiltonianVariationalAnsatzHubbard(f_state, lat_dims=[nx, ny], hamiltonian_operator=hamiltonian,
                                                 step_number=n_layer, site_labeling=labels)
    # print the terms from the split Hamiltonian
    split_terms = ansatz.get_split_terms()
    print(f"Split terms of the Hamiltonian:\n{split_terms}")
    # print the ansatz properties
    # the number of parameters corresponds to the product of n_layer and number of non-zero Hamiltonian split terms:
    # 3 * 4 = 12
    print(f"Ansatz: {ansatz.generate_report()}\n")


    # We can also impose a custom ordering of the terms in the exponents of the ansatz.

    # instantiate a Hubbard-dedicated term list to help define a custom term ordering
    helper_term_list = SplitFermionTermListHvaHubbard()
    # define custom order of terms for the exponents of HVA
    # this is a sequence made up of split terms that appeared above in the output for our Hubbard problem
    # repetition is allowed and sometime might help preserve symmetries
    term_order = [
        helper_term_list.vertical_hop_even,
        helper_term_list.vertical_hop_odd,
        helper_term_list.horizontal_hop_even,
        helper_term_list.hubbard_u,
        helper_term_list.horizontal_hop_even,
        helper_term_list.vertical_hop_odd,
        helper_term_list.vertical_hop_even
    ]

    # create the HVA
    ansatz_custom = HamiltonianVariationalAnsatzHubbard(f_state, lat_dims=[nx, ny], hamiltonian_operator=hamiltonian,
                                                 step_number=n_layer, site_labeling=labels, term_order=term_order)
    # print the terms from the split Hamiltonian
    # this step does not use the custom ordering yet, it performs the same split as before
    split_terms = ansatz_custom.get_split_terms()
    print(f"Split terms of the Hamiltonian:\n{split_terms}") # the same output as before
    # print the ansatz properties
    # custom ordering will affect this output
    # the number of parameters corresponds to n_layer * len(term_order) = 3 * 7 = 21
    print(f"Custom ansatz: {ansatz_custom.generate_report()}")