#!/usr/bin/env python3
"""Re-process the raw responses from a finished run with fixes:
- Per-call temp dir so file stem matches module name
- Tag-disjoint counting (only count tags unique to a cohort,
  ignore overlapping ones like `lambda` shared by apply+call)
"""
from __future__ import annotations
import json, re, subprocess, sys
from pathlib import Path

RUN_DIR = Path(sys.argv[1])
RESULTS = json.loads((RUN_DIR / "results.json").read_text())

REVERSE_APPLY = [("tail-apply","tail-app"),("term-constructor","term-ctor"),
                 ("pat-constructor","pat-ctor"),("constructor","ctor"),
                 ("apply","app"),("lambda","lam"),("typecon","con")]
REVERSE_CALL = [("tail-call","tail-app"),("term-constructor","term-ctor"),
                ("pat-constructor","pat-ctor"),("constructor","ctor"),
                ("call","app"),("lambda","lam"),("typecon","con")]

def reverse_sub(src, pairs):
    out = src
    for n,o in pairs:
        out = re.sub(rf"\({re.escape(n)}(?=[\s)])", f"({o}", out)
    return out

# Disjoint discriminator tags only — the head of function application
# (which IS the test variable). Ignore lambda/constructor/typecon since
# apply and call share those.
DISCRIM = {
    "classic": ["app","tail-app"],
    "apply":   ["apply","tail-apply"],
    "call":    ["call","tail-call"],
}

def count_tags(src, tags):
    return {t: len(re.findall(rf"\({re.escape(t)}(?=[\s)])", src)) for t in tags}

def extract_module_name(code):
    m = re.search(r"\(module\s+([A-Za-z_][A-Za-z0-9_]*)", code)
    return m.group(1) if m else None

def run_pipeline_fixed(code, module_name, expected_stdout, base):
    work = base / module_name
    work.mkdir(parents=True, exist_ok=True)
    src = work / f"{module_name}.ail"
    src.write_text(code)
    out = {"check_ok": False, "check_stderr": "", "build_ok": False,
           "run_ok": False, "run_stdout": "", "matches_expected": False}
    c = subprocess.run(["ail","check",str(src)], capture_output=True, text=True, timeout=30)
    out["check_stderr"] = (c.stderr or c.stdout)[:2000]
    if c.returncode != 0:
        return out
    out["check_ok"] = True
    b = subprocess.run(["ail","build",str(src),"-o",str(work/"bin")],
                       capture_output=True, text=True, timeout=60)
    if b.returncode != 0:
        out["check_stderr"] = b.stderr[:2000]
        return out
    out["build_ok"] = True
    r = subprocess.run([str(work/"bin")], capture_output=True, text=True, timeout=10)
    if r.returncode == 0:
        out["run_ok"] = True
        out["run_stdout"] = r.stdout
        out["matches_expected"] = (r.stdout == expected_stdout)
    return out

# need expected stdout per task
TASKS = {}
CMA_TASKS = Path("/home/brummel/dev/ailang/experiments/2026-05-12-cross-model-authoring/master/tasks")
for f in CMA_TASKS.glob("*.task.json"):
    t = json.loads(f.read_text())
    TASKS[t["id"]] = t["expected_stdout"]

base = RUN_DIR / "work2"
base.mkdir(exist_ok=True)

new_rows = []
for r in RESULTS:
    code = r["extracted_code"]
    cohort = r["cohort"]
    pairs = {"classic":None,"apply":REVERSE_APPLY,"call":REVERSE_CALL}[cohort]
    pipeline_input = reverse_sub(code, pairs) if pairs else code
    mod_name = extract_module_name(pipeline_input) or r["task_id"]
    pipe = run_pipeline_fixed(pipeline_input, mod_name, TASKS[r["task_id"]], base)

    own_disc = count_tags(code, DISCRIM[cohort])
    foreign_disc = {}
    for c, tags in DISCRIM.items():
        if c == cohort: continue
        foreign_disc[c] = count_tags(code, tags)

    new_rows.append({
        "cohort": cohort, "task_id": r["task_id"], "trial": r["trial"],
        "module_name_in_output": mod_name,
        "module_name_expected": r["task_id"],
        "module_match": (mod_name == r["task_id"]),
        "own_disc": own_disc,
        "foreign_disc": foreign_disc,
        "check_ok": pipe["check_ok"],
        "build_ok": pipe["build_ok"],
        "run_ok": pipe["run_ok"],
        "matches_expected": pipe["matches_expected"],
        "run_stdout": pipe["run_stdout"],
        "first_err_line": pipe["check_stderr"].split("\n")[0][:140],
    })

(RUN_DIR / "results2.json").write_text(json.dumps(new_rows, indent=2))

print("\n=== final pipeline pass per cohort ===")
by_cohort = {}
for r in new_rows:
    c = by_cohort.setdefault(r["cohort"], {"n":0,"mod":0,"check":0,"build":0,"run":0,"correct":0})
    c["n"] += 1
    if r["module_match"]: c["mod"] += 1
    if r["check_ok"]: c["check"] += 1
    if r["build_ok"]: c["build"] += 1
    if r["run_ok"]: c["run"] += 1
    if r["matches_expected"]: c["correct"] += 1
for cohort, c in by_cohort.items():
    print(f"{cohort:8s}  n={c['n']}  mod-name-OK={c['mod']}  check={c['check']}  build={c['build']}  run={c['run']}  stdout-correct={c['correct']}")

print("\n=== discriminator tag totals (head of function application only) ===")
disc_totals = {}
for r in new_rows:
    c = r["cohort"]
    own_sum = sum(r["own_disc"].values())
    foreign_sum = sum(sum(v.values()) for v in r["foreign_disc"].values())
    t = disc_totals.setdefault(c, {"own":0,"foreign":0,"by_foreign":{}})
    t["own"] += own_sum
    for fc, fv in r["foreign_disc"].items():
        t["by_foreign"][fc] = t["by_foreign"].get(fc,0) + sum(fv.values())
    t["foreign"] += foreign_sum
for cohort, t in disc_totals.items():
    print(f"{cohort:8s}  own={t['own']:3d}  foreign={t['foreign']:3d}  breakdown={t['by_foreign']}")

print("\n=== per-row failure mode ===")
for r in new_rows:
    flag = "OK" if r["matches_expected"] else ("CORRECT-NOT" if r["run_ok"] else ("RUN" if r["build_ok"] else ("BUILD" if r["check_ok"] else "CHECK")))
    print(f"[{r['cohort']:8s}] {r['task_id']:18s} t{r['trial']}  {flag:12s}  {r['first_err_line']}")