import math
import hashlib
import json
import copy
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from dataclasses import dataclass, replace, asdict
from typing import Dict, List, Tuple, Callable, Optional
from math import erf

try:
    from pymoo.core.problem import ElementwiseProblem
    from pymoo.algorithms.moo.nsga2 import NSGA2
    from pymoo.optimize import minimize
    from pymoo.operators.sampling.rnd import FloatRandomSampling
    from pymoo.operators.crossover.sbx import SBX
    from pymoo.operators.mutation.pm import PM
    PYMOO_AVAILABLE = True
except ImportError:
    PYMOO_AVAILABLE = False

import math
import hashlib
import numpy as np
from typing import Optional, Callable, Tuple

# ---------------------------------------------------------------------------
# Helpers: integer reciprocal modes + shell grouping
# ---------------------------------------------------------------------------

def _build_integer_modes_sphere(kmax: int) -> np.ndarray:
    """
    Build integer modes m = (mx,my,mz) with 0 < ||m|| <= kmax (Euclidean norm),
    excluding (0,0,0). Returned sorted lexicographically for determinism.
    """
    rng = range(-kmax, kmax + 1)
    modes = []
    k2max = kmax * kmax
    for mx in rng:
        for my in rng:
            for mz in rng:
                if mx == 0 and my == 0 and mz == 0:
                    continue
                r2 = mx * mx + my * my + mz * mz
                if r2 <= k2max:
                    modes.append((mx, my, mz))
    M = np.array(modes, dtype=np.int32)
    # deterministic ordering
    M = M[np.lexsort((M[:, 2], M[:, 1], M[:, 0]))]
    return M


