"""Numerical baseline-twin gates for non-byte-deterministic GPU CFD."""
from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np


def _relative_rms(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
    difference = np.asarray(a, float) - np.asarray(b, float)
    rms = float(np.sqrt(np.mean(difference * difference)))
    scale = max(float(np.std(a)), float(np.std(b)), 1e-12)
    return rms, rms / scale


def numerical_twin_gate(
    zero_dir: str | Path,
    twin_dir: str | Path,
    *,
    max_field_relative_rms: float = 0.01,
    max_telemetry_relative_rms: float = 0.01,
) -> dict[str, Any]:
    """Compare same-checkpoint zero branches and fail if numerical noise is material."""
    zero, twin = Path(zero_dir), Path(twin_dir)
    report: dict[str, Any] = {"thresholds": {
        "field_relative_rms": max_field_relative_rms,
        "telemetry_relative_rms": max_telemetry_relative_rms,
    }, "fields": {}, "telemetry": {}}
    with np.load(zero / "fields.npz", allow_pickle=False) as a, np.load(twin / "fields.npz", allow_pickle=False) as b:
        if not np.array_equal(a["lattice_steps"], b["lattice_steps"]):
            raise ValueError("baseline twin field timelines differ")
        for key in ("ux", "uy"):
            rms, relative = _relative_rms(a[key], b[key])
            report["fields"][key] = {"rms": rms, "relative_rms": relative,
                                      "max_abs": float(np.max(np.abs(a[key] - b[key])))}
    with np.load(zero / "telemetry.npz", allow_pickle=False) as a, np.load(twin / "telemetry.npz", allow_pickle=False) as b:
        if not np.array_equal(a["lattice_steps"], b["lattice_steps"]):
            raise ValueError("baseline twin telemetry timelines differ")
        for key in ("sensors", "forces"):
            rms, relative = _relative_rms(a[key], b[key])
            report["telemetry"][key] = {"rms": rms, "relative_rms": relative,
                                         "max_abs": float(np.max(np.abs(a[key] - b[key])))}
    field_max = max(value["relative_rms"] for value in report["fields"].values())
    telemetry_max = max(value["relative_rms"] for value in report["telemetry"].values())
    report.update({"field_max_relative_rms": field_max,
                   "telemetry_max_relative_rms": telemetry_max,
                   "passed": field_max <= max_field_relative_rms and telemetry_max <= max_telemetry_relative_rms})
    if not report["passed"]:
        raise RuntimeError(f"baseline twin numerical gate failed: {report}")
    return report
