"""CMB analysis. Unit of analysis is the item; replicates averaged within an item first."""
import json
import math
import os
import sys
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import views as V
import items as IT

R = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                  "..", "results"))

ARM_LABEL = {
    "wiki": "A shared wiki", "vector_memory": "B vector memory + timestamps",
    "bitemporal": "C bi-temporal graph", "graph_catalog": "D graph + governance catalogue",
    "dimension": "E dimension", "dim_no_gov": "F dimension minus governance",
}


def load(name):
    p = os.path.join(R, name)
    if not os.path.exists(p):
        return []
    rows = [json.loads(l) for l in open(p, encoding="utf-8")]
    if name == "grid.jsonl":
        # The first grid run asked three historical items for a version number, a field
        # only some representations carry. Those items were reworded to ask for the
        # content in force and re-run; the corrected run replaces the family here, and
        # both files are published.
        fixed = os.path.join(R, "grid-hist-v2.jsonl")
        if os.path.exists(fixed):
            rows = [r for r in rows if r["family"] != "HIST"]
            rows += [json.loads(l) for l in open(fixed, encoding="utf-8")]
        import run as RUN
        by = {i["id"]: i for i in IT.I}
        for r in rows:
            r.update(RUN.score(by[r["item"]], r.get("obj")))
    return rows


def item_means(rows, arm, metric, gate, family=None):
    acc = defaultdict(list)
    for r in rows:
        if r["arm"] != arm or not r.get(gate):
            continue
        if family and r["family"] != family:
            continue
        acc[r["item"]].append(r[metric])
    return {k: sum(v) / len(v) for k, v in acc.items()}


def mean(d):
    return sum(d.values()) / len(d) if d else None


def cluster_ci(d):
    n = len(d)
    if n < 2:
        return (None, None)
    vals = list(d.values())
    m = sum(vals) / n
    var = sum((v - m) ** 2 for v in vals) / (n - 1)
    se = math.sqrt(var / n)
    return (max(0.0, m - 1.96 * se), min(1.0, m + 1.96 * se))


def sign_test(a, b):
    keys = sorted(set(a) & set(b))
    pos = sum(1 for k in keys if b[k] > a[k])
    neg = sum(1 for k in keys if b[k] < a[k])
    n = pos + neg
    if n == 0:
        return 1.0, 0, 0
    k = min(pos, neg)
    return round(min(1.0, 2 * sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n), 5), \
        pos, neg


def pct(x):
    return "-" if x is None else f"{100 * x:5.1f}%"


def bar(t):
    print("\n" + "=" * 104)
    print(t)
    print("=" * 104)