def _group_shells_by_r2(modes: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Group modes by r2 = ||m||^2.

    Returns:
      unique_r2 : (S,) sorted unique r2 values
      mode_shell: (M,) shell index per mode in [0,S)
      shell_counts: (S,) number of modes per shell
    """
    m2 = (modes.astype(np.int64) ** 2).sum(axis=1)  # (M,)
    unique_r2, inv = np.unique(m2, return_inverse=True)
    shell_counts = np.bincount(inv, minlength=unique_r2.shape[0]).astype(np.int32)
    return unique_r2.astype(np.int64), inv.astype(np.int32), shell_counts


def _stable_frac_from_cart(X: np.ndarray, A: np.ndarray) -> np.ndarray:
    """
    Convert Cartesian positions X (N,3) to fractional F (N,3) with numerical stability.
    A is (3,3) lattice with row vectors (consistent with your code).
    """
    # Solve A^T * y = X^T => y^T = X * inv(A)^T (but solve is stabler than inv)
    FinvT = np.linalg.solve(A.T, np.eye(3, dtype=float))
    F = X @ FinvT
    # wrap robustly
    F -= np.floor(F)
    return F


def _qint_rint(x: np.ndarray, grid: float, dtype=np.int32) -> np.ndarray:
    """Quantize x by grid using round-to-nearest (robust) and return integer array."""
    return np.rint(np.asarray(x, float) / float(grid)).astype(dtype, copy=False)


def _hash_bytes_sha256(*chunks: bytes) -> str:
    h = hashlib.sha256()
    for c in chunks:
        h.update(c)
    return h.hexdigest()


# ---------------------------------------------------------------------------
#  Robust, rotation-invariant TSF hash (periodic)
#    - numerical stability
#    - cell choice invariance (metric tensor)
#    - supercell/primitive equivalence (spglib reduction optional, recommended)
#    - small relaxation tolerance (log1p + quantization)
#    - degeneracy reduction (power spectrum + weights + shell counts)
#    - performance (no loop over modes; vectorized shell aggregation)
#    - hashing efficiency (hash large blocks as digests)
# ---------------------------------------------------------------------------

def structure_factor_hash_factory_rotinv_robust(
    *,
    # Reciprocal sampling
    kmax: int = 4,
    modes: Optional[np.ndarray] = None,       # custom (M,3) integer modes; overrides kmax when given
    per_species: bool = True,                 # channels per Z (True) or aggregated (False)

    # Quantization / robustness knobs
    ps_grid: float = 5e-2,                    # grid after log1p power-shell (dimensionless)
    metric_grid: float = 5e-2,                # Å^2 grid for metric tensor G = A^T A
    e_grid: float = 2e-3,                     # eV/atom tick (optional)
    v_grid: float = 2e-2,                     # Å^3/atom tick (optional)
    include_energy: bool = True,
    include_volume: bool = True,

    # Canonicalization for supercell/primitive equivalence (recommended)
    use_spglib: bool = True,
    symprec: float = 2e-3,
    angle_tolerance: float = -1.0,
    to_primitive: bool = True,

    # Numerical stability / noise reduction
    center_fractionals: bool = True,          # remove average fractional shift (origin robustness)
    power_spectrum: bool = True,              # use |S|^2 (recommended)
    low_k_weight: bool = True,                # weight shells by 1/(1+r2) (recommended)
    normalize_mode: str = "sqrtN",            # "sqrtN" or "sqrtNz"

    # Performance
    chunk_size: int = 20000,
    debug: bool = False,
) -> Callable[[object], str]:
    """
    Rotationally invariant structure-factor hash (periodic), robust version.

    Key changes vs your baseline:
      - Cell invariance: hash metric tensor G = A^T A (quantized) instead of raw A.
      - Primitive equivalence: optional spglib.standardize_cell(to_primitive=True) by default.
      - Noise robustness: power spectrum (|S|^2), log1p compression, quantization.
      - Degeneracy reduction: include shell_counts, optional low-k weighting.
      - Performance: vectorized S(m) accumulation without looping over modes; shell aggregation via bincount.

    Assumes container.AtomPositionManager provides:
      - atomPositions (N,3) Cartesian Å
      - latticeVectors (3,3) with column vectors (a,b,c)  [your original assumption]
      - get_atomic_numbers() -> (N,)
      - optional .pbc (tuple/list of 3 bools), optional .E (total energy)
    """
    TWO_PI = 2.0 * math.pi
    EPS = 1e-12

    # Modes + shells
    if modes is None:
        M = _build_integer_modes_sphere(int(kmax))
    else:
        M = np.asarray(modes, dtype=np.int32)
        if M.ndim != 2 or M.shape[1] != 3 or M.shape[0] == 0:
            raise ValueError("modes must have shape (M,3) with M>0")
        M = M[np.lexsort((M[:, 2], M[:, 1], M[:, 0]))]

    unique_r2, mode_shell, shell_counts = _group_shells_by_r2(M)
    # Ensure we exclude r2==0 if present (should not happen)
    if unique_r2.size > 0 and unique_r2[0] == 0:
        mask = unique_r2 != 0
        # remap shells
        keep_shell_ids = np.nonzero(mask)[0]
        old_to_new = -np.ones(unique_r2.shape[0], dtype=np.int32)
        old_to_new[keep_shell_ids] = np.arange(keep_shell_ids.size, dtype=np.int32)
        mode_shell = old_to_new[mode_shell]
        good = mode_shell >= 0
        M = M[good]
        mode_shell = mode_shell[good]
        unique_r2 = unique_r2[mask]
        shell_counts = shell_counts[mask]

    # deterministic weights by shell
    if low_k_weight:
        w_shell = 1.0 / (1.0 + unique_r2.astype(np.float64))
    else:
        w_shell = np.ones_like(unique_r2, dtype=np.float64)

    M_T = M.T  # (3,M)
    M_count = int(M.shape[0])
    S_count = int(unique_r2.shape[0])

    # Precompute per-mode weights derived from shell weights (vectorized)
    w_mode = w_shell[mode_shell]  # (M,)

    # Hashing efficiency: hash big blocks as digests, then hash digests together
    def _digest(b: bytes) -> bytes:
        return hashlib.sha256(b).digest()

    def hash_fn(container) -> str:
        apm = container.AtomPositionManager

        # PBC detection
        try:
            pbc_flags = getattr(apm, "pbc", None)
            is_nonpbc = isinstance(pbc_flags, (tuple, list)) and not any(pbc_flags)
        except Exception:
            is_nonpbc = False
        if not isinstance(is_nonpbc, bool):
            is_nonpbc = False

        Z = np.asarray(apm.get_atomic_numbers(), np.int32)
        N = int(Z.shape[0])
        if N == 0:
            return _hash_bytes_sha256(b"TSFv3_rotinv_robust\0", np.array([0], dtype="<i8").tobytes())

        # Optional ticks
        try:
            E_total = float(getattr(apm, "E", 0.0))
        except Exception:
            E_total = 0.0
        e_tick = -1 if not include_energy else int(np.rint((E_total / N) / float(e_grid)))

        if is_nonpbc:
            # If you want true rotation invariance for molecules, use a distance-spectrum style hash;
            # here we keep a stable fallback but note it is not diffraction-based.
            X = np.asarray(apm.atomPositions, float)
            Xc = X - X.mean(axis=0)
            # quantize centered coordinates (robust-ish)
            Xq = _qint_rint(Xc, math.sqrt(metric_grid))  # use a length-like grid
            K = np.column_stack((Z, Xq)).astype(np.int32)
            Kuniq, counts = np.unique(K, axis=0, return_counts=True)
            meta = np.array([N, e_tick, -1], dtype="<i8").tobytes()
            return _hash_bytes_sha256(
                b"TSFv3_rotinv_robust\0",
                _digest(np.asarray(Kuniq, dtype="<i4").tobytes()),
                _digest(np.asarray(counts, dtype="<i4").tobytes()),
                meta,
            )

        # Periodic path
        A_cols = np.asarray(apm.latticeVectors, float)  # columns a,b,c
        A = A_cols.T                                   # make rows for math, consistent with your baseline
        X = np.asarray(apm.atomPositions, float)

        # Optional spglib reduction for primitive equivalence / canonical setting
        if use_spglib:
            try:
                import spglib
                # spglib expects (lattice, positions_frac, numbers) with lattice rows (3,3)
                F0 = _stable_frac_from_cart(X, A)
                A_s, F_s, Z_s = spglib.standardize_cell(
                    (A.copy(), F0.copy(), Z.copy()),
                    to_primitive=bool(to_primitive),
                    no_idealize=True,
                    symprec=float(symprec),
                    angle_tolerance=float(angle_tolerance),
                )
                A = np.asarray(A_s, float)
                F = np.asarray(F_s, float)
                F -= np.floor(F)  # wrap
                Z = np.asarray(Z_s, np.int32)
                N = int(Z.shape[0])
                if N == 0:
                    return _hash_bytes_sha256(b"TSFv3_rotinv_robust\0", np.array([0], dtype="<i8").tobytes())
            except Exception as exc:
                if debug:
                    print("spglib standardize_cell failed; using unreduced cell:", exc)
                F = _stable_frac_from_cart(X, A)
        else:
            F = _stable_frac_from_cart(X, A)

        if center_fractionals:
            # improves robustness to origin choice / numeric jitter near boundaries
            F = (F - F.mean(axis=0))
            F -= np.floor(F)

        # Species channels
        if per_species:
            species, inv = np.unique(Z, return_inverse=True)  # sorted
            C = int(species.shape[0])
        else:
            species = np.array([0], dtype=np.int32)
            inv = np.zeros(N, dtype=np.int32)
            C = 1

        # Metric tensor (cell-choice invariance)
        # G = A^T A is invariant to rotation/permutation of lattice basis (up to numerical noise)
        G = A.T @ A
        G_q = _qint_rint(G, metric_grid, dtype=np.int32)

        # Volume per atom (optional)
        v_tick = -1
        if include_volume:
            vpa = abs(np.linalg.det(A)) / max(N, 1)
            v_tick = int(np.rint(vpa / float(v_grid)))

        # Structure factor accumulation (vectorized; no Python loop over modes)
        # We compute per chunk:
        #   E = exp(i 2π F @ M_T)  -> (chunk, M)
        # then sum E by species channel.
        S_re = np.zeros((C, M_count), dtype=np.float64)
        S_im = np.zeros((C, M_count), dtype=np.float64)

        # Precompute indices per channel once (faster, deterministic)
        # If C is small (typical), looping over channels is cheap and avoids massive one-hot matrices.
        idx_by_c = [np.nonzero(inv == c)[0] for c in range(C)]

        for s in range(0, N, max(1, int(chunk_size))):
            e = min(N, s + int(chunk_size))
            Fi = F[s:e]  # (chunk,3)

            phase = (TWO_PI * (Fi @ M_T)).astype(np.float64, copy=False)  # (chunk,M)
            E = np.exp(1j * phase)  # complex128

            inv_i = inv[s:e]
            # accumulate per channel without looping over modes
            for c in range(C):
                # positions in this chunk belonging to channel c
                mask = (inv_i == c)
                if not np.any(mask):
                    continue
                Sc = E[mask].sum(axis=0)  # (M,)
                S_re[c] += Sc.real
                S_im[c] += Sc.imag

        # Normalization
        if normalize_mode == "sqrtNz":
            Nz = np.bincount(inv, minlength=C).astype(np.float64)
            denom = np.sqrt(np.maximum(Nz, 4.0))[:, None]  # clamp for robustness
        else:
            denom = np.sqrt(float(max(N, 1.0)))
            denom = np.full((C, 1), denom, dtype=np.float64)

        # Magnitude / power spectrum
        mag = np.hypot(S_re, S_im) / denom  # (C,M)
        if power_spectrum:
            spec = mag * mag  # |S|^2
        else:
            spec = mag

        # Apply low-k weighting per mode (derived from shell weights)
        if low_k_weight:
            spec = spec * w_mode[None, :]

        # Shell aggregation (vectorized via bincount; no Python dict/list indexing)
        # For each channel c:
        #   shell_sum[s] = sum_{m in shell s} spec[c,m]
        #   shell_mean = shell_sum / shell_counts
        shell_counts_f = shell_counts.astype(np.float64)
        spec_shell = np.zeros((C, S_count), dtype=np.float64)

        # Precompute bincount "x" as mode_shell, same for all channels
        ms = mode_shell

        for c in range(C):
            shell_sum = np.bincount(ms, weights=spec[c], minlength=S_count).astype(np.float64, copy=False)
            spec_shell[c] = shell_sum / np.maximum(shell_counts_f, 1.0)

        # Compress dynamic range + quantize
        # log1p makes the descriptor much more tolerant to small relaxations and spikes.
        spec_shell = np.log1p(np.maximum(spec_shell, 0.0))
        Pq = _qint_rint(spec_shell, ps_grid, dtype=np.int32)  # (C,S)

        # Explicit composition improves collision resistance
        # Keep only present species for compactness, still deterministic due to sorted `species`.
        comp = np.bincount(inv, minlength=C).astype(np.int32)

        # Serialize deterministically using digests for large blocks
        meta = np.array(
            [N, int(species.shape[0]), e_tick, v_tick, int(kmax), int(M_count), int(S_count)],
            dtype="<i8"
        ).tobytes()

        payload_parts = (
            b"TSFv3_rotinv_robust\0",
            _digest(np.asarray(G_q, dtype="<i4").tobytes()),
            _digest(np.asarray(species, dtype="<i4").tobytes()),
            _digest(np.asarray(comp, dtype="<i4").tobytes()),
            _digest(np.asarray(unique_r2, dtype="<i8").tobytes()),
            _digest(np.asarray(shell_counts, dtype="<i4").tobytes()),
            _digest(np.asarray(Pq, dtype="<i4").tobytes()),
            meta,
        )
        return _hash_bytes_sha256(*payload_parts)

    return hash_fn

# =============================================================================
#  Rotationally-invariant TSF hash (periodic) with physically-motivated grids
# =============================================================================
#
# Key upgrades vs your current version
# -----------------------------------
# 1) Shells are defined in *physical reciprocal space* using the reciprocal metric:
#       k(m) = 2π A^{-T} m,   |k|^2 = m^T G* m,  where G* = (2π)^2 A^{-1} A^{-T}
#    This is robust for non-cubic / sheared / anisotropic cells.
#
# 2) Quantization grids are per-species (and per-shell) derived from a positional
#    uncertainty model σ_s(T) ~ α sqrt(T / m_s). This links discretization to
#    mass and temperature in a mathematically defensible way.
#
# 3) Energy tick uses Δe = max(γ k_B T, σ_model_e) (per atom), consistent with
#    thermal + model uncertainty.
#
# 4) Optional extra invariant moments per shell (mean(|S|) and mean(|S|^2))
#    to reduce collisions while keeping rotational invariance.
#
# 5) Non-PBC fallback uses a sampled pair-distance spectrum (rotation/translation
#    invariant) with physically-motivated bin width.
#
# Notes
# -----
# - This is a hashing descriptor, not a metric. You should choose α, β, γ to
#   trade off stability (same structure under noise) vs discrimination.
# - You can later "learn" α, β, γ and per-species correction factors by
#   optimizing collision rates on curated positive/negative pairs.
# =============================================================================

# Reusa lo anterior:
# - atomic_mass_u
# - _build_integer_modes_sphere
# - _reciprocal_metric
# - _compute_k2_for_modes
# - _bin_shells_by_k2
# - _pair_distance_signature
# - _qint
# - _sha256_bytes
# - _energy_tick
# - _volume_tick
# - K_B_eV_per_K, TWO_PI, EPS

# ----------------------------- Physical constants -----------------------------
K_B_eV_per_K = 8.617333262145e-5  # Boltzmann constant in eV/K
TWO_PI = 2.0 * math.pi
EPS = 1e-12


# ----------------------------- Mass table (lightweight) -----------------------
# If you already have a mass provider, plug it in here.
# Fallback: approximate atomic masses for common Z. Extend as needed.
_ATOMIC_MASS_U: Dict[int, float] = {
    1: 1.00784,  2: 4.002602,
    3: 6.94,     4: 9.0121831, 5: 10.81,  6: 12.011,  7: 14.007,  8: 15.999,
    9: 18.998403163, 10: 20.1797,
    11: 22.98976928, 12: 24.305, 13: 26.9815385, 14: 28.085, 15: 30.973761998,
    16: 32.06, 17: 35.45, 18: 39.948,
    19: 39.0983, 20: 40.078,
    21: 44.955908, 22: 47.867, 23: 50.9415, 24: 51.9961, 25: 54.938044,
    26: 55.845, 27: 58.933194, 28: 58.6934, 29: 63.546, 30: 65.38,
    31: 69.723, 32: 72.630, 33: 74.921595, 34: 78.971, 35: 79.904, 36: 83.798,
    37: 85.4678, 38: 87.62, 39: 88.90584, 40: 91.224, 41: 92.90637, 42: 95.95,
    43: 98.0, 44: 101.07, 45: 102.90550, 46: 106.42, 47: 107.8682, 48: 112.414,
    49: 114.818, 50: 118.710, 51: 121.760, 52: 127.60, 53: 126.90447, 54: 131.293,
    55: 132.90545196, 56: 137.327,
    57: 138.90547, 58: 140.116, 59: 140.90766, 60: 144.242,
    61: 145.0, 62: 150.36, 63: 151.964, 64: 157.25, 65: 158.92535, 66: 162.500,
    67: 164.93033, 68: 167.259, 69: 168.93422, 70: 173.045, 71: 174.9668,
    72: 178.49, 73: 180.94788, 74: 183.84, 75: 186.207, 76: 190.23, 77: 192.217,
    78: 195.084, 79: 196.966569, 80: 200.592,
}

def atomic_mass_u(Z: int) -> float:
    return float(_ATOMIC_MASS_U.get(int(Z), 50.0))

# ----------------------------- Mode construction ------------------------------
def _build_integer_modes_sphere(kmax: int) -> np.ndarray:
    """
    Build integer modes m = (mx,my,mz) in a cubic box [-kmax..kmax]^3 excluding (0,0,0),
    then keep those with ||m|| <= kmax (Euclidean).
    """
    ks = np.arange(-kmax, kmax + 1, dtype=np.int32)
    mx, my, mz = np.meshgrid(ks, ks, ks, indexing="ij")
    M = np.stack([mx.ravel(), my.ravel(), mz.ravel()], axis=1).astype(np.int32, copy=False)
    # remove zero
    nz = np.any(M != 0, axis=1)
    M = M[nz]
    # sphere cut
    r2 = np.sum(M.astype(np.int64) ** 2, axis=1)
    M = M[r2 <= int(kmax * kmax)]
    # deterministic ordering
    M = M[np.lexsort((M[:, 2], M[:, 1], M[:, 0]))]
    return M


# ----------------------------- Reciprocal shells ------------------------------
def _reciprocal_metric(A: np.ndarray) -> np.ndarray:
    """
    Reciprocal metric G* = (2π)^2 A^{-1} A^{-T}, where A is the real-space cell matrix
    with rows as lattice vectors in Cartesian (Å).

    If you store lattice vectors as columns elsewhere, transpose before calling.
    """
    Ainv = np.linalg.inv(A)
    Gstar = (TWO_PI ** 2) * (Ainv @ Ainv.T)
    return Gstar


def _compute_k2_for_modes(M: np.ndarray, Gstar: np.ndarray) -> np.ndarray:
    """
    Compute |k(m)|^2 = m^T G* m for each integer mode m.
    """
    # k2_i = sum_{a,b} m_a G*_ab m_b
    # vectorized: (M @ Gstar) * M then sum row
    MG = M.astype(np.float64) @ Gstar
    k2 = np.einsum("ij,ij->i", MG, M.astype(np.float64))
    return k2


def _bin_shells_by_k2(k2: np.ndarray, k2_grid: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Bin modes into shells by quantized k^2:
        shell_id = round(k2 / k2_grid)
    Returns:
      - unique_shell_ids (int64, sorted)
      - shell_ptr: indices into sorted modes (CSR-like pointer)
      - order: sorting indices that sort by shell_id
    """
    if k2_grid <= 0:
        raise ValueError("k2_grid must be > 0")

    shell_id = np.floor(k2 / k2_grid + 0.5).astype(np.int64)
    order = np.argsort(shell_id, kind="mergesort")  # stable
    sid_sorted = shell_id[order]
    unique_sid, first_idx = np.unique(sid_sorted, return_index=True)
    # build pointer array: start positions of each shell + end
    shell_ptr = np.concatenate([first_idx, np.array([sid_sorted.size], dtype=first_idx.dtype)])
    return unique_sid, shell_ptr, order


def _generate_power_law_bins(kmax: float, p: float, n_bins: int) -> np.ndarray:
    """Generate power-law bins k_i = kmax * (i/n_bins)^p."""
    bins = float(kmax) * (np.linspace(0, 1, int(n_bins) + 1) ** float(p))
    bins[0] = 0.05 # Ensure first bin starts > 0 for stability
    return bins


# ----------------------------- Non-PBC fallback -------------------------------
def _pair_distance_signature(
    X: np.ndarray,
    Z: np.ndarray,
    *,
    per_species: bool,
    bin_width: float,
    max_pairs: int = 200000,
    rng_seed: int = 0,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Rotation/translation invariant signature for non-PBC:
      - sample pair distances (i<j),
      - histogram (optionally per species-pair channels).

    Returns (channels, bins) as int32 payload plus channel ids.
    """
    N = int(X.shape[0])
    if N < 2:
        return np.zeros((1, 1), dtype=np.int32), np.array([0], dtype=np.int32)

    if bin_width <= 0:
        raise ValueError("bin_width must be > 0")

    # center (translation invariance)
    Xc = X - X.mean(axis=0, keepdims=True)

    # choose pairs (avoid O(N^2) for large N)
    total_pairs = N * (N - 1) // 2
    K = min(int(max_pairs), int(total_pairs))

    rng = np.random.default_rng(int(rng_seed))
    # sample indices i<j by rejection sampling (simple, adequate for hashing)
    ii = rng.integers(0, N, size=K, endpoint=False)
    jj = rng.integers(0, N, size=K, endpoint=False)
    mask = ii != jj
    ii, jj = ii[mask], jj[mask]
    # enforce i<j
    lo = np.minimum(ii, jj)
    hi = np.maximum(ii, jj)
    ii, jj = lo, hi

    # distances
    d = np.linalg.norm(Xc[ii] - Xc[jj], axis=1)
    if d.size == 0:
        return np.zeros((1, 1), dtype=np.int32), np.array([0], dtype=np.int32)

    # bins
    dmax = float(np.max(d))
    nbins = max(4, int(math.ceil(dmax / bin_width)))
    edges = np.linspace(0.0, nbins * bin_width, nbins + 1)

    if not per_species:
        hist, _ = np.histogram(d, bins=edges)
        return hist.astype(np.int32)[None, :], np.array([0], dtype=np.int32)

    # per species-pair channels (Z_i<=Z_j canonical)
    Zi = Z[ii].astype(np.int32)
    Zj = Z[jj].astype(np.int32)
    Za = np.minimum(Zi, Zj)
    Zb = np.maximum(Zi, Zj)
    pairs = np.stack([Za, Zb], axis=1)

    uniq_pairs, inv = np.unique(pairs, axis=0, return_inverse=True)
    C = int(uniq_pairs.shape[0])
    H = np.zeros((C, nbins), dtype=np.int32)
    # fill each channel
    for c in range(C):
        dc = d[inv == c]
        if dc.size:
            hc, _ = np.histogram(dc, bins=edges)
            H[c, :] = hc.astype(np.int32, copy=False)

    # channel ids flatten (Za, Zb) pairs in payload
    chan_ids = uniq_pairs.astype(np.int32, copy=False).ravel()
    return H, chan_ids


# ----------------------------- Quantization utilities -------------------------
def _qint(x: np.ndarray, grid: float) -> np.ndarray:
    """Quantize x to nearest integer tick with step 'grid'."""
    return np.floor(np.asarray(x, float) / float(grid) + 0.5).astype(np.int32, copy=False)


def _sha256_bytes(*chunks: bytes) -> str:
    h = hashlib.sha256()
    for c in chunks:
        h.update(c)
    return h.hexdigest()



def _volume_tick(A: np.ndarray, N: int, v_grid: float) -> int:
    if v_grid <= 0:
        return -1
    vpa = abs(float(np.linalg.det(A))) / max(int(N), 1)
    return int(math.floor(vpa / float(v_grid) + 0.5))


# =============================================================================
#  1) Params extendidos: species corrections + shell band weights
# =============================================================================

@dataclass(frozen=True)
class TSFPhysParamsLearned:
    """
    Extiende TSFPhysParams con:
      - species_correction: dict Z->c_s  (multiplica sigma_s)
      - shell_band_edges: cortes en |k| (Å^-1) para asignar shell -> banda
      - shell_band_weights: pesos w_b por banda (divide ps_grid)
    """
    # Base
    T_K: float = 300.0
    alpha_A: float = 0.02
    beta: float = 2.0
    gamma: float = 0.5
    sigma_model_e_eV_per_atom: float = 0.0

    pair_bin_beta: float = 2.0
    min_ps_grid: float = 1e-4
    max_ps_grid: float = 100.0  # Increased to allow for very stable basin classification

    # Learned / optional
    species_correction: Dict[int, float] = None  # {Z: c_s}, default None => 1
    shell_band_edges: Tuple[float, ...] = (2.0, 5.0, 10.0)  # Å^-1 (example)
    shell_band_weights: Tuple[float, ...] = (1.0, 1.0, 1.0, 1.0)  # len = len(edges)+1
    
    kmax: float = 1.5
    power_law_exponent: float = 1.0 # for grid spacing
    B_factor: float = 0.0           # Debye-Waller: exp(-B*k^2)
    eta: float = 1.5                # k-weighting exponent

    def __post_init__(self):
        # dataclass frozen: cannot mutate; validation outside if needed
        pass

    def to_dict(self) -> dict:
        """Serialize params to a dictionary."""
        d = asdict(self)
        # Ensure tuples/lists are handled if any
        return d

    @classmethod
    def from_dict(cls, data: dict) -> "TSFPhysParamsLearned":
        """Reconstruct params from a dictionary."""
        # Convert dictionary keys back to integers for species_correction if needed
        if "species_correction" in data and data["species_correction"] is not None:
            data["species_correction"] = {int(k): float(v) for k, v in data["species_correction"].items()}
        
        if "shell_band_edges" in data and isinstance(data["shell_band_edges"], list):
            data["shell_band_edges"] = tuple(data["shell_band_edges"])
        
        if "shell_band_weights" in data and isinstance(data["shell_band_weights"], list):
            data["shell_band_weights"] = tuple(data["shell_band_weights"])
            
        return cls(**data)

    def save_json(self, path: str):
        """Save params to a JSON file."""
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=4)

    @classmethod
    def load_json(cls, path: str) -> "TSFPhysParamsLearned":
        """Load params from a JSON file."""
        with open(path, "r") as f:
            data = json.load(f)
        return cls.from_dict(data)


# ----------------------------- Perturbation utilities -------------------------

def thermal_perturb(structure: object, alpha: float = 0.02, T: float = 300.0) -> object:
    """
    Apply mass-dependent thermal vibrations: sigma_s = alpha * sqrt(T/m_s).
    """
    s = copy.deepcopy(structure)
    apm = s.AtomPositionManager
    Z = np.asarray(apm.get_atomic_numbers(), dtype=np.int32)
    X = np.asarray(apm.atomPositions, dtype=np.float64)
    
    masses = np.array([atomic_mass_u(int(z)) for z in Z], dtype=np.float64)
    sigmas = alpha * np.sqrt(max(T, 1e-9) / masses)
    
    noise = np.random.normal(scale=1.0, size=X.shape) * sigmas[:, None]
    
    # Check if set_atomPositions exists, otherwise set attribute directly
    if hasattr(apm, "set_atomPositions"):
        apm.set_atomPositions(X + noise)
    else:
        apm.atomPositions = X + noise
        
    return s


def strain_perturb(structure: object, max_strain: float = 0.02) -> object:
    """
    Apply a small random lattice strain.
    """
    s = copy.deepcopy(structure)
    apm = s.AtomPositionManager
    
    # Real-space lattice A (vectors as columns)
    A = np.asarray(apm.latticeVectors, dtype=np.float64)
    
    # Random strain matrix epsilon
    eps = np.random.normal(scale=max_strain, size=(3, 3))
    # Symmetric part for pure strain
    eps = 0.5 * (eps + eps.T)
    
    # New lattice A' = (I + eps) A
    A_new = (np.eye(3) + eps) @ A
    
    if hasattr(apm, "set_latticeVectors"):
        apm.set_latticeVectors(A_new)
    else:
        apm.latticeVectors = A_new
        
    # Also transform fractional positions if they exist, 
    # but here we assume atomPositions are in Cartesian (Å).
    # If they were fractional, they'd stay same. 
    # If they are Cartesian, we should ideally strain them too?
    # Usually, a volume expansion/strain scales the atoms.
    X = np.asarray(apm.atomPositions, dtype=np.float64)
    X_new = X @ (np.eye(3) + eps).T
    
    if hasattr(apm, "set_atomPositions"):
        apm.set_atomPositions(X_new)
    else:
        apm.atomPositions = X_new
        
    return s


def random_perturb(structure: object, alpha: float = 0.01, T: float = 300.0, max_strain: float = 0.0) -> object:
    """
    Combine thermal vibration and lattice strain.
    """
    s = thermal_perturb(structure, alpha=alpha, T=T)
    s = strain_perturb(s, max_strain=max_strain)
    return s


def swap_perturb(structure: object, n_swaps: int = 1) -> object:
    """
    Randomly swap species of n pairs of atoms.
    """
    s = copy.deepcopy(structure)
    apm = s.AtomPositionManager
    Z = np.array(apm.get_atomic_numbers(), dtype=np.int32)
    
    indices = np.arange(len(Z))
    for _ in range(n_swaps):
        i, j = np.random.choice(indices, size=2, replace=False)
        Z[i], Z[j] = Z[j], Z[i]
        
    # Set back
    if hasattr(apm, "_Z"):
        apm._Z = Z
    return s


def diffuse_perturb(structure: object, rms_dist: float = 1.0) -> object:
    """
    Large random displacements (diffusion/melt) to ensure distinct structure.
    """
    s = copy.deepcopy(structure)
    apm = s.AtomPositionManager
    X = np.asarray(apm.atomPositions, dtype=np.float64)
    noise = np.random.normal(scale=rms_dist, size=X.shape)
    
    if hasattr(apm, "set_atomPositions"):
        apm.set_atomPositions(X + noise)
    else:
        apm.atomPositions = X + noise
    return s


def _species_sigma_A_learned(species_Z: np.ndarray, params: TSFPhysParamsLearned) -> np.ndarray:
    """
    σ_s(T) [Å] per species from mass + temperature, with learned correction c_s:
        σ_s = alpha * sqrt(T/m_s) * c_s
    """
    T = max(float(params.T_K), 1e-9)
    masses = np.array([atomic_mass_u(int(z)) for z in species_Z], dtype=np.float64)
    base = float(params.alpha_A) * np.sqrt(T / masses)

    corr = np.ones_like(base)
    if params.species_correction:
        for i, z in enumerate(species_Z):
            corr[i] = float(params.species_correction.get(int(z), 1.0))
    return base * corr


def _shell_band_index(k_abs_shell: np.ndarray, edges: Tuple[float, ...]) -> np.ndarray:
    """
    Map each shell to a band index by |k|.
    edges = (e0, e1, ..., e_{K-2}) produce K bands.
    """
    edges_arr = np.array(edges, dtype=np.float64)
    # band = number of edges <= k
    return np.searchsorted(edges_arr, k_abs_shell, side="right").astype(np.int32)


def _energy_tick(E_total: float, N: int, params: TSFPhysParamsLearned) -> Tuple[int, float]:
    """
    Compute energy tick integer and the used grid Δe (eV/atom).
    """
    if N <= 0:
        return -1, 1.0
    e_pa = float(E_total) / float(N)
    de = max(float(params.gamma) * K_B_eV_per_K * float(params.T_K),
             float(params.sigma_model_e_eV_per_atom))
    de = max(de, 1e-12)
    e_tick = int(math.floor(e_pa / de + 0.5))
    return e_tick, de

# =============================================================================
#  2) Hash factory mod: incorporate learned sigma and shell weights
# =============================================================================
def structure_factor_hash_factory_rotinv_phys_learned(
    *,
    kmax: int = 3,
    modes: Optional[np.ndarray] = None,
    k2_grid: float = 0.01,
    fixed_k_bins: Optional[np.ndarray] = None,
    per_species: bool = True,
    lattice_grid_A: float = 2e-1,
    include_energy: bool = True,
    include_volume: bool = True,
    v_grid_A3_per_atom: float = 1e-2,
    include_power_moment: bool = True,
    phys: TSFPhysParamsLearned = TSFPhysParamsLearned(),
    use_spglib: bool = False,
    symprec: float = 1e-3,
    angle_tolerance: float = -1.0,
    chunk_modes: int = 2048,
    nonpbc_max_pairs: int = 200000,
    rng_seed: int = 0,
    debug: bool = False,
) -> Callable[[object], str]:
    """
    Rotationally-invariant TSF hash (periodic + non-PBC fallback), with:
      - per-species sigma (mass + learned correction)
      - optional fixed k-bins
      - optional power moment (mu2)
      - added shell variance: var = <|S|^2> - <|S|>^2
      - k-weighting for high-k shells (eta exponent)
    """

    # Use phys.kmax if kmax is default (3), else override
    curr_kmax = getattr(phys, "kmax", kmax) if kmax == 3 else kmax

    if modes is None:
        M = _build_integer_modes_sphere(int(curr_kmax))
    else:
        M = np.asarray(modes, dtype=np.int32)
        if M.ndim != 2 or M.shape[1] != 3 or M.shape[0] == 0:
            raise ValueError("modes must have shape (M,3) with M>0")
        M = M[np.lexsort((M[:, 2], M[:, 1], M[:, 0]))]

    M_count = int(M.shape[0])

    def hash_fn(container, return_descriptor: bool = False, fixed_k_bins=fixed_k_bins):
        apm = container.AtomPositionManager

        # -----------------------------
        # PBC detection
        # -----------------------------
        try:
            pbc_flags = getattr(apm, "pbc", None)
            is_nonpbc = isinstance(pbc_flags, (tuple, list)) and not any(pbc_flags)
        except Exception:
            is_nonpbc = False
        if not isinstance(is_nonpbc, bool):
            is_nonpbc = False

        Z = np.asarray(apm.get_atomic_numbers(), dtype=np.int32)
        N = int(Z.shape[0])
        if N == 0:
            if return_descriptor:
                return np.zeros(0, dtype=np.float64)
            return _sha256_bytes(b"TSFv4_learn\0", np.array([0], dtype="<i8").tobytes())

        X = np.asarray(apm.atomPositions, dtype=np.float64)

        # -----------------------------
        # energy tick
        # -----------------------------
        try:
            E_total = float(getattr(apm, "E", 0.0))
        except Exception:
            E_total = 0.0

        if include_energy:
            e_tick, e_grid_used = _energy_tick(E_total, N, phys)
        else:
            e_tick, e_grid_used = -1, 0.0

        # -----------------------------
        # Non-PBC fallback: pair-distance signature
        # -----------------------------
        if is_nonpbc:
            if per_species:
                species_Z, _ = np.unique(Z, return_inverse=True)
                sig = _species_sigma_A_learned(species_Z, phys)
                bin_w = float(phys.pair_bin_beta) * float(np.max(sig))
            else:
                bin_w = float(phys.pair_bin_beta) * float(phys.alpha_A) * math.sqrt(max(phys.T_K, 1e-9) / 50.0)
            bin_w = max(bin_w, 1e-4)

            H, chan_ids = _pair_distance_signature(
                X, Z,
                per_species=per_species,
                bin_width=bin_w,
                max_pairs=nonpbc_max_pairs,
                rng_seed=rng_seed,
            )

            if return_descriptor:
                return H.ravel().astype(np.float64)

            meta = np.array([N, int(chan_ids.size), int(e_tick), -1], dtype="<i8").tobytes()
            return _sha256_bytes(
                b"TSFv4_learn_nonpbc\0",
                np.asarray(chan_ids, dtype="<i4").tobytes(),
                np.asarray(H, dtype="<i4").tobytes(),
                np.array([float(bin_w), float(e_grid_used)], dtype="<f8").tobytes(),
                meta,
            )

        # -----------------------------
        # Periodic: build fractional positions
        # -----------------------------
        A_cols = np.asarray(apm.latticeVectors, dtype=np.float64)  # columns (a,b,c)
        A = A_cols.T  # rows for math
        FinvT = np.linalg.inv(A).T
        F = (X @ FinvT + EPS) % 1.0

        # Optional spglib
        if use_spglib:
            try:
                import spglib
                A_s, F_s, Z_s = spglib.standardize_cell(
                    (A.copy(), F.copy(), Z.copy()),
                    to_primitive=True, no_idealize=True,
                    symprec=symprec, angle_tolerance=angle_tolerance,
                )
                A = np.asarray(A_s, dtype=np.float64)
                F = (np.asarray(F_s, dtype=np.float64) + EPS) % 1.0
                Z = np.asarray(Z_s, dtype=np.int32)
                N = int(Z.shape[0])
            except Exception as exc:
                if debug:
                    print("spglib failed:", exc)

        # -----------------------------
        # Species channels
        # -----------------------------
        if per_species:
            species_Z, inv = np.unique(Z, return_inverse=True)
            C = int(species_Z.size)
        else:
            species_Z = np.array([0], dtype=np.int32)
            inv = np.zeros(N, dtype=np.int32)
            C = 1

        N_s = np.bincount(inv, minlength=C).astype(np.float64)
        denom = np.sqrt(np.maximum(N_s, 1.0))[:, None]  # (C,1)

        # -----------------------------
        # Reciprocal metric & mode k^2
        # -----------------------------
        Gstar = _reciprocal_metric(A)
        k2_modes = _compute_k2_for_modes(M, Gstar)  # (M_count,)
        k_abs_modes = np.sqrt(k2_modes)

        # -----------------------------
        # Precompute raw structure factors for all modes
        # -----------------------------
        atom_order = np.argsort(inv, kind="mergesort")
        Fg = F[atom_order]
        invg = inv[atom_order]
        bounds = np.concatenate(([0], np.cumsum(np.bincount(invg, minlength=C))))

        S_complex_raw = np.zeros((C, M_count), dtype=np.complex128)

        for m0 in range(0, M_count, max(1, int(chunk_modes))):
            m1 = min(M_count, m0 + int(chunk_modes))
            Mb = M.T[:, m0:m1]  # (3,B)
            phase_complex = np.exp(1j * (TWO_PI * (Fg @ Mb)))  # (N,B)

            for s in range(C):
                a = int(bounds[s]); b = int(bounds[s + 1])
                if b > a:
                    S_complex_raw[s, m0:m1] = np.sum(phase_complex[a:b, :], axis=0)

        # -----------------------------
        # Learned sigma
        # -----------------------------
        sigma_s = _species_sigma_A_learned(species_Z, phys)  # (C,)

        # -----------------------------
        # Shell binning: fixed_k_bins OR dynamic (k2_grid)
        # -----------------------------
        if fixed_k_bins is not None:
            fixed_k_bins = np.asarray(fixed_k_bins, dtype=np.float64)
            if fixed_k_bins.ndim != 1 or fixed_k_bins.size < 2:
                raise ValueError("fixed_k_bins must be 1D with length >= 2")
            if not np.all(np.diff(fixed_k_bins) > 0):
                raise ValueError("fixed_k_bins must be strictly increasing")

            shell_idx = np.searchsorted(fixed_k_bins, k_abs_modes, side="right") - 1
            valid = (shell_idx >= 0) & (shell_idx < (fixed_k_bins.size - 1))

            n_shells = int(fixed_k_bins.size - 1)
            S_shells = n_shells

            mu1 = np.zeros((C, n_shells), dtype=np.float64)
            mu2 = np.zeros((C, n_shells), dtype=np.float64) if include_power_moment else None
            shell_counts = np.zeros(n_shells, dtype=np.int64)
            k_abs_shell = np.zeros(n_shells, dtype=np.float64)

            B_factor = float(getattr(phys, "B_factor", 0.0))

            for sh in range(n_shells):
                mask = valid & (shell_idx == sh)
                if not np.any(mask):
                    continue

                S_shell = S_complex_raw[:, mask]
                k_abs_m = k_abs_modes[mask]

                # Debye-Waller damping (optional)
                if B_factor > 0.0:
                    damping = np.exp(-B_factor * (k_abs_m ** 2))
                    S_shell = S_shell * damping[None, :]

                P = np.abs(S_shell) / denom  # (C, n_modes_in_shell)

                mu1[:, sh] = np.mean(P, axis=1)
                if mu2 is not None:
                    mu2[:, sh] = np.mean(P * P, axis=1)

                shell_counts[sh] = int(np.sum(mask))
                k_abs_shell[sh] = float(np.mean(k_abs_m))

            shell_ids = np.arange(S_shells, dtype=np.int64)

        else:
            unique_sid, shell_ptr, order = _bin_shells_by_k2(k2_modes, float(k2_grid))

            k2_sorted = k2_modes[order]
            k_abs_sorted = np.sqrt(k2_sorted)

            S_shells = int(unique_sid.size)

            mu1 = np.zeros((C, S_shells), dtype=np.float64)
            mu2 = np.zeros((C, S_shells), dtype=np.float64) if include_power_moment else None
            shell_counts = np.zeros(S_shells, dtype=np.int64)
            k_abs_shell = np.zeros(S_shells, dtype=np.float64)

            B_factor = float(getattr(phys, "B_factor", 0.0))

            for sh in range(S_shells):
                sa = int(shell_ptr[sh]); sb = int(shell_ptr[sh + 1])
                if sb <= sa:
                    continue

                mode_idx = order[sa:sb]  # original indices
                S_shell = S_complex_raw[:, mode_idx]
                k2_sh = k2_sorted[sa:sb]
                k_abs_sh = k_abs_sorted[sa:sb]

                if B_factor > 0.0:
                    damping = np.exp(-B_factor * k2_sh)
                    S_shell = S_shell * damping[None, :]

                P = np.abs(S_shell) / denom
                mu1[:, sh] = np.mean(P, axis=1)
                if mu2 is not None:
                    mu2[:, sh] = np.mean(P * P, axis=1)

                shell_counts[sh] = (sb - sa)
                k_abs_shell[sh] = float(np.mean(k_abs_sh))

            shell_ids = unique_sid.astype(np.int64)

        # -----------------------------
        # Band weights (coarse)
        # -----------------------------
        band_idx = _shell_band_index(k_abs_shell, phys.shell_band_edges)
        w = np.asarray(phys.shell_band_weights, dtype=np.float64)
        if w.size != (len(phys.shell_band_edges) + 1):
            raise ValueError("shell_band_weights must have len = len(shell_band_edges)+1")
        w_shell = w[band_idx]  # (S_shells,)

        # -----------------------------
        # Quantization grid for mu1/mu2
        # -----------------------------
        ps_grid = (float(phys.beta) * np.sqrt(np.maximum(N_s, 1.0))[:, None] *
                   k_abs_shell[None, :] *
                   sigma_s[:, None]) / np.maximum(w_shell[None, :], 1e-12)
        ps_grid = np.clip(ps_grid, float(phys.min_ps_grid), float(phys.max_ps_grid))

        # -----------------------------
        # Shell variance and k-weighting
        # -----------------------------
        if include_power_moment and mu2 is not None:
            var_shell = np.maximum(mu2 - mu1 * mu1, 0.0)
        else:
            var_shell = None

        eta = float(getattr(phys, "eta", 0.0))
        if eta != 0.0:
            k_weight = (k_abs_shell + 1e-12) ** eta
            k_weight = k_weight[None, :]

            mu1 = mu1 * k_weight
            if include_power_moment and mu2 is not None:
                mu2 = mu2 * (k_weight ** 2)
                var_shell = np.maximum(mu2 - mu1 * mu1, 0.0)

        # -----------------------------
        # Return raw descriptor (optional)
        # -----------------------------
        if return_descriptor:
        
            parts = [mu1.ravel()]
        
            if include_power_moment and mu2 is not None:
                parts.append(mu2.ravel())
        
            if var_shell is not None:
                parts.append(var_shell.ravel())
        
            descriptor = np.concatenate(parts).astype(np.float64, copy=False)
        
            # width estimate from quantization grid
            width = ps_grid.ravel()
        
            if include_power_moment and mu2 is not None:
                width = np.concatenate([width, ps_grid.ravel()])
        
            if var_shell is not None:
                width = np.concatenate([width, ps_grid.ravel()])
        
            return descriptor, width

        # -----------------------------
        # Quantization
        # -----------------------------
        Pq1 = np.floor(mu1 / ps_grid + 0.5)
        Pq1 = np.clip(Pq1, -2e9, 2e9).astype(np.int32)

        Pq2 = None
        Pq_var = None

        if include_power_moment and mu2 is not None:
            ps2_grid = np.clip(ps_grid * ps_grid, phys.min_ps_grid**2, phys.max_ps_grid**2)

            Pq2 = np.floor(mu2 / ps2_grid + 0.5)
            Pq2 = np.clip(Pq2, -2e9, 2e9).astype(np.int32)

            if var_shell is not None:
                Pq_v = np.floor(var_shell / ps2_grid + 0.5)
                Pq_var = np.clip(Pq_v, -2e9, 2e9).astype(np.int32)

        # -----------------------------
        # Lattice / volume ticks
        # -----------------------------
        A_q = _qint(A, float(lattice_grid_A))
        v_tick = _volume_tick(A, N, float(v_grid_A3_per_atom)) if include_volume else -1

        meta_i8 = np.array(
            [int(N), int(C), int(e_tick), int(v_tick),
             int(curr_kmax), int(M_count), int(S_shells), int(include_power_moment)],
            dtype="<i8"
        ).tobytes()

        meta_f8 = np.array(
            [float(phys.T_K), float(phys.alpha_A), float(phys.beta), float(phys.gamma),
             float(phys.sigma_model_e_eV_per_atom), float(k2_grid),
             float(lattice_grid_A), float(v_grid_A3_per_atom), float(e_grid_used),
             float(getattr(phys, "B_factor", 0.0)), float(getattr(phys, "eta", 0.0))],
            dtype="<f8"
        ).tobytes()

        shell_ids_i8 = np.asarray(shell_ids, dtype="<i8").tobytes()
        sigma_f8 = np.asarray(sigma_s, dtype="<f8").tobytes()
        kabs_f8 = np.asarray(k_abs_shell, dtype="<f8").tobytes()
        bands_f8 = np.asarray(w, dtype="<f8").tobytes()

        chunks = [
            b"TSFv4_learn_pbc\0",
            np.asarray(A_q, dtype="<i4").tobytes(),
            np.asarray(species_Z, dtype="<i4").tobytes(),
            shell_ids_i8,
            np.asarray(Pq1, dtype="<i4").tobytes(),
        ]

        # Insert optional arrays in a deterministic order
        if Pq2 is not None:
            chunks.append(np.asarray(Pq2, dtype="<i4").tobytes())
        if Pq_var is not None:
            chunks.append(np.asarray(Pq_var, dtype="<i4").tobytes())

        chunks += [
            meta_i8,
            meta_f8,
            sigma_f8,
            kabs_f8,
            bands_f8,
        ]

        return _sha256_bytes(*chunks)

    return hash_fn


# =============================================================================
#  3) Bayesian Optimization (GP + Expected Improvement) in log-space
# =============================================================================

def _rbf_kernel(X1: np.ndarray, X2: np.ndarray, lengthscales: np.ndarray, var: float) -> np.ndarray:
    """
    ARD RBF kernel:
      k(x,x') = var * exp(-0.5 * sum_d ((x_d-x'_d)^2 / l_d^2))
    """
    X1 = np.atleast_2d(X1)
    X2 = np.atleast_2d(X2)
    ls = np.asarray(lengthscales, dtype=np.float64)
    dif = (X1[:, None, :] - X2[None, :, :]) / ls[None, None, :]
    r2 = np.sum(dif * dif, axis=-1)
    return float(var) * np.exp(-0.5 * r2)


def _gp_posterior(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    *,
    lengthscales: np.ndarray,
    var: float,
    noise: float,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Standard GP posterior mean and variance (diagonal).
    """
    K = _rbf_kernel(X_train, X_train, lengthscales, var)
    K[np.diag_indices_from(K)] += float(noise)

    # Cholesky
    L = np.linalg.cholesky(K + 1e-12 * np.eye(K.shape[0]))
    alpha = np.linalg.solve(L.T, np.linalg.solve(L, y_train))

    K_s = _rbf_kernel(X_train, X_test, lengthscales, var)  # (n, m)
    mu = K_s.T @ alpha

    v = np.linalg.solve(L, K_s)  # (n, m)
    K_ss_diag = np.full(X_test.shape[0], float(var), dtype=np.float64)  # because k(x,x)=var
    var_post = np.maximum(K_ss_diag - np.sum(v * v, axis=0), 1e-12)

    return mu, var_post


