"""Regression gate for the learning stages.

Re-scores every baked gain set with the CURRENT scoring code and compares it
to the score recorded when it was baked. Run it after any edit to
learn_balance.py, the robot spec, or the environment:

    python check_stages.py

Three findings it exists to catch, all of which happened in one evening:

  DRIFT      the same gains now score differently - the scorer, the model or
             the controller changed underneath them. The 7.5-scoring recover
             stage silently became a 0.67 one and nobody could say when.
  STALE      the baked file's parameter names no longer match the stage's
             current layout - loading it would scatter numbers into the wrong
             slots (stop's brake gain landed in recover's brace slot: a 25x
             overdose that put the robot down in under a second).
  FLAT       a stage scores identically across conditions that should differ.
             Flat is the signature of a harness bug, not of robustness: both
             times it appeared tonight the "different" conditions were never
             actually reaching the robot.
"""

import json
import sys
from pathlib import Path

import numpy as np

import learn_balance as LB

ROOT = Path(__file__).parent.parent


def main():
    ok = True
    for stage, cfg in LB.STAGES.items():
        f = ROOT / "runs" / f"{stage}_gains.json"
        if not f.exists():
            print(f"  {stage:9s} -        not baked yet")
            continue
        rec = json.load(open(f))
        names = rec.get("names") or []
        if names != cfg["names"]:
            ok = False
            print(f"  {stage:9s} STALE    baked layout {len(names)} names, "
                  f"stage now has {len(cfg['names'])} - do not load; relearn")
            continue
        sv = rec.get("score_version")
        if sv is not None and sv != LB.SCORE_VERSION:
            ok = False
            print(f"  {stage:9s} STALE    baked under score v{sv}, scorer is "
                  f"v{LB.SCORE_VERSION} - scores not comparable; relearn")
            continue
        g = np.array(cfg["seed"], dtype=float)
        g[:len(rec["gains"])] = rec["gains"]
        try:
            now = LB.stage_score(stage, g, secs=6.0)
        except Exception as e:
            ok = False
            print(f"  {stage:9s} ERROR    {type(e).__name__}: {e}")
            continue
        baked = rec.get("score")
        if baked is None:
            print(f"  {stage:9s} {now:7.2f}  (no baked score recorded)")
            continue
        # secs differ between bake and check, so allow real slack; what this
        # catches is collapse, not decimals
        drift = now - baked
        flag = "" if now > min(baked * 0.5, baked - 1.5) else "  DRIFT - investigate"
        if flag:
            ok = False
        print(f"  {stage:9s} {now:7.2f}  baked {baked:6.2f}  ({drift:+.2f}){flag}")
    print()
    print("all clear" if ok else "PROBLEMS FOUND - fix before trusting any new result")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())