#!/usr/bin/env python3
"""Compare every model across the four evaluation splits.

    python analysis/dataset_level_analysis/analyze_dataset_level.py

Reads result/<model>/<task>/baseline_<b>/{metrics.json,predictions.jsonl} and
writes the CSVs next to this script.

Raw accuracy is not comparable across the four splits: VSR is a two-way
true/false question while the three Youtube levels are four-way, and level 2
carries a structural floor of its own (see FLOORS).  Everything is therefore
also reported as *headroom captured*, (acc - floor) / (1 - floor): the share of
the distance between guessing and perfect that the model actually covers.  It
is negative when a model does worse than the floor, which happens here.
"""
from __future__ import annotations

import argparse
import collections
import csv
import json
import math
import os
import statistics

BASELINES = ["0.00", "0.05", "0.10", "0.15", "0.20", "0.25"]
TASKS = ["vsr", "youtube_level1", "youtube_level2", "youtube_level3"]

# What a model scores on each split without looking at the image: pure guessing
# among the choices.  Level 2's three distractors are drawn uniformly from the
# 23 wrong orderings, so the four options are exchangeable and every text-only
# heuristic sits at chance -- measured on the split: prefix vote 0.252, medoid
# 0.256, outlier 0.254, positional majority 0.245.  Earlier structured variants
# did NOT have this property; see disparity_qa/level2_choices.py.
FLOORS = {"vsr": 0.500, "youtube_level1": 0.250,
          "youtube_level2": 0.250, "youtube_level3": 0.250}
CHANCE = dict(FLOORS)


def load_metrics(root: str, model: str, task: str, baseline: str) -> dict:
    with open(os.path.join(root, "result", model, task,
                           f"baseline_{baseline}", "metrics.json")) as fh:
        return json.load(fh)


def load_predictions(root: str, model: str, task: str, baseline: str):
    path = os.path.join(root, "result", model, task,
                        f"baseline_{baseline}", "predictions.jsonl")
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if line:
                yield json.loads(line)


def entropy(counter: collections.Counter) -> float:
    n = sum(counter.values())
    if not n:
        return 0.0
    return -sum((v / n) * math.log2(v / n) for v in counter.values() if v)


def slope(xs, ys) -> float:
    """Least-squares slope of accuracy against baseline (per unit of baseline)."""
    mx, my = statistics.mean(xs), statistics.mean(ys)
    den = sum((x - mx) ** 2 for x in xs)
    return sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / den if den else 0.0