def _expected_improvement(mu: np.ndarray, var: np.ndarray, best: float, xi: float = 0.01) -> np.ndarray:
    """
    EI for maximization: E[max(0, f - best - xi)]
    """
    sigma = np.sqrt(np.maximum(var, 1e-12))
    imp = mu - best - float(xi)
    Z = imp / sigma

    # normal pdf/cdf
    pdf = (1.0 / np.sqrt(2.0 * np.pi)) * np.exp(-0.5 * Z * Z)
    #cdf = 0.5 * (1.0 + np.erf(Z / np.sqrt(2.0)))
    cdf = 0.5 * (1.0 + np.vectorize(erf)(Z / np.sqrt(2.0)))
    
    ei = imp * cdf + sigma * pdf
    ei[sigma < 1e-12] = 0.0
    return ei


class TSFBayesLearner:
    """
    Bayesian optimizer for TSFPhysParamsLearned in log-space.

    Objective:
      score = P_same - λ * P_diff

    Variables (log-space):
      log(alpha), log(beta), log(gamma),
      log(c_s) for selected species,
      log(w_b) for K bands

    You can clamp bounds to keep physical plausibility and stability.
    """

    def __init__(
        self,
        positive_pairs: List[Tuple[object, object]],
        negative_pairs: List[Tuple[object, object]],
        make_hash_fn: Callable[[TSFPhysParamsLearned], Callable[[object], str]],
        base_params: TSFPhysParamsLearned,
        *,
        learn_species_Z: Optional[List[int]] = None,
        learn_band_weights: bool = True,
        learn_temperature: bool = True,
        lambda_penalty: float = 1.0,
        seed: int = 0,
    ):
        self.pos = positive_pairs
        self.neg = negative_pairs
        self.make_hash_fn = make_hash_fn
        self.base = base_params
        self.lambda_penalty = float(lambda_penalty)

        self.learn_species_Z = list(learn_species_Z or [])
        self.learn_band_weights = bool(learn_band_weights)
        self.learn_temperature = bool(learn_temperature)

        self.rng = np.random.default_rng(int(seed))

        # Build parameter vector layout
        # dims: [alpha_A, beta, gamma, T_K (opt), species..., bands...]
        self.n_bands = len(self.base.shell_band_edges) + 1
        num_fixed = 3 + (1 if self.learn_temperature else 0)
        self.dim = num_fixed + len(self.learn_species_Z) + (self.n_bands if self.learn_band_weights else 0)

        # Bounds in log-space (tune to your problem)
        # alpha in [1e-4, 0.2], beta in [0.5, 20], gamma in [0.01, 5]
        self.bounds = []
        self.bounds += [(math.log(1e-5), math.log(0.5))]    # log alpha
        self.bounds += [(math.log(0.1), math.log(50.0))]    # log beta (reduced from 500)
        self.bounds += [(math.log(0.01), math.log(5.0))]    # log gamma
        
        if self.learn_temperature:
            self.bounds += [(math.log(10.0), math.log(1000.0))] # log T_K

        # c_s in [0.3, 3.0]
        for _ in self.learn_species_Z:
            self.bounds += [(math.log(0.3), math.log(3.0))]

        # w_b in [0.3, 3.0]
        if self.learn_band_weights:
            for _ in range(self.n_bands):
                self.bounds += [(math.log(0.3), math.log(3.0))]

        self.bounds = np.array(self.bounds, dtype=np.float64)

    def _scale(self, x: np.ndarray) -> np.ndarray:
        """Map log-space x to [0, 1] unit cube."""
        lo = self.bounds[:, 0]
        hi = self.bounds[:, 1]
        return (x - lo) / (hi - lo)

    def _unscale(self, x_scaled: np.ndarray) -> np.ndarray:
        """Map [0, 1] unit cube back to log-space x."""
        lo = self.bounds[:, 0]
        hi = self.bounds[:, 1]
        return lo + x_scaled * (hi - lo)

    def _vector_to_params(self, x: np.ndarray) -> TSFPhysParamsLearned:
        x = np.asarray(x, dtype=np.float64).ravel()
        alpha = float(np.exp(x[0]))
        beta = float(np.exp(x[1]))
        gamma = float(np.exp(x[2]))

        idx = 3
        if self.learn_temperature:
            T_K = float(np.exp(x[idx])); idx += 1
        else:
            T_K = self.base.T_K

        sc = dict(self.base.species_correction or {})
        for Z in self.learn_species_Z:
            sc[int(Z)] = float(np.exp(x[idx])); idx += 1

        if self.learn_band_weights:
            w = tuple(float(np.exp(v)) for v in x[idx: idx + self.n_bands])
        else:
            w = self.base.shell_band_weights

        return replace(
            self.base,
            alpha_A=alpha,
            beta=beta,
            gamma=gamma,
            T_K=T_K,
            species_correction=sc,
            shell_band_weights=w,
        )

    def _evaluate(self, params: TSFPhysParamsLearned) -> Tuple[float, float, float]:
        """
        Soft-collision objective: maximize log(P_collision) for positive pairs
        and minimize it for negative pairs.
        """
        h = self.make_hash_fn(params)

        def get_log_p_coll(pair):
            a, b = pair
            da, wa = h(a, return_descriptor=True)
            db, wb = h(b, return_descriptor=True)
            
            # Robust overlap: truncate to smaller if mismatch (rare with fixed-bins)
            n = min(len(da), len(db))
            if n == 0: return -100.0
            
            # Gaussian soft overlap: exp(-0.5 * sum((d1-d2)^2 / w^2))
            # We take the log -> -0.5 * sum((d1-d2)^2 / w^2)
            # Use mean bin width
            w_avg2 = (wa[:n]**2 + wb[:n]**2) / 2.0
            diff2 = (da[:n] - db[:n])**2
            
            log_p = -0.5 * np.sum(diff2 / np.maximum(w_avg2, 1e-10))
            return log_p

        pos_logs = [get_log_p_coll(p) for p in self.pos]
        neg_logs = [get_log_p_coll(p) for p in self.neg]

        mean_pos = np.mean(pos_logs)
        mean_neg = np.mean(neg_logs)

        # Objective: Positive pairs should collide (log_p -> 0)
        # Negative pairs should NOT collide (log_p -> -inf)
        score = mean_pos - mean_neg

        # Prior regularization
        prior_penalty = 0.01 * (
            (np.log(params.alpha_A) - np.log(self.base.alpha_A))**2 +
            (np.log(params.beta) - np.log(self.base.beta))**2 +
            (np.log(params.gamma) - np.log(self.base.gamma))**2 +
            (np.log(params.T_K / self.base.T_K))**2
        )
        score -= prior_penalty

        return float(score), float(mean_pos), float(mean_neg)


    def _sample_uniform_scaled(self, n: int) -> np.ndarray:
        return self.rng.random((int(n), self.dim))

    def optimize(
        self,
        *,
        n_init: int = 24,
        n_iter: int = 80,
        n_candidates: int = 4000,
        gp_noise: float = 1e-6,
        gp_var: float = 1.0,
        gp_lengthscale: float = 0.7,
        xi: float = 0.01,
        verbose: bool = True,
    ):
        """
        Bayesian optimization loop:
          - Random init points
          - Fit GP
          - Propose next by maximizing EI over random candidates
        """
        # Init
        X_scaled = []
        y = []
        aux = []  # store (P_same, P_diff, params)

        # Init
        X0_scaled = self._sample_uniform_scaled(int(n_init))
        for i in range(int(n_init)):
            x_log = self._unscale(X0_scaled[i])
            params = self._vector_to_params(x_log)
            score, ps, pd = self._evaluate(params)
            X_scaled.append(X0_scaled[i])
            y.append(score)
            aux.append((ps, pd, params))
            if verbose:
                print(f"[init {i:02d}] score={score:.4f} P_same={ps:.3f} P_diff={pd:.3f} "
                      f"alpha={params.alpha_A:.5f} beta={params.beta:.3f} gamma={params.gamma:.3f} "
                      f"T={params.T_K:.1f} bands={params.shell_band_weights}")

        X_scaled = np.asarray(X_scaled, dtype=np.float64)
        y = np.asarray(y, dtype=np.float64)

        best_idx = int(np.argmax(y))
        best_y = float(y[best_idx])
        best_params = aux[best_idx][2]

        # GP hyperparameters (simple, stable defaults)
        # In scaled space, lengthscales are relative to unit cube [0, 1]^D
        lengthscales = np.full(self.dim, float(gp_lengthscale), dtype=np.float64)
        var = float(gp_var)
        noise = float(gp_noise)

        # Main BO loop
        for it in range(int(n_iter)):
            # Fit GP posterior on candidates in scaled space
            Xcand_scaled = self._sample_uniform_scaled(int(n_candidates))
            mu, var_post = _gp_posterior(X_scaled, y, Xcand_scaled, lengthscales=lengthscales, var=var, noise=noise)
            ei = _expected_improvement(mu, var_post, best_y, xi=xi)

            x_next_scaled = Xcand_scaled[int(np.argmax(ei))]
            x_next_log = self._unscale(x_next_scaled)
            params_next = self._vector_to_params(x_next_log)
            score, ps, pd = self._evaluate(params_next)

            # Append
            X_scaled = np.vstack([X_scaled, x_next_scaled[None, :]])
            y = np.append(y, score)
            aux.append((ps, pd, params_next))

            if score > best_y:
                best_y = float(score)
                best_params = params_next

            if verbose:
                print(f"[bo {it:03d}] score={score:.4f} P_same={ps:.3f} P_diff={pd:.3f} "
                      f"best={best_y:.4f} | alpha={params_next.alpha_A:.5f} beta={params_next.beta:.3f} "
                      f"gamma={params_next.gamma:.3f} T={params_next.T_K:.1f} bands={params_next.shell_band_weights}")

        return best_params, (X_scaled, y, aux)


