livecodebench_difficulty_analysis.py

"""
livecodebench_difficulty_analysis.py
─────────────────────────────────────────────────────────────────────────────
Offline difficulty-stratified analysis for LiveCodeBench trials.

No API calls. Steps:
  1. Load code_generation_lite from HuggingFace → build {question_id: difficulty}
  2. Read specified trial details.json files → load {sample_id: correct}
  3. Join on ID, group by difficulty, compute per-condition accuracy
  4. Print stratified table + McNemar-style discordant pair counts for key pairs

Usage:
    # Auto-discover all livecodebench trials in trials/ and analyse them
    python livecodebench_difficulty_analysis.py

    # Or specify trial dirs explicitly (bare name or full path)
    python livecodebench_difficulty_analysis.py \\
        --trials trials/livecodebench_C1_... trials/livecodebench_C2_... ...

    # Restrict to specific conditions (useful if you have partial results)
    python livecodebench_difficulty_analysis.py --conditions C1 C2 C3 C4
"""

import argparse
import json
import os
import re
from collections import defaultdict

TRIALS_DIR = "trials"
DIFFICULTIES = ["easy", "medium", "hard"]
CONDITION_ORDER = ["C1", "C2", "C3", "C4", "C1r", "C2r", "C3r", "C4r"]


# ── 1. Load difficulty map from HuggingFace ───────────────────────────────────

def load_difficulty_map() -> dict[str, str]:
    """
    Returns {question_id (str): difficulty (str)} for all rows in
    livecodebench/code_generation_lite.
    Difficulty values are typically 'easy', 'medium', 'hard'.
    """
    print("Loading LiveCodeBench difficulty labels from HuggingFace...")
    from datasets import load_dataset
    ds = load_dataset(
        "livecodebench/code_generation_lite",
        split="test",
        trust_remote_code=True,
    )
    diff_map = {}
    for row in ds:
        qid = str(row["question_id"])
        diff = str(row.get("difficulty", "unknown")).lower().strip()
        diff_map[qid] = diff
    print(f"  Loaded {len(diff_map)} difficulty labels "
          f"({sum(1 for v in diff_map.values() if v == 'easy')} easy / "
          f"{sum(1 for v in diff_map.values() if v == 'medium')} medium / "
          f"{sum(1 for v in diff_map.values() if v == 'hard')} hard).")
    return diff_map


# ── 2. Discover / load trials ─────────────────────────────────────────────────

def _resolve_trial(path: str) -> str:
    if os.path.isdir(path):
        return path
    candidate = os.path.join(TRIALS_DIR, path)
    if os.path.isdir(candidate):
        return candidate
    raise FileNotFoundError(f"Trial not found: {path!r}")


def discover_livecodebench_trials(conditions: list[str] | None = None) -> list[str]:
    """Find all livecodebench trial dirs in TRIALS_DIR, optionally filtered."""
    if not os.path.isdir(TRIALS_DIR):
        return []
    pattern = re.compile(r"^livecodebench_([A-Za-z0-9r]+)_\d{8}T\d{6}Z$")
    dirs = []
    for name in sorted(os.listdir(TRIALS_DIR)):
        m = pattern.match(name)
        if not m:
            continue
        cond = m.group(1)
        if conditions and cond not in conditions:
            continue
        dirs.append(os.path.join(TRIALS_DIR, name))
    return dirs


def load_trial(trial_dir: str) -> tuple[str, dict[str, bool]]:
    """
    Returns (condition, {sample_id: correct}).
    Uses summary.json for condition name, details.json for per-sample results.
    """
    resolved = _resolve_trial(trial_dir)
    with open(os.path.join(resolved, "summary.json"), encoding="utf-8") as f:
        summary = json.load(f)
    with open(os.path.join(resolved, "details.json"), encoding="utf-8") as f:
        details = json.load(f)
    results = {str(r["id"]): bool(r.get("correct", False)) for r in details}
    return summary["condition"], results


# ── 3. Stratified accuracy ────────────────────────────────────────────────────

