"""Small, dependency-light helpers for mapping independent sensor clocks."""

from __future__ import annotations

from dataclasses import asdict, dataclass

import numpy as np


@dataclass(frozen=True)
class AffineClockModel:
    """Numerically stable affine map ``y = y_ref + scale * (x - x_ref)``."""

    x_ref: float
    y_ref: float
    scale: float
    sample_count: int
    inlier_count: int
    residual_std_s: float
    residual_p95_s: float

    def map(self, value: float | np.ndarray) -> float | np.ndarray:
        array = np.asarray(value, dtype=np.float64)
        mapped = self.y_ref + self.scale * (array - self.x_ref)
        return float(mapped) if array.ndim == 0 else mapped

    def inverse(self, value: float | np.ndarray) -> float | np.ndarray:
        if abs(self.scale) < 1e-12:
            raise ValueError("clock model scale is zero")
        array = np.asarray(value, dtype=np.float64)
        mapped = self.x_ref + (array - self.y_ref) / self.scale
        return float(mapped) if array.ndim == 0 else mapped

    def to_dict(self) -> dict[str, float | int]:
        return asdict(self)


def fit_affine_clock(
    x: np.ndarray,
    y: np.ndarray,
    *,
    max_iterations: int = 4,
    min_residual_gate_s: float = 5e-4,
) -> AffineClockModel:
    """Robustly fit an affine clock map while rejecting receive-time spikes.

    ``x`` and ``y`` may have large, unrelated epochs. Centering around their
    medians avoids losing precision when host UTC is around 1e9 seconds.
    """

    x_values = np.asarray(x, dtype=np.float64).reshape(-1)
    y_values = np.asarray(y, dtype=np.float64).reshape(-1)
    finite = np.isfinite(x_values) & np.isfinite(y_values)
    x_values = x_values[finite]
    y_values = y_values[finite]
    if x_values.size < 2:
        raise ValueError("need at least two finite clock samples")

    x_ref = float(np.median(x_values))
    y_ref = float(np.median(y_values))
    dx = x_values - x_ref
    dy = y_values - y_ref
    inliers = np.ones(x_values.size, dtype=bool)
    scale = 1.0
    offset = 0.0

    for _ in range(max_iterations):
        local_x = dx[inliers]
        local_y = dy[inliers]
        denom = float(local_x @ local_x)
        if denom < 1e-18:
            raise ValueError("clock samples do not span enough time")
        scale = float(local_x @ local_y / denom)
        offset = float(np.median(local_y - scale * local_x))
        residual = dy - (offset + scale * dx)
        center = float(np.median(residual[inliers]))
        mad = float(np.median(np.abs(residual[inliers] - center)))
        sigma = 1.4826 * mad
        gate = max(float(min_residual_gate_s), 6.0 * sigma)
        updated = np.abs(residual - center) <= gate
        if np.count_nonzero(updated) < 2 or np.array_equal(updated, inliers):
            break
        inliers = updated

    # Fold the small centered intercept into y_ref so map/inverse stay simple.
    y_ref += offset
    residual = y_values - (y_ref + scale * (x_values - x_ref))
    residual_inliers = residual[inliers]
    return AffineClockModel(
        x_ref=x_ref,
        y_ref=y_ref,
        scale=scale,
        sample_count=int(x_values.size),
        inlier_count=int(np.count_nonzero(inliers)),
        residual_std_s=float(np.std(residual_inliers)),
        residual_p95_s=float(np.percentile(np.abs(residual_inliers), 95.0)),
    )