# ------------------------------------------------------------
# Diagnostics & Visualization
# ------------------------------------------------------------

def plot_collision_matrix(structures: List[object], hash_fn: Callable[[object], str]):
    """
    Plots a heatmap showing hash collisions C[i,j] = (hash(i) == hash(j)).
    Ideal result is block-diagonal if structures are ordered by similarity.
    """
    n = len(structures)
    hashes = [hash_fn(s) for s in structures]
    matrix = np.zeros((n, n), dtype=int)
    for i in range(n):
        for j in range(n):
            if hashes[i] == hashes[j]:
                matrix[i, j] = 1

    plt.figure(figsize=(6, 6))
    plt.imshow(matrix, cmap="Greys", interpolation="nearest")
    plt.title("Hash Collision Matrix")
    plt.xlabel("Structure Index")
    plt.ylabel("Structure Index")
    plt.colorbar(label="Collision (1=Yes, 0=No)")
    plt.show()


def plot_descriptor_distributions(
    pos_pairs: List[Tuple[object, object]],
    neg_pairs: List[Tuple[object, object]],
    hash_fn: Callable[[object], object]
):
    """
    Plots overlapping histograms of descriptor distances for positive and negative pairs.
    hash_fn should be called with return_descriptor=True.
    """
    def get_dist(a, b):
        res_a = hash_fn(a, return_descriptor=True)
        res_b = hash_fn(b, return_descriptor=True)
        da = res_a[0] if isinstance(res_a, tuple) else res_a
        db = res_b[0] if isinstance(res_b, tuple) else res_b
        # Handle length mismatch by padding or truncating
        n = max(len(da), len(db))
        if n == 0: return 0.0
        da_p = np.pad(da, (0, n - len(da))) if len(da) < n else da[:n]
        db_p = np.pad(db, (0, n - len(db))) if len(db) < n else db[:n]
        return np.linalg.norm(da_p - db_p)

    pos_dists = [get_dist(a, b) for a, b in pos_pairs]
    neg_dists = [get_dist(a, b) for a, b in neg_pairs]

    plt.figure(figsize=(8, 5))
    plt.hist(pos_dists, bins=30, alpha=0.5, label="Positive Pairs (Same)", density=True)
    plt.hist(neg_dists, bins=30, alpha=0.5, label="Negative Pairs (Diff)", density=True)
    plt.title("Descriptor Distance Distribution")
    plt.xlabel("L2 Distance ||P_i - P_j||")
    plt.ylabel("Density")
    plt.legend()
    plt.show()