def stratified_accuracy(
        results: dict[str, bool],
        diff_map: dict[str, str],
) -> dict[str, dict]:
    """
    Returns {difficulty: {"n": int, "correct": int, "accuracy": float}}.
    Samples whose ID is not in diff_map are placed in "unknown".
    """
    buckets = defaultdict(lambda: {"n": 0, "correct": 0})
    for sid, correct in results.items():
        diff = diff_map.get(sid, "unknown")
        buckets[diff]["n"] += 1
        if correct:
            buckets[diff]["correct"] += 1
    out = {}
    for diff, counts in buckets.items():
        n = counts["n"]
        c = counts["correct"]
        out[diff] = {"n": n, "correct": c, "accuracy": c / n if n else 0.0}
    return out


# ── 4. McNemar discordant pairs (per difficulty) ──────────────────────────────

def discordant_pairs(
        results_a: dict[str, bool],
        results_b: dict[str, bool],
        diff_map: dict[str, str],
        difficulty: str,
) -> tuple[int, int]:
    """
    For shared IDs in the given difficulty bucket:
    Returns (n10, n01) where
        n10 = A correct, B wrong
        n01 = A wrong,   B correct
    """
    shared = set(results_a) & set(results_b)
    n10 = n01 = 0
    for sid in shared:
        if diff_map.get(sid, "unknown") != difficulty:
            continue
        a_ok = results_a[sid]
        b_ok = results_b[sid]
        if a_ok and not b_ok:
            n10 += 1
        if not a_ok and b_ok:
            n01 += 1
    return n10, n01


# ── 5. Pretty printing ────────────────────────────────────────────────────────

def _pct(acc: float) -> str:
    return f"{acc * 100:.1f}%"


def print_stratified_table(
        condition_data: dict[str, dict[str, dict]],
        difficulties: list[str],
):
    """
    condition_data: {condition: {difficulty: {n, correct, accuracy}}}
    """
    conds = [c for c in CONDITION_ORDER if c in condition_data]

    # Header
    col_w = 12
    header = f"{'Condition':<8}" + "".join(
        f"{'Overall':>{col_w}}"
        + "".join(f"{d.capitalize():>{col_w}}" for d in difficulties)
    )
    print(f"\n{'═' * (8 + col_w * (1 + len(difficulties)))}")
    print("  LiveCodeBench — Accuracy by Difficulty")
    print(f"{'─' * (8 + col_w * (1 + len(difficulties)))}")
    print(f"{'':8}" + f"{'Overall':>{col_w}}" +
          "".join(f"{d.capitalize():>{col_w}}" for d in difficulties))
    print(f"{'─' * (8 + col_w * (1 + len(difficulties)))}")

    for cond in conds:
        data = condition_data[cond]
        # Overall = sum across known difficulties
        total_n = sum(data[d]["n"] for d in difficulties if d in data)
        total_c = sum(data[d]["correct"] for d in difficulties if d in data)
        overall = total_c / total_n if total_n else 0.0

        row = f"{cond:<8}" + f"{_pct(overall):>{col_w}}"
        for diff in difficulties:
            if diff in data:
                acc = data[diff]["accuracy"]
                n = data[diff]["n"]
                row += f"{_pct(acc) + f' ({n})':>{col_w}}"
            else:
                row += f"{'—':>{col_w}}"
        print(row)

    print(f"{'═' * (8 + col_w * (1 + len(difficulties)))}\n")


