"""Reference partial checks for symmetry-sector construction.

These helpers implement the older, conceptually direct partial checks used as
validation baselines for the optimized iterative builders.
"""

import numpy as np
from numba import njit

__all__ = [
    "check_link_sym_partial",
    "check_nbody_sym_partial",
    "check_Z2_nbody_sym_partial",
    "check_glob_sym_partial",
]


@njit(cache=True)
def check_link_sym_partial(config, sym_op_diags, sym_sectors, pair_list):
    """Reference partial check of link symmetries for a configuration prefix."""
    n_sites = config.shape[0]
    num_lattice_directions = len(pair_list)
    for idx in range(num_lattice_directions):
        pairs_for_direction = pair_list[idx]
        num_pairs = pairs_for_direction.shape[0]
        for pair_idx in range(num_pairs):
            site_indices = pairs_for_direction[pair_idx]
            if site_indices[0] >= n_sites or site_indices[1] >= n_sites:
                continue
            sym_value = 0.0
            for op_idx in range(2):
                site_index = site_indices[op_idx]
                op_diag = sym_op_diags[idx, op_idx, site_index, :]
                sym_value += op_diag[config[site_index]]
            if not np.isclose(sym_value, sym_sectors[idx], atol=1e-13):
                return False
    return True


@njit(cache=True)
def check_nbody_sym_partial(config, sym_op_diags, sym_sectors, nsites_list):
    """Reference partial check of additive n-body symmetries."""
    n_sites = config.shape[0]
    n_symmetries = sym_op_diags.shape[0]
    for sym_idx in range(n_symmetries):
        site_indices = nsites_list[sym_idx]
        nsites_per_symmetry = len(site_indices)
        if np.any(site_indices >= n_sites):
            continue
        sym_value = 0.0
        for ii in range(nsites_per_symmetry):
            site_index = site_indices[ii]
            op_diag = sym_op_diags[sym_idx, site_index, :]
            sym_value += op_diag[config[site_index]]
        if not np.isclose(sym_value, sym_sectors[sym_idx], atol=1e-13):
            return False
    return True


@njit(cache=True)
def check_Z2_nbody_sym_partial(config, sym_op_diags, sym_sectors, nsites_list):
    """Reference partial check of multiplicative n-body symmetries."""
    n_sites = config.shape[0]
    n_symmetries = sym_op_diags.shape[0]
    for sym_idx in range(n_symmetries):
        site_indices = nsites_list[sym_idx]
        nsites_per_symmetry = len(site_indices)
        if np.any(site_indices >= n_sites):
            continue
        sym_value = 1.0
        for ii in range(nsites_per_symmetry):
            site_index = site_indices[ii]
            op_diag = sym_op_diags[sym_idx, site_index, :]
            sym_value *= op_diag[config[site_index]]
        if not np.isclose(sym_value, sym_sectors[sym_idx], atol=1e-13):
            return False
    return True


@njit(cache=True)
def check_glob_sym_partial(config, sym_op_diags, sym_sectors, sym_type_flag):
    """Reference partial global-symmetry compatibility check."""
    n_sites = config.shape[0]
    num_operators = sym_op_diags.shape[0]
    max_sites = sym_op_diags.shape[1]
    for jj in range(num_operators):
        if sym_type_flag == 0:
            operation_result = 0.0
            for kk in range(n_sites):
                operation_result += sym_op_diags[jj, kk, config[kk]]
            if n_sites < max_sites:
                if operation_result > sym_sectors[jj]:
                    return False
            else:
                if not np.isclose(operation_result, sym_sectors[jj], atol=1e-13):
                    return False
        else:
            operation_result = 1.0
            for kk in range(n_sites):
                operation_result *= sym_op_diags[jj, kk, config[kk]]
            if n_sites == max_sites:
                if not np.isclose(operation_result, sym_sectors[jj], atol=1e-13):
                    return False
    return True