def plot_roc_curve(
    pos_pairs: List[Tuple[object, object]],
    neg_pairs: List[Tuple[object, object]],
    hash_fn: Callable[[object], object]
):
    """
    Plots the ROC curve using a cumulative method (fast & robust).
    """
    def get_dist(a, b):
        res_a = hash_fn(a, return_descriptor=True)
        res_b = hash_fn(b, return_descriptor=True)
        da = res_a[0] if isinstance(res_a, tuple) else res_a
        db = res_b[0] if isinstance(res_b, tuple) else res_b
        n = max(len(da), len(db))
        if n == 0: return 0.0
        da_p = np.pad(da, (0, n - len(da))) if len(da) < n else da[:n]
        db_p = np.pad(db, (0, n - len(db))) if len(db) < n else db[:n]
        return np.linalg.norm(da_p - db_p)

    pos_dists = [get_dist(a, b) for a, b in pos_pairs]
    neg_dists = [get_dist(a, b) for a, b in neg_pairs]

    all_dists = np.array(pos_dists + neg_dists)
    labels = np.array([1] * len(pos_dists) + [0] * len(neg_dists))

    # Higher score should mean "likely Same" -> use negative distance
    scores = -all_dists
    order = np.argsort(scores)[::-1] # descending score
    labels_sorted = labels[order]
    
    # Cumulative counts
    tpr = np.cumsum(labels_sorted) / max(1, labels.sum())
    fpr = np.cumsum(1 - labels_sorted) / max(1, (1 - labels).sum())
    
    # Add origin
    tpr = np.concatenate([[0], tpr])
    fpr = np.concatenate([[0], fpr])

    plt.figure(figsize=(6, 6))
    auc = np.trapz(tpr, fpr)
    print(f"ROC AUC (Descriptor objective quality): {auc:.3f}")
    plt.plot(fpr, tpr, label=f"AUC = {auc:.3f}")
    plt.plot([0, 1], [0, 1], "k--")
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("ROC Curve (Descriptor Stability/Separation)")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.show()