def print_discordant_table(
        condition_data: dict[str, dict[str, dict]],
        all_results: dict[str, dict[str, bool]],
        diff_map: dict[str, str],
        pairs: list[tuple[str, str]],
        difficulties: list[str],
):
    """Print discordant pair counts for key condition pairs, by difficulty."""
    conds_present = set(condition_data.keys())
    valid_pairs = [(a, b) for a, b in pairs
                   if a in conds_present and b in conds_present]
    if not valid_pairs:
        return

    print(f"{'─' * 62}")
    print("  Discordant pairs (n10 / n01) by difficulty")
    print(f"  (n10: A✓B✗  n01: A✗B✓  — basis for McNemar's test)")
    print(f"{'─' * 62}")
    col_w = 16
    print(f"{'Pair':<14}" + "".join(f"{d.capitalize():>{col_w}}" for d in difficulties))
    print(f"{'─' * 62}")
    for cond_a, cond_b in valid_pairs:
        label = f"{cond_a} vs {cond_b}"
        row = f"{label:<14}"
        for diff in difficulties:
            n10, n01 = discordant_pairs(
                all_results[cond_a], all_results[cond_b], diff_map, diff
            )
            row += f"{f'{n10}/{n01}':>{col_w}}"
        print(row)
    print(f"{'─' * 62}\n")


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    global TRIALS_DIR
    TRIALS_DIR = "./trials"

    ap = argparse.ArgumentParser(
        description="Offline difficulty-stratified analysis for LiveCodeBench trials."
    )
    ap.add_argument(
        "--trials", nargs="*", default=None,
        metavar="TRIAL_DIR",
        help="Trial dirs to analyse (default: auto-discover all in trials/).",
    )
    ap.add_argument(
        "--conditions", nargs="*", default=None,
        metavar="COND",
        help="Filter to specific conditions, e.g. --conditions C1 C2 C3 C4.",
    )
    ap.add_argument(
        "--trials_dir", default=TRIALS_DIR,
        help=f"Root trials directory (default: {TRIALS_DIR!r}).",
    )
    args = ap.parse_args()

    TRIALS_DIR = args.trials_dir

    # ── Collect trial dirs ────────────────────────────────────────────────────
    if args.trials:
        trial_dirs = args.trials
    else:
        trial_dirs = discover_livecodebench_trials(args.conditions)
        if not trial_dirs:
            print(f"No livecodebench trials found in {TRIALS_DIR!r}. "
                  f"Use --trials to specify paths explicitly.")
            return
        print(f"Auto-discovered {len(trial_dirs)} trial(s):")
        for d in trial_dirs:
            print(f"  {d}")

    # ── Load trials ───────────────────────────────────────────────────────────
    all_results: dict[str, dict[str, bool]] = {}  # {condition: {id: correct}}
    for td in trial_dirs:
        try:
            cond, results = load_trial(td)
            if cond in all_results:
                print(f"  WARNING: duplicate condition {cond!r}, "
                      f"keeping {td} (last seen wins).")
            all_results[cond] = results
            print(f"  Loaded {cond}: {sum(results.values())}/{len(results)} correct")
        except Exception as e:
            print(f"  ERROR loading {td}: {e}")

    if not all_results:
        print("No valid trials loaded. Exiting.")
        return

    # ── Load difficulty map ───────────────────────────────────────────────────
    diff_map = load_difficulty_map()

    # ── Compute stratified accuracy per condition ─────────────────────────────
    condition_data: dict[str, dict[str, dict]] = {}
    for cond, results in all_results.items():
        condition_data[cond] = stratified_accuracy(results, diff_map)

    # Determine which difficulties are actually present
    present_diffs = [d for d in DIFFICULTIES
                     if any(d in condition_data[c] for c in condition_data)]

    # ── Print results ─────────────────────────────────────────────────────────
    print_stratified_table(condition_data, present_diffs)

    # Key pairs: primary C2/C3/C4 vs C3, supplementary C2r vs C3r
    key_pairs = [
        ("C2", "C3"), ("C2", "C4"), ("C4", "C3"),
        ("C2r", "C3r"), ("C2r", "C4r"), ("C4r", "C3r"),
    ]
    print_discordant_table(
        condition_data, all_results, diff_map, key_pairs, present_diffs
    )

    # ── Per-difficulty summary (delta vs C1 or C1r) ───────────────────────────
    for baseline, label in [("C1", "Primary"), ("C1r", "Supplementary")]:
        if baseline not in condition_data:
            continue
        print(f"{'─' * 62}")
        print(f"  {label} setting — delta vs {baseline} by difficulty")
        print(f"{'─' * 62}")
        col_w = 16
        print(f"{'Condition':<10}" +
              "".join(f"{d.capitalize():>{col_w}}" for d in present_diffs))
        print(f"{'─' * 62}")
        base = condition_data[baseline]
        related = [c for c in CONDITION_ORDER
                   if c in condition_data and c != baseline
                   and (c.endswith("r") == baseline.endswith("r"))]
        for cond in related:
            data = condition_data[cond]
            row = f"{cond:<10}"
            for diff in present_diffs:
                if diff in data and diff in base:
                    delta = data[diff]["accuracy"] - base[diff]["accuracy"]
                    sign = "+" if delta >= 0 else ""
                    row += f"{sign + f'{delta * 100:.1f}pp':>{col_w}}"
                else:
                    row += f"{'—':>{col_w}}"
            print(row)
        print(f"{'─' * 62}\n")


if __name__ == "__main__":
    main()