"""Per-tablet error analysis for the v13 ensemble on Hitit cuneiform OCR.

Loads the ensemble probs dump (which already carries tablet_ids, targets, and
label_to_idx), computes per-tablet top1 + top-5 confusions, and dumps a JSON
report plus a human-readable markdown table with the worst/best tablets.

Usage:
    python3 hitit_ocr/src/analysis/per_tablet_errors.py
"""

from __future__ import annotations

import json
from collections import Counter, defaultdict
from pathlib import Path

import torch

PROBS_PATH = Path("/arf/scratch/stakan/hitit-proje/hitit_ocr/runs/h100/ensemble_v13_probs.pt")
EVAL_PATH = Path("/arf/scratch/stakan/hitit-proje/hitit_ocr/runs/h100/ensemble_v13_eval.json")
LABEL_MAP_PATH = Path("/arf/scratch/stakan/hitit-proje/hitit_ocr/runs/h100/v13_label_to_idx.json")
OUT_JSON = Path("/arf/scratch/stakan/hitit-proje/hitit_ocr/runs/h100/per_tablet_error_analysis.json")
OUT_MD = Path("/arf/scratch/stakan/hitit-proje/hitit_ocr/runs/h100/per_tablet_error_analysis.md")

MIN_SAMPLES = 20  # tablets below this threshold are excluded from worst/best shortlists


def load_inputs():
    dump = torch.load(PROBS_PATH, map_location="cpu", weights_only=False)
    probs: torch.Tensor = dump["probs"]
    targets: torch.Tensor = dump["targets"]
    tablet_ids = list(dump["tablet_ids"])
    # Prefer the dump's label_to_idx (it was used to produce the predictions);
    # fall back to the on-disk file so the two agree.
    label_to_idx = dump.get("label_to_idx")
    if label_to_idx is None:
        with LABEL_MAP_PATH.open() as f:
            label_to_idx = json.load(f)
    if len(tablet_ids) != probs.shape[0]:
        raise RuntimeError(
            f"tablet_ids length {len(tablet_ids)} != probs rows {probs.shape[0]}"
        )
    if targets.shape[0] != probs.shape[0]:
        raise RuntimeError(
            f"targets length {targets.shape[0]} != probs rows {probs.shape[0]}"
        )
    return probs, targets, tablet_ids, label_to_idx


