"""Numerical primitives reconstructed from the exact judged logbook.

The judged Hugging Face revision embedded ``verify.py`` but did not preserve this
imported helper.  These functions implement the formulas described by that
logbook and by Algorithm 1 of arXiv:2507.06161.
"""

from __future__ import annotations

import numpy as np
from scipy.linalg import expm
from scipy.spatial.distance import cdist


def gaussian_kernel(points: np.ndarray, sigma: float) -> np.ndarray:
    distances = cdist(points, points, metric="sqeuclidean")
    return np.exp(-distances / (2.0 * sigma**2))


def exponential_kernel(points: np.ndarray, sigma: float) -> np.ndarray:
    distances = cdist(points, points, metric="euclidean")
    return np.exp(-distances / sigma)


def heat_kernel(laplacian: np.ndarray, time: float) -> np.ndarray:
    return expm(-time * laplacian)


def symmetric_sinkhorn(
    smoothing: np.ndarray,
    tol: float = 1e-12,
    max_iter: int = 1_000,
) -> tuple[np.ndarray, np.ndarray, int, float, list[float]]:
    """Apply the symmetric Sinkhorn update from Algorithm 1.

    The returned error curve uses the maximum absolute row-sum residual, which
    matches the stopping behavior reported in the judged logbook.
    """

    n = smoothing.shape[0]
    scaling = np.ones(n, dtype=np.float64)
    errors: list[float] = []

    for iteration in range(1, max_iter + 1):
        normalized = scaling[:, None] * smoothing * scaling[None, :]
        residual = normalized @ np.ones(n, dtype=np.float64) - 1.0
        error = float(np.max(np.abs(residual)))
        errors.append(error)
        if error < tol:
            return normalized, scaling, iteration, error, errors

        denominator = smoothing @ scaling
        scaling = np.sqrt(scaling / np.maximum(denominator, 1e-300))

    normalized = scaling[:, None] * smoothing * scaling[None, :]
    residual = normalized @ np.ones(n, dtype=np.float64) - 1.0
    error = float(np.max(np.abs(residual)))
    return normalized, scaling, max_iter, error, errors


def axiom_symmetry(operator: np.ndarray) -> float:
    return float(np.max(np.abs(operator - operator.T)))


def axiom_mass_conservation(operator: np.ndarray) -> float:
    return float(np.max(np.abs(operator @ np.ones(operator.shape[0]) - 1.0)))


def axiom_spectrum(operator: np.ndarray) -> tuple[float, float]:
    eigenvalues = np.linalg.eigvalsh((operator + operator.T) / 2.0)
    return float(eigenvalues[0]), float(eigenvalues[-1])


def axiom_positivity(operator: np.ndarray) -> float:
    mask = ~np.eye(operator.shape[0], dtype=bool)
    return float(np.min(operator[mask]))


def is_diffusion_operator(operator: np.ndarray, tol: float = 1e-6) -> bool:
    low, high = axiom_spectrum(operator)
    return bool(
        axiom_symmetry(operator) < tol
        and axiom_mass_conservation(operator) < tol
        and low >= -tol
        and high <= 1.0 + tol
        and axiom_positivity(operator) >= -tol
    )