"""Cross-baseline stability and failure-case overlap.

Reads the per-baseline prediction files written by inference/run_inference.py
and, for each (model, task), lines the samples up by sample_id to answer:

  accuracy        per baseline, plus how much it moves across the sweep.
  flip rate       fraction of samples whose ANSWER changes relative to the
                  reference baseline (default 0.00). This is the direct test of
                  the premise: a visual tower that modelled the scene perfectly
                  would answer identically at every baseline, so the flip rate
                  is the size of the effect the intervention has, independent of
                  whether the answer was right to begin with.
  stability       fraction of samples correct at every baseline / wrong at every
                  baseline / unstable (correct at some, wrong at others).
  failure overlap for every pair of baselines, |F_a & F_b| / |F_a | F_b| over
                  the failure sets, plus the raw counts. High overlap means the
                  baselines fail on the SAME samples (a shared difficulty);
                  low overlap means each baseline breaks a different subset,
                  i.e. the failures are driven by the intervention itself.

Usage
    python analysis/analyze_baseline_overlap.py --result_dir result
    python analysis/analyze_baseline_overlap.py --result_dir result --dump_per_sample

Writes analysis_summary.json and one <model>__<task>_overlap.csv per run next
to the results, and prints a readable table.
"""

import argparse
import csv
import itertools
import json
import os
import pathlib
import sys
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dataset import BASELINES  # noqa: E402


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    default_result = pathlib.Path(__file__).resolve().parent.parent / "result"
    parser.add_argument("--result_dir", default=str(default_result))
    parser.add_argument(
        "--output_dir",
        default=str(pathlib.Path(__file__).resolve().parent / "intervention_analysis"),
        help="Defaults to analysis/intervention_analysis, beside the notebook that reads it.",
    )
    parser.add_argument(
        "--reference_baseline",
        default="0.00",
        help="Baseline that flip rate and per-pair comparisons are measured against.",
    )
    parser.add_argument(
        "--dump_per_sample",
        action="store_true",
        help="Also write a per-sample correctness/answer matrix for each run.",
    )
    return parser.parse_args()


def sample_key(row):
    """A key that identifies one question, unique within a task.

    VSR reuses an image across several statements, so its `sample_id` is the
    image id and repeats (2195 rows over 1874 ids). Keying the cross-baseline
    alignment on `sample_id` alone would silently keep only the last question
    per image and drop the rest, so the question text is part of the key. The
    Youtube levels already have unique ids and are unaffected.
    """
    return (row["sample_id"], row.get("question", ""))


def load_runs(result_dir):
    """{(model, task): {baseline: {sample_key: row}}} from result/<model>/<task>/baseline_<b>/."""
    runs = defaultdict(dict)
    for metrics_path in sorted(pathlib.Path(result_dir).glob("*/*/baseline_*/metrics.json")):
        predictions_path = metrics_path.parent / "predictions.jsonl"
        if not predictions_path.is_file():
            print(f"  skipping {metrics_path.parent} (no predictions.jsonl)", file=sys.stderr)
            continue
        with metrics_path.open(encoding="utf-8") as f:
            metrics = json.load(f)
        with predictions_path.open(encoding="utf-8") as f:
            rows = {}
            for line in f:
                if not line.strip():
                    continue
                row = json.loads(line)
                rows[sample_key(row)] = row
        key = (metrics.get("model_id", metrics_path.parents[2].name), metrics.get("task", metrics_path.parents[1].name))
        runs[key][metrics["baseline"]] = {"metrics": metrics, "rows": rows}
    return runs


def analyse_run(baselines_data, reference_baseline):
    """Accuracy / flip rate / stability / pairwise failure overlap for one run."""
    baselines = [b for b in BASELINES if b in baselines_data]
    common = set.intersection(*(set(baselines_data[b]["rows"]) for b in baselines))
    dropped = {b: len(baselines_data[b]["rows"]) - len(common) for b in baselines}

    correct = {b: {sid for sid in common if baselines_data[b]["rows"][sid]["correct"]} for b in baselines}
    answers = {b: {sid: baselines_data[b]["rows"][sid]["prediction_norm"] for sid in common} for b in baselines}
    failures = {b: common - correct[b] for b in baselines}

    per_baseline = []
    reference = reference_baseline if reference_baseline in baselines else baselines[0]
    for b in baselines:
        flipped = {sid for sid in common if answers[b][sid] != answers[reference][sid]}
        per_baseline.append(
            {
                "baseline": b,
                "num_samples": len(common),
                "accuracy": len(correct[b]) / len(common) if common else 0.0,
                "num_failures": len(failures[b]),
                "flip_rate_vs_reference": len(flipped) / len(common) if common else 0.0,
                "unparsed_rate": baselines_data[b]["metrics"].get("unparsed_rate"),
            }
        )

    always_correct = set.intersection(*(correct[b] for b in baselines)) if baselines else set()
    always_wrong = set.intersection(*(failures[b] for b in baselines)) if baselines else set()
    unstable = common - always_correct - always_wrong

    pairwise = []
    for a, b in itertools.combinations(baselines, 2):
        intersection = failures[a] & failures[b]
        union = failures[a] | failures[b]
        agree = sum(1 for sid in common if answers[a][sid] == answers[b][sid])
        pairwise.append(
            {
                "baseline_a": a,
                "baseline_b": b,
                "failures_a": len(failures[a]),
                "failures_b": len(failures[b]),
                "failures_shared": len(intersection),
                "failure_jaccard": len(intersection) / len(union) if union else 1.0,
                # Of the samples a fails, how many does b also fail.
                "failure_recall_a_in_b": len(intersection) / len(failures[a]) if failures[a] else float("nan"),
                # The same measure over the samples each baseline gets RIGHT. It is
                # not redundant with the failure Jaccard: with 4-way answers the
                # success sets are small and the failure sets large, so the two
                # move independently and a task can share its successes while its
                # failures scatter (or the reverse).
                "successes_a": len(correct[a]),
                "successes_b": len(correct[b]),
                "successes_shared": len(correct[a] & correct[b]),
                "success_jaccard": (
                    len(correct[a] & correct[b]) / len(correct[a] | correct[b])
                    if (correct[a] | correct[b]) else 1.0
                ),
                "answer_agreement": agree / len(common) if common else 0.0,
            }
        )

    return {
        "baselines": baselines,
        "reference_baseline": reference,
        "num_common_samples": len(common),
        "dropped_per_baseline": dropped,
        "per_baseline": per_baseline,
        "accuracy_spread": (
            max(r["accuracy"] for r in per_baseline) - min(r["accuracy"] for r in per_baseline)
            if per_baseline
            else 0.0
        ),
        "stability": {
            "always_correct": len(always_correct),
            "always_wrong": len(always_wrong),
            "unstable": len(unstable),
            "always_correct_rate": len(always_correct) / len(common) if common else 0.0,
            "always_wrong_rate": len(always_wrong) / len(common) if common else 0.0,
            "unstable_rate": len(unstable) / len(common) if common else 0.0,
        },
        "pairwise": pairwise,
        "_common": common,
        "_answers": answers,
        "_correct": correct,
    }


