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

# ---------------------------------------------------------------------------
# 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)
    damping_factor: float = 0.5,              # exponential damping np.exp(-damping * |k|)
    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 = False,
    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,                # applies exponential damping based on damping_factor
    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 = np.exp(-damping_factor * np.sqrt(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

def structure_factor_vector_factory(
    *,
    # 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
    damping_factor: float = 0.5,              # exponential damping np.exp(-damping * |k|)
    metric_grid: float = 5e-2,                # Å^2 grid for metric tensor G = A^T A (not heavily used in vector math, kept for logic parity)
    include_energy: bool = True,
    include_volume: bool = True,

    # Canonicalization for supercell/primitive equivalence (recommended)
    use_spglib: bool = False,
    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,                # applies exponential damping based on damping_factor
    normalize_mode: str = "sqrtN",            # "sqrtN" or "sqrtNz"

    # Performance
    chunk_size: int = 20000,
    debug: bool = False,
    return_raw: bool = False,  # If True, returns dict with raw components instead of 1D array
) -> Callable[[object], Union[np.ndarray, Dict[str, np.ndarray]]]:
    """
    Returns a continuous vector extraction function instead of a discrete SHA-256 hash.
    This continuous vector (np.ndarray of fixed size) is suitable for FAISS or KD-Tree indexing
    to perform rapid Euclidean distance searches over 10^8 structures.
    """
    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
        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 = np.exp(-damping_factor * np.sqrt(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,)

    def vec_fn(container) -> np.ndarray:
        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 np.zeros(S_count, dtype=np.float32)

        if is_nonpbc:
            # Simple fallback for molecular structures (not diffraction based)
            X = np.asarray(apm.atomPositions, float)
            Xc = X - X.mean(axis=0)
            D = np.linalg.norm(Xc, axis=1)
            H, _ = np.histogram(D, bins=S_count, range=(0, kmax))
            return H.astype(np.float32)

        # 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 np.zeros(S_count, dtype=np.float32)
            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:
            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
        
        # Structure factor accumulation
        S_re = np.zeros((C, M_count), dtype=np.float64)
        S_im = np.zeros((C, M_count), dtype=np.float64)

        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]
            for c in range(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
        shell_counts_f = shell_counts.astype(np.float64)
        spec_shell = np.zeros((C, S_count), dtype=np.float64)
        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)

        # Do not compress dynamic range (no log1p) to keep Euclidean distances accurate
        spec_shell = np.maximum(spec_shell, 0.0)

        # Build feature vector combining spec_shell and invariant metric
        eigvals = np.linalg.eigvalsh(G) # invariant intrinsic cell metric, sort order natively
        # ... Unroll and concatenate ...
        # If the user wants the raw continuous tensor without flattening to a specific format
        if return_raw:
            return {
                "S_re": S_re, 
                "S_im": S_im, 
                "spec_shell": spec_shell, 
                "eigvals_raw": np.linalg.eigvalsh(G)
            }
            
        S_global = spec_shell.sum(axis=0)
        vector = np.concatenate([S_global, eigvals])
        
        return vector.astype(np.float32)

    return vec_fn

if __name__ == "__main__":
    from scipy.spatial import cKDTree
    import time
    
    # Mock classes to simulate the expected container structure
    class MockAtomPositionManager:
        def __init__(self, positions, lattice, numbers):
            self.atomPositions = np.array(positions)
            self.latticeVectors = np.array(lattice)
            self.numbers = np.array(numbers)
            self.pbc = [True, True, True]
            self.E = -10.0
            
        def get_atomic_numbers(self):
            return self.numbers
            
    class MockContainer:
        def __init__(self, apm):
            self.AtomPositionManager = apm

    print("--- Evaluación de Búsqueda Vectorial Continua ---")
    m_trials = 100
    N = 100
    lattice = np.eye(3) * 10.0
    
    # Generate the vector extractor function
    # We use kmax=9 and damping=1.0 for a long, rich, but robust vector
    vec_fn = structure_factor_vector_factory(kmax=9, damping_factor=1.0)
    
    clean_vectors = []
    low_noise_vectors = []
    high_noise_vectors = []
    
    start_t = time.time()
    for i in range(m_trials):
        np.random.seed(i)
        positions = np.random.uniform(0, 10.0, size=(N, 3))
        numbers = np.random.choice([1, 8, 6], size=N)
        
        # Clean
        apm_clean = MockAtomPositionManager(positions, lattice, numbers)
        container_clean = MockContainer(apm_clean)
        clean_vectors.append(vec_fn(container_clean))
        
        # Low noise (0.5A normal)
        noise_low = np.random.normal(0, 0.5, size=(N, 3))
        apm_low = MockAtomPositionManager(positions + noise_low, lattice, numbers)
        container_low = MockContainer(apm_low)
        low_noise_vectors.append(vec_fn(container_low))
        
        # High noise -> different structure entirely (Uniform 2.2A)
        noise_high = np.random.uniform(-2.2, 2.2, size=(N, 3))
        apm_high = MockAtomPositionManager(positions + noise_high, lattice, numbers)
        container_high = MockContainer(apm_high)
        high_noise_vectors.append(vec_fn(container_high))
        
    ext_time = time.time() - start_t
    print(f"Extracción de {m_trials*3} vectores tomó {ext_time:.3f} s ({(ext_time)/(m_trials*3)*1000:.2f} ms por vector)")
    
    # Construcción de base de datos
    clean_vectors = np.array(clean_vectors)
    low_noise_vectors = np.array(low_noise_vectors)
    high_noise_vectors = np.array(high_noise_vectors)
    
    dim = clean_vectors.shape[1]
    print(f"Dimensión del vector D = {dim}")
    
    # We use cKDTree for log(N) massive queries
    tree = cKDTree(clean_vectors)
    
    # Medir distancias a sus correspondientes (diagonal)
    dist_low = np.linalg.norm(clean_vectors - low_noise_vectors, axis=1)
    dist_high = np.linalg.norm(clean_vectors - high_noise_vectors, axis=1)
    
    max_dist_low = np.max(dist_low)
    min_dist_high = np.min(dist_high)
    
    print("\n--- Estadísticas de Distancia L2 ---")
    print(f"Media Distancia Ruido BAJO (0.5A) : {np.mean(dist_low):.4f}  [Max: {max_dist_low:.4f}]")
    print(f"Media Distancia Ruido ALTO (2.2A) : {np.mean(dist_high):.4f}  [Min: {min_dist_high:.4f}]")
    
    # Simulamos threshold = (max_low + min_high) / 2
    umbral = (max_dist_low + min_dist_high) / 2
    print(f"\n=> Umbral de separación propuesto: {umbral:.4f}")
    
    if max_dist_low < min_dist_high:
        print("¡Separación perfecta lograda! Todos los 'low noise' están más cerca que el peor 'high noise'.")
    else:
        print("Hay solapamiento en las distribuciones con este kmax/damping.")
        
    # Example database query
    # We look for how many neighbors are within the radius determined by the threshold
    start_q = time.time()
    res_low = tree.query_ball_point(low_noise_vectors, r=umbral)
    res_high = tree.query_ball_point(high_noise_vectors, r=umbral)
    q_time = time.time() - start_q
    
    correct_accept = sum(1 for hits in res_low if len(hits) > 0)
    correct_reject = sum(1 for hits in res_high if len(hits) == 0)
    
    print(f"\nConsulta al árbol KDTree (100 estructuras) tomó: {q_time*1000:.2f} ms")
    print(f"Evaluación Final de Base de Datos ({m_trials} trials):")
    print(f"  - Repetidos detectados con éxito (Ruido Bajo): {correct_accept}/{m_trials}")
    print(f"  - Nuevos detectados con éxito (Ruido Alto) : {correct_reject}/{m_trials}")