def main() -> int:
    here = os.path.dirname(os.path.abspath(__file__))
    p = argparse.ArgumentParser()
    p.add_argument("--root", default=os.path.abspath(os.path.join(here, "..", "..")))
    p.add_argument("--out_dir", default=here)
    args = p.parse_args()

    models = sorted(m for m in os.listdir(os.path.join(args.root, "result"))
                    if not m.startswith("_")
                    and os.path.isdir(os.path.join(args.root, "result", m)))

    acc = {}          # (model, task, baseline) -> metrics dict
    for m in models:
        for t in TASKS:
            for b in BASELINES:
                acc[(m, t, b)] = load_metrics(args.root, m, t, b)

    # ---------------------------------------------------------- per baseline
    rows = []
    for m in models:
        for t in TASKS:
            for b in BASELINES:
                d = acc[(m, t, b)]
                rows.append({
                    "model": m, "task": t, "baseline": b,
                    "n": d["num_samples"], "accuracy": round(d["accuracy"], 4),
                    "num_correct": d["num_correct"],
                    "unparsed_rate": round(d["unparsed_rate"], 4),
                    "scoring": d["scoring"],
                    "chance": CHANCE[t], "floor": FLOORS[t],
                    "headroom_captured": round(
                        (d["accuracy"] - FLOORS[t]) / (1 - FLOORS[t]), 4),
                })
    write_csv(os.path.join(args.out_dir, "accuracy_by_baseline.csv"), rows)

    # ------------------------------------------------------- per model/task
    summary = []
    for m in models:
        for t in TASKS:
            a = [acc[(m, t, b)]["accuracy"] for b in BASELINES]
            xs = [float(b) for b in BASELINES]
            summary.append({
                "model": m, "task": t, "n": acc[(m, t, "0.00")]["num_samples"],
                "scoring": acc[(m, t, "0.00")]["scoring"],
                "chance": CHANCE[t], "floor": FLOORS[t],
                "acc_b0": round(a[0], 4), "acc_mean": round(statistics.mean(a), 4),
                "acc_min": round(min(a), 4), "acc_max": round(max(a), 4),
                "acc_range": round(max(a) - min(a), 4),
                "acc_std": round(statistics.pstdev(a), 4),
                "slope_per_baseline": round(slope(xs, a), 4),
                "delta_b025_minus_b0": round(a[-1] - a[0], 4),
                "headroom_captured_mean": round(
                    (statistics.mean(a) - FLOORS[t]) / (1 - FLOORS[t]), 4),
                "above_floor": statistics.mean(a) > FLOORS[t],
            })
    write_csv(os.path.join(args.out_dir, "summary_by_model_task.csv"), summary)

    # --------------------------------------------------------------- biases
    bias = []
    for m in models:
        for t in TASKS:
            for b in BASELINES:
                pred = collections.Counter()
                ref = collections.Counter()
                n = 0
                for r in load_predictions(args.root, m, t, b):
                    pred[r["prediction_norm"]] += 1
                    ref[r["reference_norm"]] += 1
                    n += 1
                top_label, top_n = pred.most_common(1)[0]
                # what a model would score by always naming its own modal answer
                majority_acc = ref[top_label] / n if n else 0.0
                bias.append({
                    "model": m, "task": t, "baseline": b, "n": n,
                    "n_distinct_predictions": len(pred),
                    "modal_prediction": top_label,
                    "modal_share": round(top_n / n, 4),
                    "prediction_entropy_bits": round(entropy(pred), 4),
                    "reference_entropy_bits": round(entropy(ref), 4),
                    "always_modal_accuracy": round(majority_acc, 4),
                    "accuracy": round(acc[(m, t, b)]["accuracy"], 4),
                    "prediction_counts": json.dumps(dict(sorted(pred.items()))),
                })
    write_csv(os.path.join(args.out_dir, "prediction_bias.csv"), bias)

    ladder = level2_ladder(args.root, models)
    if ladder:
        write_csv(os.path.join(args.out_dir, "level2_structure_use.csv"), ladder)

    report(models, acc, summary, bias, ladder)
    print(f"\nwrote {3 + bool(ladder)} CSVs to {args.out_dir}")
    return 0


# the split lives in the repo, one level above this experiment directory
LEVEL2_TEST = os.path.join("..", "02_data", "Youtube_self_depth_QA",
                           "level2", "test.jsonl")


def _prefix_survivors(options):
    """The options a majority vote on successive positions cannot rule out.

    A regression guard.  With random distractors this should leave most items
    undecided and score at chance; if it ever climbs toward 0.5 again, the
    option construction has reacquired a structural leak.
    """
    alive = list(options)
    seqs = {lab: txt.split(" > ") for lab, txt in options.items()}
    for k in range(4):
        if len(alive) <= 1:
            break
        counts = collections.Counter(seqs[a][k] for a in alive)
        top = counts.most_common(1)[0][1]
        winners = [l for l, v in counts.items() if v == top]
        if len(winners) == 1:
            alive = [a for a in alive if seqs[a][k] == winners[0]]
    return alive