def print_performance_table(name: str, hash_fn: Callable, pos_pairs: list, neg_pairs: list):
    """
    Computes and prints a confusion matrix table: TP, FP, TN, FN.
    """
    tp = sum(1 for a, b in pos_pairs if hash_fn(a) == hash_fn(b))
    fn = len(pos_pairs) - tp
    
    fp = sum(1 for a, b in neg_pairs if hash_fn(a) == hash_fn(b))
    tn = len(neg_pairs) - fp
    
    total = len(pos_pairs) + len(neg_pairs)
    acc = (tp + tn) / total if total > 0 else 0
    
    print(f"\n--- Performance Table: {name} ---")
    print(f"{'Metric':<20} | {'Count':<10} | {'Rate':<10}")
    print("-" * 45)
    print(f"{'True Positives (TP)':<20} | {tp:<10} | {tp/len(pos_pairs)*100 if len(pos_pairs)>0 else 0:>8.1f}%")
    print(f"{'False Negatives (FN)':<20} | {fn:<10} | {fn/len(pos_pairs)*100 if len(pos_pairs)>0 else 0:>8.1f}%")
    print(f"{'True Negatives (TN)':<20} | {tn:<10} | {tn/len(neg_pairs)*100 if len(neg_pairs)>0 else 0:>8.1f}%")
    print(f"{'False Positives (FP)':<20} | {fp:<10} | {fp/len(neg_pairs)*100 if len(neg_pairs)>0 else 0:>8.1f}%")
    print("-" * 45)
    print(f"{'Overall Accuracy':<20} | {'':<10} | {acc*100:>8.1f}%")
    print("-" * 45)

# ------------------------------------------------------------
# Dataset Generation
# ------------------------------------------------------------

class HashDatasetBuilder:
    """
    Automates the generation of positive (perturbed) and negative (cross-basin)
    pairs from a set of prototype structures.
    """
    def __init__(self, prototypes: List[object]):
        self.prototypes = prototypes

    def build(self, n_perturb: int = 10, n_hard: int = 5, alpha: float = 0.02, T: float = 300.0, max_strain: float = 0.01):
        """
        Builds a dataset with positive pairs and both easy/hard negative pairs.
        """
        pos_pairs = []
        neg_pairs = []

        # 1. Positive pairs (Internal basin distortions)
        for s in self.prototypes:
            for _ in range(n_perturb):
                pert = random_perturb(s, alpha=alpha, T=T, max_strain=max_strain)
                pos_pairs.append((s, pert))

        # 2. Hard Negatives (Same composition/cell, different structure)
        for s in self.prototypes:
            for _ in range(n_hard):
                # Species swap
                s_swap = swap_perturb(s, n_swaps=1)
                neg_pairs.append((s, s_swap))
                # Local diffusion
                s_diff = diffuse_perturb(s, rms_dist=1.0)
                neg_pairs.append((s, s_diff))

        # 3. Easy Negatives (Cross-basin comparisons)
        import itertools
        cross_pairs = list(itertools.combinations(self.prototypes, 2))
        
        # Balance neg_pairs to match pos_pairs if needed
        target_neg = len(pos_pairs)
        remaining = target_neg - len(neg_pairs)
        
        if remaining > 0:
            if len(cross_pairs) > remaining:
                idx = np.random.choice(len(cross_pairs), size=remaining, replace=False)
                neg_pairs.extend([cross_pairs[i] for i in idx])
            else:
                neg_pairs.extend(cross_pairs)

        return pos_pairs, neg_pairs


# =============================================================================
#  4) Example wiring
# =============================================================================

def make_hash_fn_learned(params: TSFPhysParamsLearned, fixed_k_bins=None) -> Callable[[object], str]:
    return structure_factor_hash_factory_rotinv_phys_learned(
        phys=params,
        kmax=5,
        k2_grid=0.001,
        fixed_k_bins=fixed_k_bins,
        per_species=True,
        include_power_moment=True,
        chunk_modes=2048,
        use_spglib=False,
    )


# Example usage (replace structure objects with your containers):
#
# base = TSFPhysParamsLearned(
#     T_K=300.0,
#     alpha_A=0.02,
#     beta=2.0,
#     gamma=0.5,
#     shell_band_edges=(2.0, 5.0, 10.0),
#     shell_band_weights=(1.0, 1.0, 1.0, 1.0),
#     species_correction={},
# )
#
# learner = TSFBayesLearner(
#     positive_pairs=pos_pairs,
#     negative_pairs=neg_pairs,
#     make_hash_fn=make_hash_fn_learned,
#     base_params=base,
#     learn_species_Z=[1, 8, 28, 26],  # e.g. H,O,Ni,Fe
#     learn_band_weights=True,
#     lambda_penalty=1.0,
# )
#
# best_params, trace = learner.optimize(n_init=24, n_iter=80, n_candidates=4000)
# print(best_params)

# ------------------------------------------------------------
# Mock container (compatible with your hash function)
# ------------------------------------------------------------

class AtomPositionManagerMock:
    def __init__(self, positions, Z, lattice):
        self.atomPositions = positions
        self._Z = Z
        self.latticeVectors = lattice
        self.E = 0.0
        self.pbc = (True, True, True)

    def get_atomic_numbers(self):
        return self._Z


class StructureMock:
    def __init__(self, positions, Z, lattice):
        self.AtomPositionManager = AtomPositionManagerMock(positions, Z, lattice)


def build_structure(seed=0, noise=0.0, n_atoms=8):
    rng = np.random.default_rng(seed)
    # Randomize lattice to ensure basins are distinct in cell-space too
    lattice = np.eye(3) * 10.0 + rng.normal(scale=1.0, size=(3,3))
    positions = rng.random((n_atoms, 3)) * 5.0
    positions += noise * rng.normal(size=positions.shape)
    
    # Generate Z by cycling through common elements [Ni, O, H]
    base_Z = [28, 8, 1]
    Z = np.array([base_Z[i % len(base_Z)] for i in range(n_atoms)], dtype=np.int32)
    return StructureMock(positions, Z, lattice)



# ------------------------------------------------------------
#  Precomputed Data for fast evolutionary calibration
# ------------------------------------------------------------

@dataclass
class PrecomputedHashData:
    """
    Stores pre-calculated structure factors and lattice metadata for a structure
    to skip the Fourier sum O(N*M) during evolutionary optimization.
    """
    S_complex: np.ndarray    # (C, M) complex structure factors
    k2: np.ndarray           # (M,) square reciprocal magnitudes magnitudes
    species_Z: np.ndarray    # (C,) unique atomic numbers
    N: int                   # number of atoms
    E: float                 # total energy
    V: float                 # volume
    A: np.ndarray            # lattice matrix (rows are vectors)
    denom: np.ndarray        # (C, 1) normalization: sqrt(max(N_s, 1.0))

def precompute_hash_data(container: object, kmax: float = 6.0) -> PrecomputedHashData:
    """
    Pre-calculates all heavy Fourier sums for a container up to a safe kmax.
    """
    apm = container.AtomPositionManager
    Z = np.asarray(apm.get_atomic_numbers(), dtype=np.int32)
    N = int(Z.shape[0])
    X = np.asarray(apm.atomPositions, dtype=np.float64)
    
    A_cols = np.asarray(apm.latticeVectors, dtype=np.float64)
    A = A_cols.T
    FinvT = np.linalg.inv(A).T
    F = (X @ FinvT + EPS) % 1.0
    
    species_Z, inv = np.unique(Z, return_inverse=True)
    C = int(species_Z.size)
    N_s = np.bincount(inv, minlength=C).astype(np.float64)
    denom = np.sqrt(np.maximum(N_s, 1.0))[:, None]
    
    # Modes sphere: build enough modes to cover kmax for THIS cell
    # A_{inv} maps k to integer modes. We need to find m such that |2pi Ainv.T m| <= kmax
    # A more conservative bound: build a cube then sphere-cut.
    max_m = int(math.ceil(kmax * np.max(np.linalg.norm(A, axis=0)) / TWO_PI))
    M = _build_integer_modes_sphere(max_m)
    
    Gstar = _reciprocal_metric(A)
    k2_modes = _compute_k2_for_modes(M, Gstar)
    mask = k2_modes <= (kmax * kmax)
    M = M[mask]
    k2 = k2_modes[mask]
    M_count = int(M.shape[0])
    
    # Fourier sum
    atom_order = np.argsort(inv, kind="mergesort")
    Fg = F[atom_order]
    invg = inv[atom_order]
    bounds = np.concatenate(([0], np.cumsum(np.bincount(invg, minlength=C))))
    
    S_complex = np.zeros((C, M_count), dtype=np.complex128)
    Mb_T = M.T
    # Chunking to avoid memory issues if large
    chunk_size = 1024
    for i in range(0, M_count, chunk_size):
        end = min(i + chunk_size, M_count)
        Mb_chunk = Mb_T[:, i:end]
        phase_complex = np.exp(1j * (TWO_PI * (Fg @ Mb_chunk)))
        for s in range(C):
            a, b = int(bounds[s]), int(bounds[s + 1])
            if b > a:
                S_complex[s, i:end] = np.sum(phase_complex[a:b, :], axis=0)
            
    try:
        E = float(getattr(apm, "E", 0.0))
    except:
        E = 0.0
    V = abs(float(np.linalg.det(A)))
    
    return PrecomputedHashData(
        S_complex=S_complex, k2=k2, species_Z=species_Z,
        N=N, E=E, V=V, A=A, denom=denom
    )

