r"""1D and 2D Hubbard driver example"""

from inquanto.express import (
    DriverGeneralizedHubbard1D,
    DriverGeneralizedHubbard2D,
    SiteLabelling2D,
)
from inquanto.mappings import QubitMappingJordanWigner

# Define the parameters of the Hubbard model: on-site energies "e", hopping "t", on-site repulsion "U"
# and nearest neighbor repulsion "V"
e, t, u, v = [0.0, -1.0, 2.0, 0.0]
# Define the number of sites for 1D Hubbard model
n = 2
# Initialise the Hubbard driver with periodic boundary conditions ("ring=True")
dgh_1 = DriverGeneralizedHubbard1D(n, e, t, u, v, ring=True)

# Get the Hamiltonian operator "fop_1", fermion space and initial fermion state for VQE
fop_1, fermion_space_1, fermion_state_1 = dgh_1.get_system()
# Define the fermion to qubit mapping
mapping = QubitMappingJordanWigner()
# Map the Hamiltonian operator to the qubit operator
qop_1 = mapping.operator_map(fop_1)


# Print the outputs
print("Inquanto's express 2 site Hubbard model:")
print("Fermion Hamiltonian:", fop_1)
print("Qubit Hamiltonian", qop_1)
print("# terms:", len(qop_1.terms))


# Use the same parameters as above to create a 2D Hubbard model
# Define the 2D lattice size for Hubbard model
nx, ny = [2, 2]
# Set the periodic boundary conditions to True
pbc = True
# Define the scheme for site labelling as the standard row-major ordered grid labelling of sites
site_labelling = SiteLabelling2D.grid
# Initialize the 2D Hubbard model
dgh_2 = DriverGeneralizedHubbard2D(
    nx, ny, e, t, u, v, pbc=pbc, site_labelling=site_labelling
)

# Get the Hamiltonian operator "fop_2", fermion space and initial fermion state for VQE
fop_2, fermion_space_2, fermion_state_2 = dgh_2.get_system()
# Map the Hamiltonian operator to the qubit operator using the same Jordan Wigner mapping as above
qop_2 = mapping.operator_map(fop_2)

# Print the outputs
print("\nInquanto's express 2x2 site Hubbard model:")
print("Fermion Hamiltonian:", fop_2.df())
print("Qubit Hamiltonian", qop_2.df())
print("# terms:", len(qop_2.terms))