"""
Safety / sanity gate for PCLC candidate models.

A model MUST pass this before it is exported for the app (population prior)
or shipped to a phone (per-user fine-tune). For electrical stimulation we never
let an unvetted model reach hardware.

Checks (over a dense grid of scaled inputs):
  1. All outputs finite (no NaN/Inf).
  2. Normalised outputs stay within [0,1].
  3. Decoded physical params stay within hard safety limits.
  4. Mean intensity is non-decreasing in pain  (sane dose-response).
  5. (fine-tune only) Mean output drift from the population prior is bounded
     — a user model should personalise, not mutate into something alien.
"""

import numpy as np

from pclc_model import D_IN, y_to_params


def _input_grid(n=7):
    # inputs are already scaled to [0,1]
    axes = [np.linspace(0.0, 1.0, n) for _ in range(D_IN)]
    mesh = np.array(np.meshgrid(*axes)).reshape(D_IN, -1).T
    return mesh.astype(np.float32)


def validate(model, out_cfg, prior_model=None, max_mean_drift=0.45, grid_n=7):
    """
    Returns (ok: bool, report: dict).
    Pass prior_model for fine-tuned models to enforce a bounded-drift check.
    """
    report = {}
    ok = True

    X = _input_grid(grid_n)
    Y = model.forward_np(X)

    # 1. finite
    if not np.all(np.isfinite(Y)):
        return False, {"error": "non-finite outputs"}

    # 2. normalised range
    report["y_min"], report["y_max"] = float(Y.min()), float(Y.max())
    if Y.min() < -1e-3 or Y.max() > 1.0 + 1e-3:
        ok = False
        report["range_fail"] = True

    # 3. physical hard limits
    P = np.array([y_to_params(y, out_cfg) for y in Y])  # (N, 5): freq,I,pw,heat,timer
    limits = [
        ("freq",  out_cfg["freqHardMin"], out_cfg["freqHardMax"]),
        ("I",     out_cfg["iHardMin"],    out_cfg["iHardMax"]),
        ("pw",    out_cfg["pwHardMin"],   out_cfg["pwHardMax"]),
        ("heat",  out_cfg["tHardMin"],    out_cfg["tHardMax"]),
        ("timer", out_cfg["tmHardMin"],   out_cfg["tmHardMax"]),
    ]
    for k, (name, lo, hi) in enumerate(limits):
        col = P[:, k]
        report[f"{name}_range"] = [float(col.min()), float(col.max())]
        if col.min() < lo - 1e-3 or col.max() > hi + 1e-3:
            ok = False
            report[f"{name}_limit_fail"] = True

    # 4. dose-response: mean intensity non-decreasing in pain (pain is input idx 3)
    pains = np.linspace(0.0, 1.0, 6)
    mean_I = []
    for p in pains:
        Xp = X.copy()
        Xp[:, 3] = p
        Pp = np.array([y_to_params(y, out_cfg) for y in model.forward_np(Xp)])
        mean_I.append(float(np.mean(Pp[:, 1])))
    report["mean_intensity_by_pain"] = mean_I
    # allow tiny numerical wiggle (-0.5 mA) but no real reversal
    if np.any(np.diff(mean_I) < -0.5):
        ok = False
        report["intensity_monotonic_fail"] = True

    # 5. bounded drift vs prior (fine-tune)
    if prior_model is not None:
        Yp = prior_model.forward_np(X)
        drift = float(np.mean(np.abs(Y - Yp)))
        report["mean_drift_from_prior"] = drift
        if drift > max_mean_drift:
            ok = False
            report["excessive_drift"] = True

    report["passed"] = ok
    return ok, report


if __name__ == "__main__":
    # Smoke test on a fresh (untrained) model — exercises the gate plumbing.
    import argparse
    from pclc_model import PCLCPolicy, load_scaler

    ap = argparse.ArgumentParser()
    ap.add_argument("--scaler", default="scaler.json")
    args = ap.parse_args()

    _, _, cfg = load_scaler(args.scaler)
    ok, rep = validate(PCLCPolicy(), cfg)
    print("passed:", ok)
    for k, v in rep.items():
        print(f"  {k}: {v}")