"""SU(2) lattice gauge-theory model helpers."""

import numpy as np
from numba import typed
from scipy.sparse import csr_matrix, identity
from scipy.sparse.linalg import norm as spnorm
from .quantum_model import QuantumModel
from edlgt.modeling import (
    LocalTerm,
    TwoBodyTerm,
    PlaquetteTerm,
    QMB_hamiltonian,
    QMB_state,
)
from edlgt.modeling import check_link_symmetry, staggered_mask, get_origin_surfaces
from edlgt.operators import (
    SU2_dressed_site_operators,
    SU2_gauge_invariant_states,
    SU2_gen_dressed_site_operators,
)
from edlgt.symmetries import build_parity_operator
import logging

logger = logging.getLogger(__name__)
__all__ = ["SU2_Model"]


class SU2_Model(QuantumModel):
    """SU(2) lattice gauge model with hardcoded and generalized operator sets."""

    def __init__(
        self,
        spin,
        pure_theory,
        bg_list=None,
        sectors=None,
        use_generic_model=False,
        **kwargs,
    ):
        """Initialize the SU(2) model and construct its symmetry sector.

        Parameters
        ----------
        spin : float
            Gauge-link spin representation.
        pure_theory : bool
            If ``True``, exclude matter fields.
        bg_list : list, optional
            Optional background-charge specification.
        sectors : list, optional
            Global symmetry-sector labels (used when matter is present).
        use_generic_model : bool, optional
            If ``True``, force the generalized operator construction.
        **kwargs
            Arguments forwarded to :class:`~edlgt.models.quantum_model.QuantumModel`.
        """
        # Initialize base class with the common parameters
        super().__init__(**kwargs)
        self.spin = spin
        self.pure_theory = pure_theory
        self.background = max(bg_list) if bg_list is not None else 0
        self.bg_list = bg_list if self.background != 0 else None
        self.use_generic_model = use_generic_model
        pure_label = "pure" if self.pure_theory else "with matter"
        logger.info(f"----------------------------------------------------")
        msg = f"({self.dim}+1)D SU(2) LGT {pure_label} j={spin}"
        logger.info(msg)
        if self.bg_list is not None:
            logger.info(f"background charges: {self.bg_list}")
        logger.info(f"----------------------------------------------------")
        # -------------------------------------------------------------------------------
        # Acquire gauge invariant basis and states
        self.gauge_basis, self.singlet_states = SU2_gauge_invariant_states(
            self.spin,
            self.pure_theory,
            lattice_dim=self.dim,
            background=self.background,
        )
        self.gauge_states = {}
        for key in self.singlet_states:
            self.gauge_states[key] = np.array(
                [singlet.J_config for singlet in self.singlet_states[key]], dtype=object
            )
        # -------------------------------------------------------------------------------
        # Acquire operators
        if self.spin < 1 and not self.use_generic_model:
            ops = SU2_dressed_site_operators(
                self.spin,
                self.pure_theory,
                lattice_dim=self.dim,
                background=self.background,
            )
        else:
            # Acquire operators
            ops = SU2_gen_dressed_site_operators(
                self.spin,
                self.pure_theory,
                lattice_dim=self.dim,
                background=self.background,
            )
        # Initialize the operators, local dimension and lattice labels
        self.project_operators(ops, bg_sector_list=self.bg_list)
        # -------------------------------------------------------------------------------
        # GLOBAL SYMMETRIES
        if self.pure_theory:
            global_ops = None
            global_sectors = None
            self.zero_density = None
        else:
            global_ops = [self.ops["N_tot"]]
            global_sectors = sectors
            self.zero_density = True if sectors[0] == self.n_sites else False
        # -------------------------------------------------------------------------------
        # LINK SYMMETRIES
        link_ops = [
            [self.ops[f"T2_p{d}"], -self.ops[f"T2_m{d}"]] for d in self.directions
        ]
        link_sectors = [0 for _ in self.directions]
        # -------------------------------------------------------------------------------
        # SU2 ELECTRIC-FLUX “NBODY” SYMMETRIES
        # only in the pure (no-matter) theory, more than 1D, *and* PBC
        # Constrain, for each direction, the parity of the face/line through the origin:
        # 3D: (Ex → yz-face at x=0) (Ey → xz-face at y=0) (Ez → xy-face at z=0)
        # 2D: (Ex → y-axis at x=0) (Ey → x-axis at y=0)
        if self.pure_theory and not any(self.has_obc):
            logger.info("fixing surface parity fluxes")
            # one flux‐constraint per cartesian direction
            nbody_sectors = list(np.ones(self.dim, dtype=float))
            nbody_ops = []
            nbody_sites_list = typed.List()
            surfaces = get_origin_surfaces(self.lvals)
            nbody_sym_type = "Z"
            if self.dim == 2:
                # in 2D we have two lines through (0,0):
                line_of = {"x": "y", "y": "x"}
                for dir in self.directions:
                    sites = np.array(surfaces[line_of[dir]][1], dtype=np.uint8)
                    nbody_sites_list.append(sites)
                    nbody_ops.append(self.ops[f"P_p{dir}"])
            elif self.dim == 3:
                # in 3D we have three faces through (0,0,0):
                face_of = {"x": "yz", "y": "xz", "z": "xy"}
                for dir in self.directions:
                    sites = np.array(surfaces[face_of[dir]][1], dtype=np.uint8)
                    nbody_sites_list.append(sites)
                    logger.debug(f"{dir} sites: {sites} {surfaces[face_of[dir]][0]}")
                    nbody_ops.append(self.ops[f"P_p{dir}"])
        else:
            # no electric‐flux constraint in 1D, or in OBC or with matter
            nbody_sectors = None
            nbody_ops = None
            nbody_sites_list = None
            nbody_sym_type = None
        # -------------------------------------------------------------------------------
        # GET SYMMETRY SECTOR
        self.get_abelian_symmetry_sector(
            global_ops=global_ops,
            global_sectors=global_sectors,
            link_ops=link_ops,
            link_sectors=link_sectors,
            nbody_ops=nbody_ops,
            nbody_sectors=nbody_sectors,
            nbody_sites_list=nbody_sites_list,
            nbody_sym_type=nbody_sym_type,
        )
        if self.sector_configs is None:
            raise ValueError("No configurations found for the given symmetry sectors")

    def build_Hamiltonian(self, g, m=None, theta=0.0, lambda_noise=0.0):
        """Dispatch to the hardcoded or generalized SU(2) Hamiltonian builder."""
        if self.spin < 1 and not self.use_generic_model:
            self.build_base_Hamiltonian(g, m, theta, lambda_noise)
        else:
            self.build_gen_Hamiltonian(g, m)

    def build_base_Hamiltonian(self, g, m=None, theta=0.0, lambda_noise=0.0):
        """Assemble the hardcoded low-spin SU(2) Hamiltonian.

        Parameters
        ----------
        g : float
            Gauge coupling.
        m : float, optional
            Bare mass parameter.
        theta : float, optional
            Topological-angle parameter.
        lambda_noise : float, optional
            Optional Gauss-law-violating noise strength.
        """
        logger.info(f"----------------------------------------------------")
        logger.info("BUILDING (H)ardocore (G)luon J=1/2 HAMILTONIAN")
        logger.info(f"g={g}, m={m}, theta={theta}, lambda_noise={lambda_noise}")
        logger.info(f"----------------------------------------------------")
        # Hamiltonian Coefficients
        self.SU2_Hamiltonian_couplings(g, m, theta)
        h_terms = {}
        # ---------------------------------------------------------------------------
        # ELECTRIC ENERGY
        op_name = "E_square"
        h_terms[op_name] = LocalTerm(self.ops[op_name], op_name, **self.def_params)
        self.H.add_term(h_terms[op_name].get_Hamiltonian(strength=self.coeffs["E"]))
        # ---------------------------------------------------------------------------
        # PLAQUETTE TERM: MAGNETIC INTERACTION
        if self.dim > 1:
            op_names_list = ["C_px,py", "C_py,mx", "C_my,px", "C_mx,my"]
            op_list = [self.ops[op] for op in op_names_list]
            h_terms["plaq_xy"] = PlaquetteTerm(
                ["x", "y"], op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["plaq_xy"].get_Hamiltonian(
                    strength=self.coeffs["B"], add_dagger=True
                )
            )
        if self.dim == 3:
            # XZ Plane
            op_names_list = ["C_px,pz", "C_pz,mx", "C_mz,px", "C_mx,mz"]
            op_list = [self.ops[op] for op in op_names_list]
            h_terms["plaq_xz"] = PlaquetteTerm(
                ["x", "z"], op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["plaq_xz"].get_Hamiltonian(
                    strength=self.coeffs["B"], add_dagger=True
                )
            )
            # YZ Plane
            op_names_list = ["C_py,pz", "C_pz,my", "C_mz,py", "C_my,mz"]
            op_list = [self.ops[op] for op in op_names_list]
            h_terms["plaq_yz"] = PlaquetteTerm(
                ["y", "z"], op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["plaq_yz"].get_Hamiltonian(
                    strength=self.coeffs["B"], add_dagger=True
                )
            )
        # ---------------------------------------------------------------------------
        if not self.pure_theory:
            # -----------------------------------------------------------------------
            # STAGGERED MASS TERM
            for site in ["even", "odd"]:
                h_terms[f"N_{site}"] = LocalTerm(
                    self.ops["N-1"], "N-1", **self.def_params
                )
                self.H.add_term(
                    h_terms[f"N_{site}"].get_Hamiltonian(
                        self.coeffs[f"m_{site}"], staggered_mask(self.lvals, site)
                    )
                )
            # -----------------------------------------------------------------------
            # HOPPING
            for d in self.directions:
                for site in ["even", "odd"]:
                    op_names_list = [f"Qp{d}_dag", f"Qm{d}"]
                    op_list = [self.ops[op] for op in op_names_list]
                    # Define the Hamiltonian term
                    h_terms[f"{d}_hop_{site}"] = TwoBodyTerm(
                        d, op_list, op_names_list, **self.def_params
                    )
                    mask = staggered_mask(self.lvals, site)
                    self.H.add_term(
                        h_terms[f"{d}_hop_{site}"].get_Hamiltonian(
                            strength=self.coeffs[f"t{d}_{site}"],
                            add_dagger=True,
                            mask=mask,
                        )
                    )
        # -------------------------------------------------------------------------------
        # TOPOLOGICAL TERM
        if self.dim == 3 and np.abs(self.coeffs["theta"]) > 10e-10:
            logger.info("Adding topological term")
            # XY Plane
            op_names_list = ["EzC_px,py", "C_py,mx", "C_my,px", "C_mx,my"]
            op_list = [self.ops[op] for op in op_names_list]
            h_terms["Ez_Bxy"] = PlaquetteTerm(
                ["x", "y"], op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["Ez_Bxy"].get_Hamiltonian(
                    strength=self.coeffs["theta"], add_dagger=True
                )
            )
            # XZ Plane
            op_names_list = ["EyC_px,pz", "C_pz,mx", "C_mz,px", "C_mx,mz"]
            op_list = [self.ops[op] for op in op_names_list]
            h_terms["Ey_Bxz"] = PlaquetteTerm(
                ["x", "z"], op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["Ey_Bxz"].get_Hamiltonian(
                    strength=self.coeffs["theta"], add_dagger=True
                )
            )
            # YZ Plane
            op_names_list = ["ExC_py,pz", "C_pz,my", "C_mz,py", "C_my,mz"]
            op_list = [self.ops[op] for op in op_names_list]
            h_terms["Ex_Byz"] = PlaquetteTerm(
                ["y", "z"], op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["Ex_Byz"].get_Hamiltonian(
                    strength=self.coeffs["theta"], add_dagger=True
                )
            )
        # RANDOM NOISE VIOLATING GAUSS LAW
        noise_requirements = [self.background > 0, lambda_noise > 0]
        noise_requirements += [not obc for obc in self.has_obc]
        if np.all(noise_requirements):
            logger.info("SU2 GAUSS LAW VIOLATING TERM")
            # HOPPING
            op_names_list = ["Vpx_dag", "Vmx"]
            op_list = [self.ops[op] for op in op_names_list]
            # Define the Hamiltonian term
            h_terms["V_hop"] = TwoBodyTerm(
                "x", op_list, op_names_list, **self.def_params
            )
            self.H.add_term(
                h_terms["V_hop"].get_Hamiltonian(
                    strength=-lambda_noise, add_dagger=True
                )
            )
        self.H.build(self.ham_format)

    def build_gen_Hamiltonian(self, g, m=None):
        """Assemble the generalized SU(2) Hamiltonian.

        Parameters
        ----------
        g : float
            Gauge coupling.
        m : float, optional
            Bare mass parameter.
        """
        logger.info(f"----------------------------------------------------")
        logger.info(f"BUILDING generalized HAMILTONIAN with g={g}, m={m}")
        # Hamiltonian Coefficients
        self.SU2_Hamiltonian_couplings(g, m)
        h_terms = {}
        # ---------------------------------------------------------------------------
        # ELECTRIC ENERGY
        op_name = "E_square"
        h_terms[op_name] = LocalTerm(self.ops[op_name], op_name, **self.def_params)
        self.H.add_term(h_terms[op_name].get_Hamiltonian(strength=self.coeffs["E"]))
        # ---------------------------------------------------------------------------
        if not self.pure_theory:
            # -----------------------------------------------------------------------
            # STAGGERED MASS TERM
            for site in ["even", "odd"]:
                h_terms[f"N_{site}"] = LocalTerm(
                    self.ops["N_tot"], "N_tot", **self.def_params
                )
                self.H.add_term(
                    h_terms[f"N_{site}"].get_Hamiltonian(
                        self.coeffs[f"m_{site}"], staggered_mask(self.lvals, site)
                    )
                )
            # --------------------------------------------------------------------
            # Generalized HOPPING
            for d in self.directions:
                for site in ["even", "odd"]:
                    hopping_terms = [
                        [f"Q1_p{d}_dag", f"Q2_m{d}"],
                        [f"Q2_p{d}_dag", f"Q1_m{d}"],
                    ]
                    for ii, op_names_list in enumerate(hopping_terms):
                        op_list = [self.ops[op] for op in op_names_list]
                        # Define the Hamiltonian term
                        h_terms[f"{d}{ii}_hop_{site}"] = TwoBodyTerm(
                            d, op_list, op_names_list, **self.def_params
                        )
                        mask = staggered_mask(self.lvals, site)
                        self.H.add_term(
                            h_terms[f"{d}{ii}_hop_{site}"].get_Hamiltonian(
                                strength=self.coeffs[f"t{d}_{site}"],
                                add_dagger=True,
                                mask=mask,
                            )
                        )
        # -------------------------------------------------------------------------------
        # PLAQUETTE TERM: MAGNETIC INTERACTION
        plaq_list = []
        if self.dim == 3:
            plaquette_directions = ["xy", "xz", "yz"]
        elif self.dim == 2:
            plaquette_directions = ["xy"]
        plaquette_set = [
            ["AB", "AB", "AB", "AB"],
            ["AA", "AB", "BB", "AB"],
            ["AB", "AB", "AA", "BB"],
            ["AA", "AB", "BA", "BB"],
            ["AB", "BB", "AB", "AA"],
            ["AA", "BB", "BB", "AA"],
            ["AB", "BB", "AA", "BA"],
            ["AA", "BB", "BA", "BA"],
            ["BB", "AA", "AB", "AB"],
            ["BA", "AA", "BB", "AB"],
            ["BB", "AA", "AA", "BB"],
            ["BA", "AA", "BA", "BB"],
            ["BB", "BA", "AB", "AA"],
            ["BA", "BA", "BB", "AA"],
            ["BB", "BA", "AA", "BA"],
            ["BA", "BA", "BA", "BA"],
        ]
        if self.dim > 1:
            for pdir in plaquette_directions:
                for p_set in plaquette_set:
                    # DEFINE THE LIST OF CORNER OPERATORS
                    op_names_list = [
                        f"C{p_set[0]}_p{pdir[0]},p{pdir[1]}",
                        f"C{p_set[1]}_p{pdir[1]},m{pdir[0]}",
                        f"C{p_set[2]}_m{pdir[1]},p{pdir[0]}",
                        f"C{p_set[3]}_m{pdir[0]},m{pdir[1]}",
                    ]
                    # CORRESPONDING LIST OF OPERATORS
                    op_list = [self.ops[op] for op in op_names_list]
                    # DEFINE THE PLAQUETTE CLASS
                    plaq_name = f"P{pdir}_" + "".join(p_set)
                    h_terms[plaq_name] = PlaquetteTerm(
                        [pdir[0], pdir[1]],
                        op_list,
                        op_names_list,
                        print_plaq=False,
                        **self.def_params,
                    )
                    # ADD THE HAMILTONIAN TERM
                    self.H.add_term(
                        h_terms[plaq_name].get_Hamiltonian(
                            strength=self.coeffs["B"], add_dagger=True
                        )
                    )
                    # ADD THE PLAQUETTE TO THE LIST OF OBSERVABLES
                    plaq_list.append(plaq_name)
        self.H.build(self.ham_format)

    def check_symmetries(self):
        """Check link-symmetry constraints on measured observables."""
        # CHECK LINK SYMMETRIES
        for ax in self.directions:
            check_link_symmetry(
                ax,
                self.obs_list[f"T2_p{ax}"],
                self.obs_list[f"T2_m{ax}"],
                value=0,
                sign=-1,
            )

    def overlap_QMB_state(self, name):
        """Return predefined benchmark SU(2) basis configurations.

        Parameters
        ----------
        name : str
            Label of a reference configuration.

        Returns
        -------
        numpy.ndarray
            Configuration in the model basis.
        """
        # POLARIZED AND BARE VACUUM in 1D
        if len(self.lvals) == 1:
            if name == "V":
                if self.spin < 1:
                    s1, s2, L, R = 0, 4, 0, 2
                elif self.spin == 1:
                    s1, s2, L, R = 0, 7, 0, 2
                elif self.spin > 1:
                    s1, s2, L, R = 0, 10, 0, 2
            elif name == "PV":
                if self.spin < 1:
                    s1, s2, L, R = 1, 5, 1, 1
                elif self.spin == 1:
                    s1, s2, L, R = 1, 8, 1, 1
                elif self.spin > 1:
                    s1, s2, L, R = 1, 11, 1, 1
            elif name == "T":
                s1, s2, L, R = 6, 12, 1, 1
            elif name == "M":
                s1, s2, L, R = 2, 3, 1, 1
            elif name == "B":
                s1, s2, L, R = 7, 11, 1, 1
            else:
                s1, s2, L, R = 0, 0, 0, 0
            config_state = [s1 if ii % 2 == 0 else s2 for ii in range(self.n_sites)]
            if name == "DW":
                s1, s2, L, R = 0, 4, 0, 2
                config_state = [
                    s1 if ii < self.n_sites // 2 else s2 for ii in range(self.n_sites)
                ]
            elif name == "DW2":
                s1, s2, L, R = 0, 4, 0, 2
                config_state = [
                    s1 if (ii // 2) % 2 == 0 else s2 for ii in range(self.n_sites)
                ]
            if self.has_obc[0]:
                config_state[0] = L
                config_state[-1] = R
        # POLARIZED AND BARE VACUUM in 2D
        else:
            if not self.pure_theory:
                if self.has_obc[0]:
                    if name == "V":
                        config_state = [0, 9, 0, 4, 4, 0, 9, 0]
                    elif name == "PV1":
                        config_state = [1, 12, 3, 5, 5, 2, 11, 1]
                    elif name == "PV2":
                        config_state = [1, 11, 1, 5, 5, 3, 10, 1]
                else:
                    config_state = [0, 9, 0, 9, 9, 0, 9, 0]
            else:
                if name == "V":
                    config_state = [0 for _ in range(self.n_sites)]
        config_state = config_state
        return np.array(config_state)

    def SU2_Hamiltonian_couplings(self, g, m=None, theta=0.0):
        """Set SU(2) Hamiltonian couplings from physical parameters.

        Parameters
        ----------
        g : float
            Gauge coupling.
        m : float, optional
            Bare mass parameter (used when matter fields are present).
        theta : float, optional
            Topological-angle parameter.

        Returns
        -------
        None
            Couplings are stored in ``self.coeffs``.

        Notes
        -----
        In the current convention, the Hamiltonian is rescaled so that the
        hopping term is dimensionless (as used in the project implementation and
        in PRX Quantum 5, 040309).

        The rescaling summary used in the code is:

        - hopping: original coupling ``1/2`` -> ``2 * sqrt(2)``
        - electric term: original ``g0^2 / 2`` -> ``8 g^2 / 3``
        - magnetic term: rescaled convention factor applied in the model
        - the symbol ``g`` in this implementation is used as the rescaled
          coupling (effectively a ``g^2`` convention)

        DFL-project convention (reference values often used in scripts):

        - ``E = 8 * g / 3``
        - ``B = -3 / g``
        - ``t = 2 * sqrt(2)``

        String-breaking convention (alternative reference choice):

        - ``E = g``
        - ``B = -1``
        - ``t = 1``
        """
        if self.dim == 1:
            E = 8 * g / 3
            B = 0
        else:
            E = g / 2
            B = -1 / (2 * g)
        # Dictionary with Hamiltonian COEFFICIENTS
        self.coeffs = {
            "g": g,
            "E": E,  # ELECTRIC FIELD coupling
            "B": B,  # MAGNETIC FIELD coupling
            "theta": -theta * g,  # THETA TERM coupling
        }
        if not self.pure_theory and m is not None:
            # The correct hopping in original units should be 1/2
            t = 2 * np.sqrt(2)
            self.coeffs |= {
                "tx_even": -complex(0, t),  # x HOPPING (EVEN SITES)
                "tx_odd": -complex(0, t),  # x HOPPING (ODD SITES)
                "ty_even": -t,  # y HOPPING (EVEN SITES)
                "ty_odd": t,  # y HOPPING (ODD SITES)
                "tz_even": -t,  # z HOPPING (EVEN SITES)
                "tz_odd": t,  # z HOPPING (ODD SITES)
                "m": m,
                "m_odd": -m,  # EFFECTIVE MASS for ODD SITES
                "m_even": m,  # EFFECTIVE MASS for EVEN SITES
            }

    def build_local_Hamiltonian(self, g, m, R0):
        """Build a local effective Hamiltonian around one 1D site.

        Parameters
        ----------
        g : float
            Gauge coupling.
        m : float
            Bare mass parameter.
        R0 : int
            Central lattice site.
        """
        logger.info(f"----------------------------------------------------")
        logger.info(f"BUILDING local HAMILTONIAN around {R0}")
        # -------------------------------------------------------------------------------
        if self.dim > 1:
            raise ValueError(f"Local Hamiltonian valid only for D=1, got {self.dim}")
        self.Hlocal = QMB_hamiltonian(self.lvals, size=self.sector_dim)
        # -------------------------------------------------------------------------------
        # Hamiltonian Coefficients
        self.SU2_Hamiltonian_couplings(g, m)
        # -------------------------------------------------------------------------------
        hterms = {}
        # ELECTRIC HAMILTONIAN
        hterms["E2"] = LocalTerm(self.ops["E_square"], "E_square", **self.def_params)
        # MASS TERM
        hterms["N"] = LocalTerm(self.ops["N-1"], "N-1", **self.def_params)
        # HOPPING
        op_names_list = ["Qpx_dag", "Qmx"]
        op_list = [self.ops[op] for op in op_names_list]
        hterms["hop"] = TwoBodyTerm("x", op_list, op_names_list, **self.def_params)
        # DAGGER HOPPING
        op_names_list = ["Qpx", "Qmx_dag"]
        op_list = [self.ops[op] for op in op_names_list]
        hterms["hop_dag"] = TwoBodyTerm("x", op_list, op_names_list, **self.def_params)
        hop_coeff = self.coeffs["tx_even"]
        hop_coeff_half = 0.5 * hop_coeff
        # ---------------------------------------------------------------------------
        self.Hlocal.add_term(
            hterms["E2"].get_Hamiltonian(self.coeffs["E"], self.get_mask([R0]))
        )
        # ---------------------------------------------------------------------------
        self.Hlocal.add_term(
            hterms["N"].get_Hamiltonian(
                ((-1) ** R0) * self.coeffs["m"], self.get_mask([R0])
            )
        )
        # ---------------------------------------------------------------------------
        # Add the hopping term (j,j+1)
        self.Hlocal.add_term(
            hterms["hop"].get_Hamiltonian(hop_coeff_half, mask=self.get_mask([R0]))
        )
        # ---------------------------------------------------------------------------
        # Add the term (j-1,j)
        self.Hlocal.add_term(
            hterms["hop"].get_Hamiltonian(
                hop_coeff_half,
                mask=self.get_mask([(R0 - 1) % self.n_sites]),
            )
        )
        # ---------------------------------------------------------------------------
        # Add the hermitian conjugate
        # Add the term (j,j+1)
        self.Hlocal.add_term(
            hterms["hop_dag"].get_Hamiltonian(-hop_coeff_half, mask=self.get_mask([R0]))
        )
        # ---------------------------------------------------------------------------
        # Add the term (j-1,j)
        self.Hlocal.add_term(
            hterms["hop_dag"].get_Hamiltonian(
                -hop_coeff_half,
                mask=self.get_mask([(R0 - 1) % self.n_sites]),
            )
        )

    def get_mask(self, sites_list):
        """Build a boolean mask selecting the given sites."""
        mask = np.zeros(self.lvals, dtype=bool)
        for site in sites_list:
            mask[site] = True
        return mask

    def local_parity_labels(self, wrt_site):
        """
        Local action of pure spatial inversion (left-right swap)
        on the 6d dressed-site SU(2) basis.

        Basis labels:
        0: V, J=(0,0)
        1: V, J=(1/2,1/2)  (link singlet)
        2: (1/2,0,1/2)     (matter + right link)
        3: (1/2,1/2,0)     (matter + left link)
        4: P, J=(0,0)
        5: P, J=(1/2,1/2)  (link singlet)

        Returns
        -------
        tuple
            ``(loc_perm, loc_phase)`` where ``loc_perm`` contains the mapped
            local basis indices and ``loc_phase`` contains the corresponding
            phases (``+1`` or ``-1``).
        """
        if self.dim > 1:
            raise ValueError(f"Does not work in D={self.dim}>1")
        if self.spin > 0.5:
            raise ValueError(f"Does not work spin j={self.spin}>0.5")
        logger.info(f"----------------------------------------------------")
        if wrt_site:
            # local map for "parity" (P_loc)
            loc_perm = np.array([0, 1, 3, 2, 4, 5], dtype=np.int32)
            loc_phase = np.array([1, -1, 1, -1, 1, -1], dtype=np.int32)
        else:
            # local map for "parity + particle-hole" (or C ∘ P_loc)
            loc_perm = np.array([4, 5, 3, 2, 0, 1], dtype=np.int32)
            loc_phase = np.array([1, -1, 1, 1, 1, -1], dtype=np.int32)
        return loc_perm, loc_phase

    def get_parity_inversion_operator(self, wrt_site):
        """Construct the parity inversion operator in the current sector."""
        # INVERSION SYMMETRY
        if self.dim < 1:
            raise NotImplementedError(f"Cannot apply PARITY in D={self.dim}>1")
        # Acquire the phases and the permutation of the basis
        if wrt_site:
            inversion_wrt = f"site {self.n_sites//2-1}"
        else:
            inversion_wrt = f"bond ({self.n_sites//2-1}, {self.n_sites//2})"
        loc_perm, loc_phase = self.local_parity_labels(wrt_site)
        wrt_site = np.uint8(0 if wrt_site else 1)
        logger.info(f"----------------------------------------------------")
        logger.info(f"Parity symmetry: inversion wrt {inversion_wrt}")
        # Build the projector
        r, c, d = build_parity_operator(
            self.sector_configs, loc_perm, loc_phase, wrt_site
        )
        shape = (self.sector_dim, self.sector_dim)
        logger.info(f"r {r.shape} c {c.shape}, d {d.shape}, {shape}")
        self.parityOP = csr_matrix((d, (r, c)), shape=shape)

    def check_parity_inversion_operator(self):
        """Validate algebraic and dynamical properties of the parity operator."""
        if self.momentum_basis is not None:
            # Build the projector from the momentum sector to the global one
            Pk = self._basis_Pk_as_csr()
            # Check the commutator between the Hamiltonian H and parity P
            err_PH = spnorm(
                self.parityOP @ Pk @ self.H.Ham @ Pk.getH()
                - Pk @ self.H.Ham @ Pk.getH() @ self.parityOP
            )
        else:
            err_PH = spnorm(self.parityOP @ self.H.Ham - self.H.Ham @ self.parityOP)
        logger.info(f"|PH-HP| = {err_PH}")
        # -----------------------------------------------------------
        I = identity(self.sector_dim)
        # P^2 = I
        err_P2 = spnorm(self.parityOP @ self.parityOP - I)
        logger.info(f"|P^2-I| = {err_P2}")
        # Hermitian: P^† = P
        err_herm = spnorm(self.parityOP.getH() - self.parityOP)
        logger.info(f"|P^†-P| = {err_herm}")
        # Unitary: P^† P = I
        err_unit = spnorm(self.parityOP.getH() @ self.parityOP - I)
        logger.info(f"|P^†P-I| = {err_unit}")
        # nnz per row
        nnz_per_row = np.diff(self.parityOP.indptr)
        # nnz per col
        nnz_per_col = np.diff(self.parityOP.tocsc().indptr)
        logger.info(f"row nnz counts: {np.unique(nnz_per_row)}")
        logger.info(f"col nnz counts: {np.unique(nnz_per_col)}")
        # values
        unique_vals = np.unique(np.round(self.parityOP.data, 8))
        logger.info(f"unique data values in P: {unique_vals}")
        if self.momentum_basis is None and self.sector_configs is not None:
            for ii, cfg in enumerate(self.sector_configs):
                s1 = self.get_qmb_state_from_configs([cfg])
                Ss1 = QMB_state(s1, self.lvals, self.loc_dims)
                Ps1 = QMB_state(self.parityOP @ s1, self.lvals, self.loc_dims)
                exp_s1 = Ss1.expectation_value(self.H.Ham)
                exp_Ps1 = Ps1.expectation_value(self.H.Ham)
                Hpsi = self.H.Ham @ s1  # H|psi>
                Ppsi = self.parityOP @ s1  # P|psi>
                HPpsi = self.H.Ham @ Ppsi  # HP|psi>
                PHpsi = self.parityOP @ Hpsi  # PH|psi>
                diff = HPpsi - PHpsi
                absnorm = np.linalg.norm(diff)
                if np.abs(exp_Ps1 - exp_s1) > 1e-14:
                    logger.info(f"{exp_s1}")
                    logger.info(f"{exp_Ps1}")
                    raise ValueError(f"config {ii} {cfg} not symmetrized")
                if np.abs(absnorm) > 1e-14:
                    logger.info("==============================================")
                    logger.info("State i")
                    Ss1.get_state_configurations(1e-3, self.sector_configs)
                    logger.info("State P|i>")
                    Ps1.get_state_configurations(1e-3, self.sector_configs)
                    A = QMB_state(HPpsi, self.lvals, self.loc_dims)
                    logger.info("State HP|i>")
                    A.get_state_configurations(1e-3, self.sector_configs)
                    B = QMB_state(PHpsi, self.lvals, self.loc_dims)
                    logger.info("State PH|i>")
                    B.get_state_configurations(1e-3, self.sector_configs)
                    logger.info(f"VDOT{np.vdot(PHpsi,HPpsi)}")
                    logger.info(f"Relative |HP-PH|={absnorm}")
                    raise ValueError(f"config {ii} {cfg} not symmetrized")

    def print_state_config(self, config, amplitude=None):
        """Log a readable per-site decomposition of an SU(2) basis configuration."""
        logger.info(f"----------------------------------------------------")
        msg = f"CONFIG {config}"
        if amplitude is not None:
            msg += f" |psi|^2={np.abs(amplitude)**2:.8f}"
        logger.info(msg)
        logger.info(f"----------------------------------------------------")
        # Choose width (sign + digits).
        max_abs = 1
        max_abs = max(max_abs, int(abs(getattr(self, "spin", 1))))
        max_abs = max(max_abs, int(abs(getattr(self, "background", 0))))
        entry_width = len(str(max_abs)) + 2
        logger.info(f"{'':>19s}{self._format_header(entry_width)}")
        for ii, cfg_idx in enumerate(config):
            loc_basis_state = self.gauge_states_per_site[ii][cfg_idx]
            state_str = "[" + " ".join(f"{val}" for val in loc_basis_state) + " ]"
            logger.info(f"SITE {ii:>2d} state {cfg_idx:>3d}: {state_str}")

    def _local_state_labels(self) -> list[str]:
        """Return column labels used by :meth:`print_state_config`."""
        labels: list[str] = []
        # Background first (only if present in your effective gauge states)
        if getattr(self, "background", 0) > 0:
            labels.append("bg")
        # Matter occupation (only if not pure theory)
        if not getattr(self, "pure_theory", True):
            labels.append("M")
        # Links: -x,-y,-z,+x,+y,+z up to dim
        dim = int(getattr(self, "dim", 0))
        dirs = "xyz"[:dim]
        labels += [f"-{d}" for d in dirs] + [f"+{d}" for d in dirs]
        return labels

    def _format_header(self, entry_width: int) -> str:
        """Format the header row for :meth:`print_state_config`."""
        labels = self._local_state_labels()
        header = "[" + " ".join(f"{lab:>{entry_width}s}" for lab in labels) + " ]"
        return header
