r"""Load chemistry and condensed matter Hamiltonians from HamLib."""

from inquanto.extensions.hamlib import HamLibDatasetLoader
import numpy as np

path0 = "https://portal.nersc.gov/cfs/m888/dcamps/hamlib/"

# Example chemistry Hamiltonian
url_chem = (
    path0 + "chemistry/electronic/hydrogen_data/H2_linear/ES_H2_linear_R0.5_sto-6g.zip"
)
hld_chem = HamLibDatasetLoader(url_chem)
hld_chem.print_hdf5_structure()
print("Qubit operator:", hld_chem.get_qubit_operator("ham_JW"))
print("Fermion operator:", hld_chem.get_fermion_operator("ham_molec"))
hld_chem.print_hdf5_ham_properties()

# Printing available condensed matter Hamiltonians
path_condensed = path0 + "condensedmatter/fermihubbard/"
csv = path_condensed + "condensedmatter_fermihubbard.csv"
hld_csv = HamLibDatasetLoader(csv)
ind_range = np.s_[0:10]
str_search = "Lx-10_Ly-10_Lz-10"
hld_csv.print_keys_csv(index_range=ind_range, string_search=str_search)

# Example condensed matter Hamiltonian: qubit operator
url_qubit_condensed = path_condensed + "FH_D-1.zip"
hld_qubit_condensed = HamLibDatasetLoader(url_qubit_condensed)
qubit_model = "fh-graph-1D-grid-nonpbc-qubitnodes_Lx-2_U-2_enc-jw"
qop = hld_qubit_condensed.get_qubit_operator(qubit_model)
print("Qubit operator for model:", qubit_model)
print(qop)

# Example condensed matter Hamiltonian: fermion operator
url_fermion_condensed = path_condensed + "all-fh-symbolic-unencoded.zip"
hld_fermion_condensed = HamLibDatasetLoader(url_fermion_condensed)
fermion_model = "graph-1D-grid-nonpbc-qubitnodes_Lx-2_U-2"
fop = hld_fermion_condensed.get_fermion_operator(fermion_model)
print("Fermion operator for model:", fermion_model)
print(fop)