def evaluate_precomputed_descriptor(data: PrecomputedHashData, phys: TSFPhysParamsLearned, bins: np.ndarray):
    """
    Fast vectorized evaluation of the descriptor using precomputed S(k).
    Returns a hash-ready payload (e_tick, Pq_bytes).
    """
    # 1. Slice by current kmax
    k_abs = np.sqrt(data.k2)
    mask = k_abs <= phys.kmax
    S_raw = data.S_complex[:, mask]
    k_abs_m = k_abs[mask]
    k2_m = data.k2[mask]
    
    # 2. Debye-Waller damping
    if phys.B_factor > 0:
        damping = np.exp(-float(phys.B_factor) * k2_m)
        S_raw = S_raw * damping[None, :]
        
    P = np.abs(S_raw) / data.denom
    
    # 3. Binning (Fixed Bins)
    shell_idx = np.searchsorted(bins, k_abs_m) - 1
    valid = (shell_idx >= 0) & (shell_idx < len(bins)-1)
    
    n_shells = len(bins) - 1
    mu1 = np.zeros((data.S_complex.shape[0], n_shells), dtype=np.float64)
    mu2 = np.zeros((data.S_complex.shape[0], n_shells), dtype=np.float64)
    k_abs_shell = np.zeros(n_shells, dtype=np.float64)
    
    for s_bin in range(n_shells):
        bmask = (shell_idx == s_bin) & valid
        if not np.any(bmask): continue
        Pb = P[:, bmask]
        mu1[:, s_bin] = np.mean(Pb, axis=1)
        mu2[:, s_bin] = np.mean(Pb**2, axis=1) # Always compute for fast eval consistency
        k_abs_shell[s_bin] = np.mean(k_abs_m[bmask])
        
    # Shell variance
    var_shell = mu2 - mu1 * mu1
    
    # k-weighting
    eta = float(phys.eta)
    k_w = (k_abs_shell + 1e-12) ** eta
    mu1 *= k_w[None, :]
    mu2 *= k_w[None, :] ** 2
    var_shell *= k_w[None, :] ** 2
        
    # 4. Quantization grid
    sigma_s = _species_sigma_A_learned(data.species_Z, phys)
    band_idx = _shell_band_index(k_abs_shell, phys.shell_band_edges)
    w = np.array(phys.shell_band_weights, dtype=np.float64)
    w_shell = w[band_idx]
    
    # ps_grid calculation
    ps_grid = (float(phys.beta) * data.denom * k_abs_shell[None, :] * sigma_s[:, None]) / np.maximum(w_shell[None, :], 1e-12)
    ps_grid = np.clip(ps_grid, float(phys.min_ps_grid), float(phys.max_ps_grid))
    ps2_grid = np.clip(ps_grid**2, phys.min_ps_grid**2, phys.max_ps_grid**2)
    
    # 5. Final quant
    Pq1 = np.floor(mu1 / ps_grid + 0.5)
    Pq1 = np.clip(Pq1, -2e9, 2e9).astype(np.int32)
    
    Pq2 = np.floor(mu2 / ps2_grid + 0.5)
    Pq2 = np.clip(Pq2, -2e9, 2e9).astype(np.int32)
    
    Pq_var = np.floor(var_shell / ps2_grid + 0.5)
    Pq_var = np.clip(Pq_var, -2e9, 2e9).astype(np.int32)
    
    # Energy tick
    de = max(float(phys.gamma) * K_B_eV_per_K * float(phys.T_K), float(phys.sigma_model_e_eV_per_atom))
    e_tick = int(math.floor((data.E/data.N) / max(de, 1e-12) + 0.5))
    
    # Return as flattened bytes to ensure all learned bits are captured in the surrogate
    return (e_tick, Pq1.tobytes(), Pq2.tobytes(), Pq_var.tobytes())



def evaluate_pairs(hash_fn, positive_pairs, negative_pairs):
    """
    Computes collision rates for positive and negative pairs.
    """
    pos_hits = sum(1 for a, b in positive_pairs if hash_fn(a) == hash_fn(b))
    neg_hits = sum(1 for a, b in negative_pairs if hash_fn(a) == hash_fn(b))
    
    ps = pos_hits / max(1, len(positive_pairs))
    pd = neg_hits / max(1, len(negative_pairs))
    return ps, pd

def print_performance_table(title: str, hash_fn, pos_pairs, neg_pairs):
    ps, pd = evaluate_pairs(hash_fn, pos_pairs, neg_pairs)
    print(f"\n--- Performance Table: {title} ---")
    print(f"  Positives: {len(pos_pairs):4d} | TP: {ps:6.2%} (Stability)")
    print(f"  Negatives: {len(neg_pairs):4d} | FP: {pd:6.2%} (Separation)")
    print(f"  -------------------------------------------")


def calibrate_hash_stability(
    prototypes: List[object],
    noise_level: Dict[str, float],
    target_tp: float = 1.0,
    n_gen: int = 20,
    pop_size: int = 5,
    n_bins: int = 15
) -> TSFPhysParamsLearned:
    """
    Optimized evolutionary calibration using precomputed structure factors.
    """
    if not PYMOO_AVAILABLE:
        print("  [WARNING] pymoo not available. Falling back to simplified logic.")
        T = noise_level.get("T", 300.0)
        alpha = noise_level.get("alpha", 0.125)
        return TSFPhysParamsLearned(T_K=T, alpha_A=alpha*4.0, beta=100.0, gamma=10.0)

    print(f"\n--- Optimized Evolutionary Stability Calibration ---")
    T_target = noise_level.get("T", 300.0)
    alpha_target = noise_level.get("alpha", 0.125)
    
    # 1. Generate Dataset
    builder = HashDatasetBuilder(prototypes)
    pos_pairs, neg_pairs = builder.build(n_perturb=20, n_hard=6, T=T_target, alpha=alpha_target, max_strain=0.0)

    # 2. Precompute Structure Factors (THE SPEEDUP)
    print(f"  [PRECOMPUTE] Calculating Fourier sums for unique structures...")
    unique_structs = []
    seen_ids = {}
    
    def get_sid(s):
        oid = id(s)
        if oid not in seen_ids:
            seen_ids[oid] = len(unique_structs)
            unique_structs.append(s)
        return seen_ids[oid]

    pos_idx = [(get_sid(a), get_sid(b)) for a, b in pos_pairs]
    neg_idx = [(get_sid(a), get_sid(b)) for a, b in neg_pairs]
    
    # Precalculate up to 10.0 kmax to be safe for any kmax optimization
    cache = [precompute_hash_data(s, kmax=10.0) for s in unique_structs]
    print(f"  [PRECOMPUTE] Done. Cached {len(cache)} structures.")

    # Shared parameters
    all_Z = set()
    for c in cache:
        all_Z.update(c.species_Z.tolist())
    species_list = sorted(list(all_Z))
    n_species = len(species_list)
    band_edges = (1.5, 3.0, 6.0)
    n_bands = len(band_edges) + 1

    # Variables
    n_fixed = 8
    n_var = n_fixed + n_species + n_bands
    # [0: log_alpha_A, 1: log_beta, 2: log_gamma, 3: log_T_K, 4: log_kmax, 5: p, 6: log_B, 7: eta]
    xl = np.array([np.log(1e-4), np.log(1.0),   np.log(0.1),  np.log(10.0), np.log(0.5), 0.5, np.log(1e-6), 0.0])
    xu = np.array([np.log(2.0),  np.log(20000.0), np.log(100.0), np.log(3000.0), np.log(10.0), 3.0, np.log(10.0), 4.0])
    xl = np.concatenate([xl, np.full(n_species, np.log(0.3)), np.full(n_bands, np.log(0.3))])
    xu = np.concatenate([xu, np.full(n_species, np.log(3.0)), np.full(n_bands, np.log(3.0))])

    class VectorizedStabilityProblem(ElementwiseProblem):
        def __init__(self):
            super().__init__(n_var=n_var, n_obj=3, xl=xl, xu=xu)

        def _evaluate(self, x, out, *args, **kwargs):
            # Unpack individual
            alpha_A = np.exp(x[0]); beta = np.exp(x[1]); gamma = np.exp(x[2]); T_K = np.exp(x[3])
            kmax = np.exp(x[4]); p = x[5]; B = np.exp(x[6]); eta = x[7]
            sc = {species_list[i]: np.exp(x[n_fixed + i]) for i in range(n_species)}
            wb = tuple(np.exp(x[n_fixed + n_species + i]) for i in range(n_bands))
            
            phys = TSFPhysParamsLearned(
                T_K=T_K, alpha_A=alpha_A, beta=beta, gamma=gamma,
                kmax=kmax, power_law_exponent=p, B_factor=B, eta=eta,
                species_correction=sc, shell_band_edges=band_edges, shell_band_weights=wb
            )
            bins = _generate_power_law_bins(kmax, p, n_bins)
            
            # Evaluate all unique structures' quantized descriptors
            results = [evaluate_precomputed_descriptor(c, phys, bins) for c in cache]
            
            # Calculate TP/FP using indices
            pos_hits = sum(1 for ia, ib in pos_idx if results[ia] == results[ib])
            neg_hits = sum(1 for ia, ib in neg_idx if results[ia] == results[ib])
            
            tp = pos_hits / len(pos_idx)
            fp = neg_hits / len(neg_idx)
            
            out["F"] = [1.0 - tp, fp, x[1] + x[4]*0.5]

    res = minimize(VectorizedStabilityProblem(),
                   NSGA2(pop_size=pop_size, eliminate_duplicates=True),
                   termination=('n_gen', n_gen),
                   seed=1, verbose=False)

    if res.F is None or len(res.F) == 0:
         return TSFPhysParamsLearned(T_K=T_target, alpha_A=alpha_target*2)

    F, X = res.F, res.X
    candidates = [(i, 1.0 - F[i, 0], F[i, 1], F[i, 2]) for i in range(len(F)) if (1.0 - F[i, 0]) >= target_tp]
    if not candidates:
        best_idx = np.argsort(F[:, 0])[0]
    else:
        good_fp = [c for c in candidates if c[2] < 0.1]
        best_c = min(good_fp if good_fp else candidates, key=lambda x: x[3])
        best_idx = best_c[0]
        print(f"  [SUCCESS] TP={best_c[1]:.2%}, FP={best_c[2]:.2%}")

    best_x = X[best_idx]
    return TSFPhysParamsLearned(
        alpha_A=np.exp(best_x[0]), beta=np.exp(best_x[1]) * 1.15, gamma=np.exp(best_x[2]) * 1.1,
        T_K=np.exp(best_x[3]), kmax=np.exp(best_x[4]), power_law_exponent=best_x[5], 
        B_factor=np.exp(best_x[6]), eta=best_x[7],
        species_correction={species_list[i]: np.exp(best_x[n_fixed + i]) for i in range(n_species)},
        shell_band_edges=band_edges,
        shell_band_weights=tuple(np.exp(best_x[n_fixed + n_species + i]) for i in range(n_bands))
    )


def find_structures_duplicates(structures: List[object], hash_fn: Callable) -> Dict[str, List[int]]:
    """
    Groups indices of structures that produce the same hash.
    Useful for identifying repeated structures in a dataset.
    
    Returns:
        A dictionary mapping hash strings to lists of integer indices.
    """
    groups = {}
    for i, s in enumerate(structures):
        h = hash_fn(s)
        if h not in groups:
            groups[h] = []
        groups[h].append(i)
    return groups


# ------------------------------------------------------------
# Manual Usage Example: Directly using physical parameters
# ------------------------------------------------------------

def demonstrate_physical_parameters():
    """
    Shows how to use the hash with direct, hand-tuned physical parameters
    across multiple structures (N-structures) and reports a confusion matrix.
    """
    print("\n=== ADVANCED DEMO: N-Structures & Confusion Matrix ===")
    
    # 1. Setup Parameters (Hand-tuned for stability)
    sc = {28: 1.0, 8: 1.2, 1: 0.8}
    phys = TSFPhysParamsLearned(
        T_K=400.0,      # As requested by user
        alpha_A=0.125,  # Slightly broader search for thermal stability
        beta=25.0,      # Coarse quantization to group basins
        gamma=0.5,
        sigma_model_e_eV_per_atom=0.01,
        species_correction=sc
    )

    # 2. Build Hash Function
    fixed_bins = np.linspace(0.1, 5.0, 51)
    hash_fn = make_hash_fn_learned(phys, fixed_k_bins=None)

    # 3. Create N-Structures and Perturbations
    n_protos = 5
    n_pert = 5
    seeds = [123, 456, 789, 101, 202]
    protos = [build_structure(seed=s, n_atoms=200) for s in seeds]
    
    builder = HashDatasetBuilder(protos)
    # Generate positive pairs (perturbations) and negative pairs (cross-basin)
    # Using alpha=0.003 for perturbations to simulate typical thermal noise
    pos_pairs, neg_pairs = builder.build(n_perturb=n_pert, n_hard=0, T=300.0, alpha=0.003, max_strain=0.0)

    # 4. Diagnostics
    print(f"Physical Settings: T={phys.T_K}K, alpha_A={phys.alpha_A}A, beta={phys.beta}")
    print(f"Dataset: {n_protos} basins, {len(pos_pairs)} positive pairs, {len(neg_pairs)} negative pairs")
    
    # Print the confusion matrix table
    print_performance_table("ADVANCED PHYSICAL DEMO", hash_fn, pos_pairs, neg_pairs)

    # 5. Quick check on a single pair (as previously requested)
    p1 = protos[0]
    p1_pert = random_perturb(p1, T=300.0, alpha=0.003)
    h1 = hash_fn(p1)
    h1_p = hash_fn(p1_pert)
    print(f"\nStability Check (Match?): {'YES' if h1 == h1_p else 'NO'}")


