#!/usr/bin/env python3
"""Tier 1 convergence exit check (Plan 5 Task 4).

Consumes the aggregated metrics JSON produced by
`scripts/aggregate-multi-seed-metrics.py` (Plan 5 Task 1B.2).

Pass criteria (per spec §2 Tier 1 — convergence):
  1. std/mean ≤ 0.15 on stable epochs (>= warmup_end) for `val_sharpe`,
     `avg_q_value`, `train_loss`.
  2. avg_q_value.max ≤ 500 across all epochs/streams (no fold-1 explosion;
     the fold-1 explosion bug from task #31 produced max_q ~ 1e4+).
  3. Q-saturation placeholder (see A.6 audit — full check landed elsewhere).
  4. Hot-path DtoH placeholder (enforced by pre-commit hook).

Exits 0 on PASS, 1 on FAIL. Each check prints a `PASS:` / `FAIL:` line.
"""
from __future__ import annotations

import argparse
import json
import sys
from statistics import mean


def _stable(series, warmup_end):
    """Return entries with epoch >= warmup_end."""
    return [x for x in series if x["epoch"] >= warmup_end]


def check_std_over_mean(agg, warmup_end):
    """std/mean ≤ 0.15 on stable epochs for the three Tier 1 metrics.

    Defensive: if a metric is missing from the aggregate, return False with
    a clear message so the run can be flagged for re-aggregation rather
    than silently passing.
    """
    for metric in ("val_sharpe", "avg_q_value", "train_loss"):
        series = agg.get(metric)
        if not series:
            return False, f"{metric} missing from aggregate JSON"
        stable = _stable(series, warmup_end)
        if not stable:
            return False, (
                f"No stable epochs for {metric} past warmup_end={warmup_end}"
            )
        m = mean(x["mean"] for x in stable)
        s = mean(x["std"] for x in stable)
        ratio = abs(s / (m + 1e-9)) if abs(m) > 1e-6 else float("inf")
        if ratio > 0.15:
            return False, (
                f"{metric} std/mean = {ratio:.3f} > 0.15 on stable epochs"
            )
    return True, "std/mean <= 0.15 on all Tier 1 metrics"


def check_no_fold1_explosion(agg):
    """avg_q_value.max ≤ 500 — catches the fold-1 loss explosion (task #31)."""
    series = agg.get("avg_q_value", [])
    if not series:
        # Defensive: missing series cannot prove there was no explosion;
        # failing is the safe direction.
        return False, "avg_q_value missing from aggregate JSON"
    max_q = max((x["max"] for x in series), default=0)
    if abs(max_q) > 500:
        return False, f"avg_q_value max={max_q}, possible fold-1 explosion"
    return True, "no fold-1 explosion detected"


def check_no_q_saturation(_agg):
    """Q-saturation placeholder.

    Full per-bin saturation check (q_full / q_half / q_quarter within 1%
    of ±abs_half) is tracked under the A.6 audit rather than embedded in
    the tier1 script — the relevant signal lives in HEALTH_DIAG `mag`
    block and is best validated end-to-end on a real run.
    """
    return True, "no Q-saturation check (placeholder — see A.6 audit)"


def check_no_hot_path_dtoh(_agg):
    """Hot-path DtoH enforcement is a pre-commit responsibility.

    `scripts/gpu-hotpath-guard.sh` rejects any new `DeviceCopy`-style
    DtoH inside the per-step training path. Reaching this script implies
    that guard already passed, so we surface the invariant here as PASS
    rather than re-deriving it from runtime metrics.
    """
    return True, "hot-path audit enforced by pre-commit"


def main():
    ap = argparse.ArgumentParser(
        description="Tier 1 convergence exit check (Plan 5 Task 4).",
    )
    ap.add_argument(
        "metrics_json",
        help="Aggregated metrics JSON (output of aggregate-multi-seed-metrics.py).",
    )
    ap.add_argument(
        "--warmup-end",
        type=int,
        default=15,
        help="Epoch at which warmup ends (stable epochs are > this); default 15.",
    )
    args = ap.parse_args()

    try:
        with open(args.metrics_json) as fh:
            data = json.load(fh)
    except (OSError, json.JSONDecodeError) as exc:
        print(f"FAIL: cannot read {args.metrics_json}: {exc}")
        sys.exit(1)

    agg = data.get("aggregates")
    if not isinstance(agg, dict):
        print("FAIL: missing or malformed `aggregates` block in metrics JSON")
        sys.exit(1)

    checks = [
        check_std_over_mean(agg, args.warmup_end),
        check_no_fold1_explosion(agg),
        check_no_q_saturation(agg),
        check_no_hot_path_dtoh(agg),
    ]

    all_pass = all(ok for ok, _ in checks)
    for ok, msg in checks:
        print(("PASS" if ok else "FAIL") + ": " + msg)
    sys.exit(0 if all_pass else 1)


if __name__ == "__main__":
    main()