def write_pairwise_csv(path, result):
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "baseline_a", "baseline_b",
                "failures_a", "failures_b", "failures_shared",
                "failure_jaccard", "failure_recall_a_in_b",
                "successes_a", "successes_b", "successes_shared", "success_jaccard",
                "answer_agreement",
            ],
        )
        writer.writeheader()
        writer.writerows(result["pairwise"])


def write_per_sample_csv(path, result):
    baselines = result["baselines"]
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(
            ["sample_id", "question"]
            + [f"answer_{b}" for b in baselines]
            + [f"correct_{b}" for b in baselines]
            + ["num_correct", "stability"]
        )
        for sid in sorted(result["_common"]):
            flags = [sid in result["_correct"][b] for b in baselines]
            total = sum(flags)
            stability = "always_correct" if total == len(baselines) else (
                "always_wrong" if total == 0 else "unstable"
            )
            writer.writerow(
                list(sid)
                + [result["_answers"][b][sid] for b in baselines]
                + [int(flag) for flag in flags]
                + [total, stability]
            )


def main():
    args = parse_args()
    runs = load_runs(args.result_dir)
    if not runs:
        raise SystemExit(
            f"No results under {args.result_dir}; expected "
            "<result_dir>/<model>/<task>/baseline_<b>/{metrics.json,predictions.jsonl}"
        )

    output_dir = pathlib.Path(args.output_dir or os.path.join(args.result_dir, "analysis"))
    output_dir.mkdir(parents=True, exist_ok=True)

    summary = []
    for (model_id, task), baselines_data in sorted(runs.items()):
        result = analyse_run(baselines_data, args.reference_baseline)
        tag = f"{model_id.replace('/', '_')}__{task}"

        write_pairwise_csv(output_dir / f"{tag}_failure_overlap.csv", result)
        if args.dump_per_sample:
            write_per_sample_csv(output_dir / f"{tag}_per_sample.csv", result)

        print(f"\n=== {model_id}  |  {task}  |  n={result['num_common_samples']} ===")
        print(f"{'baseline':>10} {'accuracy':>10} {'failures':>10} {'flip vs ' + result['reference_baseline']:>14}")
        for row in result["per_baseline"]:
            print(
                f"{row['baseline']:>10} {row['accuracy']:>10.4f} {row['num_failures']:>10d} "
                f"{row['flip_rate_vs_reference']:>14.4f}"
            )
        stability = result["stability"]
        print(
            f"  accuracy spread {result['accuracy_spread']:.4f} | "
            f"always correct {stability['always_correct_rate']:.4f} | "
            f"always wrong {stability['always_wrong_rate']:.4f} | "
            f"unstable {stability['unstable_rate']:.4f}"
        )
        print("  failure-set overlap (Jaccard):")
        for pair in result["pairwise"]:
            print(
                f"    {pair['baseline_a']} vs {pair['baseline_b']}: "
                f"J={pair['failure_jaccard']:.4f} shared={pair['failures_shared']:>5d} "
                f"(|F_a|={pair['failures_a']}, |F_b|={pair['failures_b']}) "
                f"answer agreement={pair['answer_agreement']:.4f}"
            )

        summary.append(
            {
                "model_id": model_id,
                "task": task,
                **{k: v for k, v in result.items() if not k.startswith("_")},
            }
        )

    summary_path = output_dir / "analysis_summary.json"
    with summary_path.open("w", encoding="utf-8") as f:
        json.dump({"result_dir": os.path.abspath(args.result_dir), "runs": summary}, f, indent=2)
    print(f"\nSummary saved to {summary_path}")


if __name__ == "__main__":
    main()