def demonstrate_stability_calibration():
    """
    Shows how calibrate_hash_stability can find the right parameters
    to ensure hashes "classify" the same structures together.
    """


    print("\n=== DEMO: Multi-Parameter Stability Calibration ===")
    
    # Realistic large system
    n_atoms = 200
    protos = [build_structure(seed=i, n_atoms=n_atoms) for i in range(10)]
    noise = {"T": 300.0, "alpha": 0.01, "max_strain": 0.0} 

    '''
    test_noise = {"T": 1000.0, "alpha": 0.01, "max_strain": 0.0} 
    test_list = []
    test_list.append(protos[0])
    test_list.extend([random_perturb(protos[0], **test_noise) for _ in range(4)])
    print( np.max(test_list[0].AtomPositionManager.atomPositions- test_list[1].AtomPositionManager.atomPositions) )
    asdf
    '''

    # Calibrate stability (now learns kmax, p, etc.)
    stable_phys = calibrate_hash_stability(protos, noise, target_tp=1.0)
    
    # Verify with Duplicate detection
    # Use learned power-law bins
    fixed_bins = _generate_power_law_bins(stable_phys.kmax, stable_phys.power_law_exponent, 15)
    
    def classification_hash(s):
        return structure_factor_hash_factory_rotinv_phys_learned(
            phys=stable_phys, kmax=stable_phys.kmax, k2_grid=0.1, fixed_k_bins=fixed_bins,
            include_energy=True,
            include_volume=False # Ignore volume for basin ID
        )(s)
    
    # Test set: MULTI-BASIN list (Variants of Basin 0 AND Basin 1)
    test_noise = {"T": 1000.0, "alpha": 0.01, "max_strain": 0.0} 
    test_list = []
    test_list.append(protos[0])
    test_list.extend([random_perturb(protos[0], **test_noise) for _ in range(4)])
    test_list.append(protos[1])
    test_list.extend([random_perturb(protos[1], **test_noise) for _ in range(4)])
    
    groups = find_structures_duplicates(test_list, classification_hash)
    

    print("\nVerification on Multi-Basin List (Basin 0 variants + Basin 1 variants):")
    n_groups = len(groups)
    print(f"  Number of unique hash groups found: {n_groups}")
    
    # Audit group membership
    for h, idxs in groups.items():
        print(f"  Group [Hash {h[:8]}...]: Indices {idxs}")
    
    if n_groups == 2:
        print("  [SUCCESS] Multi-parameter calibration achieved 100% basin stability AND separation.")
    else:
        print(f"  [FAILED] Found {n_groups} groups - Resolution or Stability issue.")


def demonstrate_duplication_identification():
    """
    Uses find_structures_duplicates to group identical/repeated structures.
    """
    print("\n=== DEMO: Duplication Identification ===")
    
    # 1. Setup Parameters
    sc = {28: 1.0, 8: 1.2, 1: 0.8}
    phys = TSFPhysParamsLearned(
        T_K=300.0,      # As requested by user
        alpha_A=0.015,  # Slightly broader search for thermal stability
        beta=25.0,      # Coarse quantization to group basins
        gamma=0.5,
        sigma_model_e_eV_per_atom=0.01,
        species_correction=sc
    )

    fixed_bins = np.linspace(0.1, 5.0, 51)
    hash_fn = make_hash_fn_learned(phys, fixed_k_bins=None)

    # 2. Build a list with known duplicates (using larger structures)
    n_atoms = 16
    p1 = build_structure(seed=13, n_atoms=n_atoms)
    p1_copy = copy.deepcopy(p1) # Exact copy
    p1_pert = random_perturb(p1, T=300.0, alpha=0.003, max_strain=0.00) # Perturbed
    p2 = build_structure(seed=34, n_atoms=n_atoms)
    p3 = build_structure(seed=34, n_atoms=n_atoms)
    p4 = build_structure(seed=34, n_atoms=n_atoms)
    p5 = build_structure(seed=34, n_atoms=n_atoms)
    print(f"Sample displacement (P1 vs P1_pert): {np.linalg.norm(p1.AtomPositionManager.atomPositions - p1_pert.AtomPositionManager.atomPositions):.4f} A")

    structures = [p1, p1_copy, p1_pert, p2, p3, p4, p5]
    print(f"Testing list of {len(structures)} structures...")

    # 3. Find duplicates
    duplicates = find_structures_duplicates(structures, hash_fn)
    
    # 4. Report
    unique_count = len(duplicates)
    print(f"Found {unique_count} unique hash groups.")
    
    for h, indices in duplicates.items():
        if len(indices) > 1:
            print(f"Group [Hash {h[:8]}...]: Indices {indices} (REPEATED)")
        else:
            print(f"Group [Hash {h[:8]}...]: Index {indices} (UNIQUE)")


# ------------------------------------------------------------
# End-to-End Recipe: Data Augmentation, Training, & Persistence
# ------------------------------------------------------------

def run_experiment():
    """
    Complete end-to-end recipe:
    1. Augmentation from prototypes (Train/Val split)
    2. Bayesian Optimization on Training Set
    3. Diagnostics (ROC, Collision Matrix)
    4. Confusion Matrix for Training and Validation
    5. Save/Reload verification
    """
    print("=== Step 1: Data Augmentation (Train/Val Split) ===")
    # Define 6 prototypes
    seeds = [101, 202, 303, 404, 505, 606]
    protos = [build_structure(seed=s) for s in seeds]
    # Split: 4 for training, 2 for validation
    train_protos = protos[:4]
    val_protos = protos[4:]
    
    builder_train = HashDatasetBuilder(train_protos)
    builder_val = HashDatasetBuilder(val_protos)
    
    # Define fixed binning for consistent descriptor diagnostics
    fixed_bins = np.linspace(0.1, 5.0, 51) # 50 shells up to kmax=5
    
    # Generate datasets (disable strain noise for now: max_strain=0.0)
    pos_train, neg_train = builder_train.build(n_perturb=15, n_hard=2, T=300.0, alpha=0.003, max_strain=0.0)
    pos_val, neg_val = builder_val.build(n_perturb=15, n_hard=2, T=300.0, alpha=0.003, max_strain=0.0)
    
    # Customized factory for this experiment (uses fixed_bins)
    def my_make_hash(p):
        return make_hash_fn_learned(p, fixed_k_bins=fixed_bins)
    
    print(f"Training set: {len(pos_train)} positive, {len(neg_train)} negative pairs (from 4 basins)")
    print(f"Validation set: {len(pos_val)} positive, {len(neg_val)} negative pairs (from 2 basins)")

    print("\n=== Step 2: Bayesian Optimization (Training) ===")
    base_params = TSFPhysParamsLearned()
    learner = TSFBayesLearner(
        pos_train,
        neg_train,
        make_hash_fn=my_make_hash,
        base_params=base_params,
        learn_species_Z=[1, 8, 28],
        learn_temperature=True,
        lambda_penalty=0.1
    )

    best_params, trace = learner.optimize(n_init=20, n_iter=80, n_candidates=3000)
    
    print("\n=== Step 3: Diagnostic Metrics ===")
    best_hash_fn = my_make_hash(best_params)
    
    # Summary of progress
    pos_rates = np.array([t[0] for t in trace[2]])
    neg_rates = np.array([t[1] for t in trace[2]])
    plt.figure(figsize=(8,4))
    plt.plot(pos_rates, label="Training Stability (P_same)")
    plt.plot(neg_rates, label="Training Uniqueness (P_diff)")
    plt.title("TSF Training Progress")
    plt.legend(); plt.show()

    # Advanced plots (on Training set)
    plot_collision_matrix(train_protos + [random_perturb(p, alpha=0.01) for p in train_protos[:2]], best_hash_fn)
    plot_roc_curve(pos_train, neg_train, best_hash_fn)

    print("\n=== Step 4: Performance Tables (Confusion Matrices) ===")
    print_performance_table("TRAINING SET", best_hash_fn, pos_train, neg_train)
    print_performance_table("VALIDATION SET (Unseen Basins)", best_hash_fn, pos_val, neg_val)

    print("\n=== Step 5: Parameter Persistence (Save/Load) ===")
    json_path = "tsf_params_best.json"
    best_params.save_json(json_path)
    print(f"Saved best parameters to {json_path}")
    
    reloaded_params = TSFPhysParamsLearned.load_json(json_path)
    final_hash = my_make_hash(reloaded_params)
    
    # Test on a completely 'fresh' perturbation from validation set
    pname_val = val_protos[0]
    p_fresh = random_perturb(pname_val, T=300.0, alpha=0.01) 
    h_orig = final_hash(pname_val)
    h_pert = final_hash(p_fresh)
    
    print(f"\nFinal Deployment Test (Validation Basin):")
    print(f"  Original Prototype Hash: {h_orig}")
    print(f"  Fresh Perturbation Hash: {h_pert}")
    print(f"  Match: {'YES' if h_orig == h_pert else 'NO'}")

# ------------------------------------------------------------
# Run test
# ------------------------------------------------------------

def train_hash_calibrator(
    train_structures: List[object],
    *,
    learn_species=[1, 8, 28],
    n_perturb=15,
    n_hard=3,
    T=1000.0,
    alpha=0.01,
    max_strain=0.0
):
    """
    Train the TSF hash parameters using a dataset of structures.

    Returns:
        best_params
        trained_hash_fn
    """

    print("Building training dataset...")

    builder = HashDatasetBuilder(train_structures)
    pos_pairs, neg_pairs = builder.build(
        n_perturb=n_perturb,
        n_hard=n_hard,
        T=T,
        alpha=alpha,
        max_strain=max_strain
    )

    print(f"Training pairs:")
    print(f"   positives: {len(pos_pairs)}")
    print(f"   negatives: {len(neg_pairs)}")

    base_params = TSFPhysParamsLearned()

    learner = TSFBayesLearner(
        positive_pairs=pos_pairs,
        negative_pairs=neg_pairs,
        make_hash_fn=make_hash_fn_learned,
        base_params=base_params,
        learn_species_Z=learn_species,
        learn_band_weights=True,
        learn_temperature=True,
        lambda_penalty=0.1
    )

    best_params, trace = learner.optimize(
        n_init=20,
        n_iter=60,
        n_candidates=2000
    )

    print("\nTraining complete.")
    print(best_params)

    trained_hash = make_hash_fn_learned(best_params)

    return best_params, trained_hash

def predict_structure_similarity(
    structures: List[object],
    hash_fn: Callable[[object], str]
):
    """
    Predict which structures are identical using a trained hash.

    Returns:
        groups: {hash : [indices]}
    """

    hashes = [hash_fn(s) for s in structures]

    groups = {}

    for i, h in enumerate(hashes):
        if h not in groups:
            groups[h] = []
        groups[h].append(i)

    return groups

def classify_pairs(
    structures: List[object],
    hash_fn: Callable[[object], str]
):
    """
    Returns pairwise classification matrix:
    1 = same structure
    0 = different
    """

    n = len(structures)
    hashes = [hash_fn(s) for s in structures]

    matrix = np.zeros((n, n), dtype=int)

    for i in range(n):
        for j in range(n):

            if hashes[i] == hashes[j]:
                matrix[i, j] = 1

    return matrix




if __name__ == "__main__":
    # ============================
    # TRAINING
    # ============================

    train_structures = [build_structure(seed=i, n_atoms=10) for i in range(6)]

    best_params, hash_fn = train_hash_calibrator(train_structures)


    # ============================
    # PREDICTION
    # ============================

    test_structures = [
        build_structure(seed=1, n_atoms=10),
        random_perturb(build_structure(seed=1, n_atoms=10), T=300),
        build_structure(seed=5, n_atoms=10),
    ]
    print( test_structures[0].AtomPositionManager.atomPositions - test_structures[1].AtomPositionManager.atomPositions )
    groups = predict_structure_similarity(test_structures, hash_fn)

    print("\nStructure groups:")
    for h, idx in groups.items():
        print(h[:10], "->", idx)
        
    # 1. Stability Calibration (Finds parameters for 100% basin-matching)
    #demonstrate_stability_calibration()

    # 2. Physical Parameters Demo (High-N structures)
    #demonstrate_physical_parameters()
    # 4. End-to-End Experiment (Bayesian Optimization + T learning)
    #run_experiment()