def main():
    grid = load("grid.jsonl")
    if not grid:
        raise SystemExit("no grid.jsonl")

    bar(f"CMB grid: {len(grid)} calls, {len(set(r['item'] for r in grid))} items, "
        f"{len(set(r['arm'] for r in grid))} arms")
    print(f"parse success: {100*sum(r['parsed'] for r in grid)/len(grid):.1f}%")

    print(f"\n{'arm':34} {'answer':>8} {'95% CI':>12} {'owner':>8} {'policy':>8} "
          f"{'source':>8} {'abstain':>8} {'stale':>8} {'rival':>8} {'fabric':>8}")
    print("-" * 122)
    summary = {}
    for a in V.ARMS:
        m1 = item_means(grid, a, "answer_correct", "scored_answer")
        lo, hi = cluster_ci(m1)
        summary[a] = {
            "answer": mean(m1), "ci": (lo, hi),
            "owner": mean(item_means(grid, a, "owner_correct", "scored_owner")),
            "policy": mean(item_means(grid, a, "policy_correct", "scored_policy")),
            "source": mean(item_means(grid, a, "source_correct", "scored_source")),
            "abstain": mean(item_means(grid, a, "abstained", "scored_abst")),
            "false_abstain": mean(item_means(grid, a, "false_abstention",
                                             "scored_answer")),
            "stale": mean(item_means(grid, a, "stale_correct", "scored_stale")),
            "rival": mean(item_means(grid, a, "overclaim", "scored_answer")),
            "fabric": mean(item_means(grid, a, "fabricated", "scored_answer")),
            "context_chars": len(V.context_for(a)),
        }
        s = summary[a]
        ci = f"[{100*lo:.0f}-{100*hi:.0f}]" if lo is not None else "-"
        print(f"{ARM_LABEL[a]:34} {pct(s['answer']):>8} {ci:>12} {pct(s['owner']):>8} "
              f"{pct(s['policy']):>8} {pct(s['source']):>8} {pct(s['abstain']):>8} "
              f"{pct(s['stale']):>8} {pct(s['rival']):>8} {pct(s['fabric']):>8}")

    bar("Answer accuracy by family")
    fams = IT.FAMILIES
    print(f"{'arm':34} " + " ".join(f"{f:>8}" for f in fams))
    print("-" * 122)
    byfam = {}
    for a in V.ARMS:
        byfam[a] = {}
        row = []
        for f in fams:
            metric, gate = ("abstained", "scored_abst") if f == "GAP" else \
                           ("answer_correct", "scored_answer")
            d = item_means(grid, a, metric, gate, family=f)
            byfam[a][f] = mean(d)
            row.append(pct(mean(d)))
        print(f"{ARM_LABEL[a]:34} " + " ".join(f"{v:>8}" for v in row))

    bar("Paired comparisons on answer accuracy (exact sign test over item means)")
    pairs = [("wiki", "dimension"), ("vector_memory", "dimension"),
             ("bitemporal", "dimension"), ("graph_catalog", "dimension"),
             ("dim_no_gov", "dimension"), ("bitemporal", "graph_catalog"),
             ("bitemporal", "dim_no_gov"), ("vector_memory", "bitemporal"),
             ("wiki", "graph_catalog")]
    tests = {}
    for a, b in pairs:
        da = item_means(grid, a, "answer_correct", "scored_answer")
        db = item_means(grid, b, "answer_correct", "scored_answer")
        p, pos, neg = sign_test(da, db)
        tests[f"{a}_vs_{b}"] = {"p": p, "better": pos, "worse": neg}
        print(f"{ARM_LABEL[a]:34} -> {ARM_LABEL[b]:34} {pct(mean(da))} -> "
              f"{pct(mean(db))}  better:{pos:2} worse:{neg:2}  p={p}")

    bar("The governance ablation, with its control family")
    abl = {}
    for fam in ("AUTH", "CONF", "ROUTE", "HIST"):
        da = item_means(grid, "dim_no_gov", "answer_correct", "scored_answer",
                        family=fam)
        db = item_means(grid, "dimension", "answer_correct", "scored_answer",
                        family=fam)
        p, _pos, _neg = sign_test(da, db)
        abl[fam] = {"ablated": mean(da), "full": mean(db), "p": p}
        role = "control" if fam == "HIST" else "target"
        print(f"{fam:6} ({role:7}) removing governance fields: {pct(mean(da))} vs "
              f"{pct(mean(db))}  p={p}")

    bar("Fidelity check: the two independent ablations of the same field classes")
    dc = item_means(grid, "bitemporal", "answer_correct", "scored_answer")
    df = item_means(grid, "dim_no_gov", "answer_correct", "scored_answer")
    p, pos, neg = sign_test(dc, df)
    print(f"C bi-temporal graph {pct(mean(dc))} vs F dimension minus governance "
          f"{pct(mean(df))}  better:{pos} worse:{neg}  p={p}")
    print("If the reimplementation of the competitor shape is fair, these should agree.")

    scope = load("scope.jsonl")
    scope_out = {}
    if scope:
        bar("Cross-team scope track: 30 release decisions per arm per replicate")
        print(f"{'arm':34} {'precision':>10} {'recall':>9} {'accuracy':>10} "
              f"{'critical over-shares':>21}")
        print("-" * 122)
        for a in V.ARMS:
            s = [r for r in scope if r["arm"] == a]
            if not s:
                continue
            tp = sum(r["tp"] for r in s); fp = sum(r["fp"] for r in s)
            tn = sum(r["tn"] for r in s); fn = sum(r["fn"] for r in s)
            tot = tp + fp + tn + fn
            crit = sum(r["critical"] for r in s)
            scope_out[a] = {"precision": tp / (tp + fp) if tp + fp else None,
                            "recall": tp / (tp + fn) if tp + fn else None,
                            "accuracy": (tp + tn) / tot if tot else None,
                            "critical": crit, "decisions": tot}
            o = scope_out[a]
            print(f"{ARM_LABEL[a]:34} {pct(o['precision']):>10} {pct(o['recall']):>9} "
                  f"{pct(o['accuracy']):>10} {str(crit) + ' of ' + str(tot):>21}")

    rest = load("restate.jsonl")
    rest_out = {}
    if rest:
        bar("Restatement track: one team restates another team's record")
        print(f"{'arm':34} {'value':>8} {'from':>8} {'by team':>9} {'closed id':>10} "
              f"{'history':>9} {'others kept':>12} {'all three':>10}")
        print("-" * 122)
        for a in V.ARMS:
            s = [r for r in rest if r["arm"] == a]
            if not s:
                continue
            n = len(s)
            rest_out[a] = {k: sum(r[k] for r in s) / n for k in
                           ("value_ok", "from_ok", "team_ok", "close_ok", "history_ok",
                            "other_ok", "well_formed")}
            rest_out[a]["chains"] = n
            o = rest_out[a]
            print(f"{ARM_LABEL[a]:34} {pct(o['value_ok']):>8} {pct(o['from_ok']):>8} "
                  f"{pct(o['team_ok']):>9} {pct(o['close_ok']):>10} "
                  f"{pct(o['history_ok']):>9} {pct(o['other_ok']):>12} "
                  f"{pct(o['well_formed']):>10}")

    bar("Context size")
    for a in V.ARMS:
        print(f"{ARM_LABEL[a]:34} {summary[a]['context_chars']:6} chars")

    out = {"summary": summary, "by_family": byfam, "tests": tests, "ablation": abl,
           "fidelity": {"bitemporal": mean(dc), "dim_no_gov": mean(df), "p": p},
           "scope": scope_out, "restate": rest_out,
           "calls": {"grid": len(grid), "scope": len(scope), "restate": len(rest)}}
    dest = os.path.join(R, "analysis.json")
    with open(dest, "w", encoding="utf-8") as fh:
        json.dump(out, fh, indent=2, ensure_ascii=False, default=str)
    print("\nwritten:", dest)


if __name__ == "__main__":
    main()