def main() -> None:
    probs, targets, tablet_ids, label_to_idx = load_inputs()

    idx_to_label = {int(v): k for k, v in label_to_idx.items()}
    preds = probs.argmax(dim=1)
    correct = preds.eq(targets)
    global_top1 = float(correct.float().mean().item())

    # Cross-check against the existing eval json.
    try:
        with EVAL_PATH.open() as f:
            eval_json = json.load(f)
        expected_top1 = float(eval_json.get("ensemble_top1", -1))
    except Exception:
        expected_top1 = -1.0

    # Group sample indices by tablet.
    by_tablet: dict[str, list[int]] = defaultdict(list)
    for i, tid in enumerate(tablet_ids):
        by_tablet[str(tid)].append(i)

    per_tablet_records: list[dict] = []
    for tid, idxs in by_tablet.items():
        n = len(idxs)
        tcorrect = int(correct[idxs].sum().item())
        acc = tcorrect / n
        # Top-5 confusion pairs (only wrong predictions).
        conf_counter: Counter = Counter()
        for i in idxs:
            if not bool(correct[i].item()):
                true_lbl = idx_to_label[int(targets[i].item())]
                pred_lbl = idx_to_label[int(preds[i].item())]
                conf_counter[(true_lbl, pred_lbl)] += 1
        top5_conf = [
            [t, p, int(c)] for (t, p), c in conf_counter.most_common(5)
        ]
        per_tablet_records.append(
            {
                "tablet_id": tid,
                "n": n,
                "top1": acc,
                "n_correct": tcorrect,
                "top5_confusions": top5_conf,
            }
        )

    # Sort: worst first (lowest accuracy, break ties by higher n so the shortlist
    # surfaces statistically meaningful tablets).
    eligible = [r for r in per_tablet_records if r["n"] >= MIN_SAMPLES]
    worst_sorted = sorted(eligible, key=lambda r: (r["top1"], -r["n"]))
    best_sorted = sorted(eligible, key=lambda r: (-r["top1"], -r["n"]))
    worst_10 = worst_sorted[:10]
    best_5 = best_sorted[:5]

    # Also keep a default stable sort for the full list (by n desc, then acc asc).
    per_tablet_records.sort(key=lambda r: (-r["n"], r["top1"]))

    report = {
        "global_top1": global_top1,
        "expected_top1": expected_top1,
        "n_samples": int(probs.shape[0]),
        "n_tablets": len(by_tablet),
        "min_samples_for_shortlist": MIN_SAMPLES,
        "n_tablets_eligible": len(eligible),
        "per_tablet": per_tablet_records,
        "worst_10": worst_10,
        "best_5": best_5,
    }

    OUT_JSON.parent.mkdir(parents=True, exist_ok=True)
    with OUT_JSON.open("w") as f:
        json.dump(report, f, indent=2, ensure_ascii=False)

    # Console summary.
    print(
        f"Global top1 = {global_top1:.4f} (expected {expected_top1:.4f}) over "
        f"{probs.shape[0]} samples across {len(by_tablet)} tablets "
        f"({len(eligible)} eligible at n>={MIN_SAMPLES})."
    )
    print("\nWorst 10 tablets (n>=20):")
    print(f"  {'tablet_id':<12}{'n':>6}{'acc':>8}   dominant confusion (true->pred, count)")
    for r in worst_10:
        dom = r["top5_confusions"][0] if r["top5_confusions"] else ["-", "-", 0]
        print(
            f"  {r['tablet_id']:<12}{r['n']:>6}{r['top1']:>8.3f}   "
            f"{dom[0]}->{dom[1]} ({dom[2]})"
        )
    print("\nBest 5 tablets (n>=20):")
    print(f"  {'tablet_id':<12}{'n':>6}{'acc':>8}")
    for r in best_5:
        print(f"  {r['tablet_id']:<12}{r['n']:>6}{r['top1']:>8.3f}")

    # Markdown.
    md_lines: list[str] = []
    md_lines.append("# Per-tablet error analysis (v13 ensemble)\n")
    md_lines.append(
        f"- Global top1: **{global_top1:.4f}** (sanity vs `ensemble_v13_eval.json`: {expected_top1:.4f})"
    )
    md_lines.append(
        f"- Samples: {probs.shape[0]} | Tablets: {len(by_tablet)} "
        f"| Eligible (n>={MIN_SAMPLES}): {len(eligible)}"
    )
    md_lines.append(
        f"- Architectures in ensemble: dinov3_vitl14 x2, convnextv2_large, dinov3_vitb14\n"
    )
    md_lines.append("## Worst 10 tablets\n")
    md_lines.append("| tablet_id | n | top1 | dominant confusion (true -> pred x count) |")
    md_lines.append("|---|---:|---:|---|")
    for r in worst_10:
        if r["top5_confusions"]:
            t, p, c = r["top5_confusions"][0]
            dom = f"{t} -> {p} x{c}"
        else:
            dom = "(none)"
        md_lines.append(
            f"| {r['tablet_id']} | {r['n']} | {r['top1']:.3f} | {dom} |"
        )
    md_lines.append("\n## Best 5 tablets\n")
    md_lines.append("| tablet_id | n | top1 |")
    md_lines.append("|---|---:|---:|")
    for r in best_5:
        md_lines.append(f"| {r['tablet_id']} | {r['n']} | {r['top1']:.3f} |")

    # 3-sentence commentary.
    if worst_10 and best_5:
        w = worst_10[0]
        b = best_5[0]
        w_conf = w["top5_confusions"][0] if w["top5_confusions"] else None
        conf_str = (
            f"with the dominant confusion being {w_conf[0]} -> {w_conf[1]} "
            f"({w_conf[2]}/{w['n'] - w['n_correct']} of its errors)"
            if w_conf
            else "without a single dominant confusion"
        )
        md_lines.append("\n## Commentary\n")
        md_lines.append(
            f"The worst tablet (`{w['tablet_id']}`, n={w['n']}) lands at top1={w['top1']:.3f}, "
            f"{conf_str}, hinting at either a distinct visual regime (erosion, lighting, "
            f"scribe hand) or noisy labels rather than bulk class imbalance. "
            f"In contrast the best tablet (`{b['tablet_id']}`, n={b['n']}) reaches top1={b['top1']:.3f}, "
            f"so ensemble capacity is clearly sufficient when the input distribution matches training. "
            f"Error mass concentrates on a handful of tablets: the bottom-10 eligible tablets contribute "
            f"{sum(r['n'] - r['n_correct'] for r in worst_10)} of "
            f"{int((~correct).sum().item())} total errors, which argues for tablet-aware augmentation or "
            f"pseudo-labeling (CoTTA / relight / stroke-aux) targeted at those IDs."
        )

    with OUT_MD.open("w") as f:
        f.write("\n".join(md_lines) + "\n")

    print(f"\nWrote JSON: {OUT_JSON}")
    print(f"Wrote MD:   {OUT_MD}")


if __name__ == "__main__":
    main()