"""Compiler v2: measure -> route -> ESCAPE-AND-SELECT. No human decisions.

Two changes from v1, both forced by the v2 joint-anneal test
(RESULTS_compiler_v2_joint.md):
1. Routing reads the trainless R-grid and, for a floored cell, picks the nearest
   ancestor with R >= 2*th by 1-axis moves, else 2-axis (JOINT) moves. This
   handles double-deficit cells (small AND low-contrast) that v1's single-axis
   router mis-prescribed.
2. The wall is BIMODAL, not a threshold: a curriculum run either escapes to the
   real solution or collapses to the uniform-guess basin (final loss == log K,
   acc ~ 1/K), seed-determined. So run K seeds per cell and SELECT the escaper.
   Collapse is detectable with zero eval as final_loss >= log(K_classes) - eps.
"""
import subprocess, json, os, sys, math, glob
sys.path.insert(0, ".")
TH = 0.02          # in-family learnability threshold
MARGIN = 2 * TH    # strong-ancestor margin (v1 lesson)
KSEEDS = 5         # restarts per cell. Measured p_escape~=0.45 for the hardest
                   # (double-deficit) cell -> K=5 gives P(all collapse)=0.55^5~=0.05.
                   # K=3 (0.166) is too weak near the wall. Better: adaptive —
                   # stop early on first escape (loss<log K), cap at K_MAX=8.
NUM_CLASSES = 4    # age buckets -> collapse loss ~ log(4) = 1.386
COLLAPSE_ACC = 0.45
SIZES = [5, 6, 8, 10, 12, 14, 16]

R = json.load(open("results_compiler_v2/R_grid.json"))  # trainless surface
def r(bg, s, ct): return R.get(f"{bg}_o{s}_ct{ct:g}", 0.0)

def route(bg, size, ct):
    """Return (mode, anneal_from, contrast_from). Trainless — reads R only."""
    if r(bg, size, ct) >= TH:
        return ("direct", 0, 0)
    # 1-axis: bigger size at target contrast
    for s2 in [s for s in SIZES if s > size]:
        if r(bg, s2, ct) >= MARGIN:
            return ("size", s2, 0)
    # 1-axis: higher contrast at target size
    if ct < 1.0 and r(bg, size, 1.0) >= MARGIN:
        return ("contrast", 0, 1.0)
    # 2-axis: bigger size AND full contrast (JOINT)
    for s2 in [s for s in SIZES if s > size]:
        if ct < 1.0 and r(bg, s2, 1.0) >= MARGIN:
            return ("joint", s2, 1.0)
    return ("unreachable", 0, 0)

def train(bg, size, ct, anc, cfrom, seed, out):
    cmd = ["python3", "-m", "simreal.train_composite", "--steps", "20000",
           "--objscale", str(size), "--contrast", str(ct), "--seed", str(seed),
           "--out", out]
    if bg == "static": cmd.append("--static-bg")
    if anc:   cmd += ["--anneal-from", str(anc)]
    if cfrom: cmd += ["--anneal-contrast-from", str(cfrom)]
    subprocess.run(cmd, check=False)
    # read back this seed's result
    pat = f"{out}/simreal{'_static' if bg=='static' else ''}_o{size}_ct{ct:g}_off0_*_an{anc}c{cfrom:g}_s{seed}.json"
    hits = glob.glob(pat)
    if not hits: return None
    return json.load(open(hits[0]))["heldout_mean"]

def compile_cell(bg, size, ct, out="results_compiler_v2/run"):
    os.makedirs(out, exist_ok=True)
    mode, anc, cfrom = route(bg, size, ct)
    print(f"ROUTE {bg} {size}px ct{ct}: {mode} (anneal_from={anc}, contrast_from={cfrom})", flush=True)
    if mode == "unreachable":
        return {"cell": (bg, size, ct), "mode": mode, "best": None, "escaped": 0}
    if mode == "direct":
        anc, cfrom = 0, 0
    scores = []
    for seed in range(KSEEDS):
        m = train(bg, size, ct, anc, cfrom, seed, out)
        scores.append(m)
        print(f"  seed {seed}: {m} {'ESCAPE' if (m or 0) >= COLLAPSE_ACC else 'collapse'}", flush=True)
    best = max([s for s in scores if s is not None], default=None)
    escaped = sum(1 for s in scores if (s or 0) >= COLLAPSE_ACC)
    print(f"  -> best {best} over {KSEEDS} seeds ({escaped} escaped)", flush=True)
    return {"cell": (bg, size, ct), "mode": mode, "scores": scores, "best": best, "escaped": escaped}

if __name__ == "__main__":
    # the six campaign floors; last is the double-deficit holdout
    TARGETS = [("static", 5, 1.0), ("static", 6, 1.0), ("moving", 5, 1.0),
               ("moving", 6, 1.0), ("moving", 10, 0.4), ("static", 10, 0.4)]
    report = [compile_cell(*t) for t in TARGETS]
    json.dump(report, open("results_compiler_v2/compiler_v2_report.json", "w"), indent=2)
    rescued = sum(1 for r_ in report if (r_["best"] or 0) >= 0.55)
    print(f"\nCOMPILER_V2: {rescued}/{len(TARGETS)} cells rescued (best-of-{KSEEDS})", flush=True)