"""Output checks shared by every problem's reference.py.

compare(out, ref, tol) passes when:
- the whole output's relative L2 error, |out - ref| / |ref|, is at most tol;
- every row's is at most ROW_FACTOR * tol. A row is the last dimension. A row's error is
  taken against its own norm, or against ROW_NORM_FLOOR times the average row norm if
  that's larger.
"""

from __future__ import annotations

import torch

ROW_FACTOR = 5.0
# A row's error is measured against its own norm, but not below this fraction of the
# typical row norm, so rows that are legitimately near zero can't blow the ratio up.
ROW_NORM_FLOOR = 0.1


def relative_l2(a: torch.Tensor, b: torch.Tensor) -> float:
    """|a - b| / |b|."""
    return float((a.float() - b.float()).norm() / b.float().norm().clamp(min=1e-12))


def query_precision_bound(e4m3_err: float, factor: float, floor: float, cap: float, slack: float = 0.0) -> float:
    """Attention's bound for one case: `factor` times the reference's own error with queries
    rounded to FP8 E4M3 (`e4m3_err`), plus `slack`, within [floor, cap]."""
    return min(cap, max(floor, factor * e4m3_err + slack))


def compare(output, ref: torch.Tensor, tol: float) -> tuple[bool, str]:
    if not isinstance(output, torch.Tensor):
        return False, f"expected a tensor, got {type(output).__name__}"
    if output.shape != ref.shape or output.dtype != ref.dtype:
        return False, f"expected {tuple(ref.shape)} {ref.dtype}, got {tuple(output.shape)} {output.dtype}"
    out = output.float()
    if not torch.isfinite(out).all():
        return False, "output contains NaN or Inf"
    r = ref.float()
    diff = out - r
    err = float(diff.norm() / r.norm().clamp(min=1e-12))
    if err > tol:
        return False, f"relative L2 error {err:.3e} > {tol:g}"

    rows_diff = diff.reshape(-1, r.shape[-1]).norm(dim=-1)
    rows_ref = r.reshape(-1, r.shape[-1]).norm(dim=-1)
    floor = (ROW_NORM_FLOOR * rows_ref.mean()).clamp(min=1e-12)
    row_err = rows_diff / torch.maximum(rows_ref, floor)
    worst = int(row_err.argmax())
    worst_err = float(row_err[worst])
    if worst_err > ROW_FACTOR * tol:
        index = tuple(int(i) for i in torch.unravel_index(torch.tensor(worst), r.shape[:-1])) if r.dim() > 1 else ()
        return False, (f"relative L2 error {err:.3e}, but row {index} is off by {worst_err:.3e}"
                       f" > {ROW_FACTOR:g} x {tol:g}")
    return True, f"relative L2 error {err:.3e} (worst row {worst_err:.3e})"