def level2_ladder(root: str, models):
    """How often each model's answer lands inside the structure-only survivor set."""
    path = os.path.abspath(os.path.join(root, LEVEL2_TEST))
    if not os.path.exists(path):
        print(f"[warn] level-2 split not found at {path}; skipping ladder analysis")
        return []
    test = {}
    with open(path) as fh:
        for line in fh:
            rec = json.loads(line)
            if rec.get("options"):
                test[rec["sample_id"]] = rec["options"]
    if not test:
        return []
    out = []
    for m in models:
        for b in BASELINES:
            inside = hit = n = 0
            for r in load_predictions(root, m, "youtube_level2", b):
                opts = test.get(r["sample_id"])
                if not opts:
                    continue
                alive = _prefix_survivors(opts)
                n += 1
                inside += int(r["prediction_norm"] in alive)
                hit += int(r["correct"])
            out.append({
                "model": m, "baseline": b, "n": n,
                "answer_in_survivor_set": round(inside / n, 4) if n else 0.0,
                "accuracy": round(hit / n, 4) if n else 0.0,
                # of the answers that did land in the survivor set, how many were
                # right: this is the coin flip the structure cannot decide
                "accuracy_within_survivors": round(hit / inside, 4) if inside else 0.0,
            })
    return out


def write_csv(path: str, rows) -> None:
    with open(path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)


def report(models, acc, summary, bias, ladder=()) -> None:
    short = {m: m.replace("Qwen_Qwen", "").replace("-Instruct", "") for m in models}
    idx = {(r["model"], r["task"]): r for r in summary}

    print("=" * 96)
    print("ACCURACY  (mean over the six baselines; floor = what the split gives "
          "away for free)")
    print("=" * 96)
    print(f'{"model":22s}' + "".join(f"{t.replace('youtube_','yt_'):>16s}" for t in TASKS))
    print(f'{"floor":22s}' + "".join(f"{FLOORS[t]:>16.3f}" for t in TASKS))
    print("-" * 96)
    for m in models:
        print(f"{short[m]:22s}" + "".join(
            f'{idx[(m,t)]["acc_mean"]:>16.3f}' for t in TASKS))
    print()
    print("HEADROOM CAPTURED  (acc - floor) / (1 - floor);  negative = below floor")
    print("-" * 96)
    for m in models:
        cells = []
        for t in TASKS:
            v = idx[(m, t)]["headroom_captured_mean"]
            cells.append(f"{v:>+15.3f} " if v < 0 else f"{v:>16.3f}")
        print(f"{short[m]:22s}" + "".join(cells))
    print()
    print("BASELINE SENSITIVITY  range over 0.00-0.25  /  slope per unit baseline")
    print("-" * 96)
    for m in models:
        print(f"{short[m]:22s}" + "".join(
            f'{idx[(m,t)]["acc_range"]:>7.3f}/{idx[(m,t)]["slope_per_baseline"]:>+8.3f}'
            for t in TASKS))
    print()
    print("ANSWER COLLAPSE  modal predicted label, its share, and what always "
          "answering it would score (baseline 0.00)")
    print("-" * 96)
    bidx = {(r["model"], r["task"], r["baseline"]): r for r in bias}
    for m in models:
        cells = []
        for t in TASKS:
            r = bidx[(m, t, "0.00")]
            cells.append(f'{r["modal_prediction"]:>7s}{r["modal_share"]:>6.2f}'
                         f'{r["always_modal_accuracy"]:>7.3f}')
        print(f"{short[m]:22s}" + "".join(cells))
    print("  (reference labels are near-uniform, so 'always modal' sits at chance)")

    if ladder:
        print()
        print("LEVEL 2 REGRESSION GUARD -- can a text-only majority vote still help?")
        print("-" * 96)
        lidx = {(r["model"], r["baseline"]): r for r in ladder}
        for m in models:
            r = lidx[(m, "0.00")]
            print(f'{short[m]:22s} answer inside the vote survivors '
                  f'{r["answer_in_survivor_set"]:.3f}   accuracy {r["accuracy"]:.3f}')
        print("  distractors are random, so the vote rarely decides anything; a")
        print("  survivor rate near 1.0 means it eliminated nothing.")


if __name__ == "__main__":
    raise SystemExit(main())