#!/usr/bin/env python3
"""Build 2-row collages for the top confusion pairs from conformal_v13.json.

Usage: python confusion_pair_collage.py
Outputs: runs/h100/confusion_pairs/pair_{A}_vs_{B}.png and README.md
"""
from __future__ import annotations

import json
import random
import re
from collections import defaultdict
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

ROOT = Path("/arf/scratch/stakan/hitit-proje")
CONFORMAL = ROOT / "hitit_ocr/runs/h100/conformal_v13.json"
MANIFEST = ROOT / "datasets/sources/hitit_local/manifest_v13_ultimate.jsonl"
OUT_DIR = ROOT / "hitit_ocr/runs/h100/confusion_pairs"

TOP_K = 15
N_PER_ROW = 8
MIN_SAMPLES = 4
TILE = 128
HEADER_H = 40
ROW_LABEL_W = 80
SEED = 42


def safe_name(s: str) -> str:
    return re.sub(r"[^A-Za-z0-9]+", "_", s).strip("_") or "X"


def load_pairs() -> list[tuple[str, str, int]]:
    data = json.loads(CONFORMAL.read_text())
    pairs = sorted(data.get("top_confusion_pairs", []), key=lambda p: -int(p["count"]))
    out, seen = [], set()
    for p in pairs:
        a, b = p["true"], p["pred"]
        key = tuple(sorted([a, b]))
        if key in seen:
            continue
        seen.add(key)
        out.append((a, b, int(p["count"])))
        if len(out) >= TOP_K:
            break
    return out


def index_manifest(needed: set[str]) -> dict[str, list[str]]:
    by_label: dict[str, list[str]] = defaultdict(list)
    with MANIFEST.open("r") as fh:
        for line in fh:
            if not line.strip():
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                continue
            lbl = rec.get("unified_label") or rec.get("label_norm")
            path = rec.get("path")
            if lbl in needed and path:
                by_label[lbl].append(path)
    return by_label


def load_tile(p: str) -> Image.Image | None:
    try:
        im = Image.open(p).convert("RGB")
    except Exception:
        return None
    im.thumbnail((TILE, TILE), Image.LANCZOS)
    canvas = Image.new("RGB", (TILE, TILE), (255, 255, 255))
    canvas.paste(im, ((TILE - im.width) // 2, (TILE - im.height) // 2))
    return canvas


def get_font(size: int):
    for c in ("/usr/share/fonts/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"):
        try:
            return ImageFont.truetype(c, size)
        except Exception:
            pass
    return ImageFont.load_default()


def make_row(paths: list[str]) -> list[Image.Image]:
    tiles = []
    for p in paths:
        t = load_tile(p)
        if t is not None:
            tiles.append(t)
        if len(tiles) >= N_PER_ROW:
            break
    return tiles


def build_collage(a: str, b: str, count: int, ap: list[str], bp: list[str]) -> Image.Image:
    w, h = ROW_LABEL_W + N_PER_ROW * TILE, HEADER_H + 2 * TILE
    img = Image.new("RGB", (w, h), (245, 245, 245))
    draw = ImageDraw.Draw(img)
    tf, lf = get_font(22), get_font(18)
    draw.rectangle([0, 0, w, HEADER_H], fill=(30, 30, 30))
    draw.text((10, 8), f"Confusion pair: {a}  vs  {b}   (count={count})", fill=(255, 255, 255), font=tf)
    for ri, (lbl, tiles) in enumerate(((a, make_row(ap)), (b, make_row(bp)))):
        y0 = HEADER_H + ri * TILE
        draw.rectangle([0, y0, ROW_LABEL_W, y0 + TILE], fill=(60, 60, 60))
        draw.text((8, y0 + TILE // 2 - 10), lbl, fill=(255, 255, 255), font=lf)
        for i, t in enumerate(tiles):
            img.paste(t, (ROW_LABEL_W + i * TILE, y0))
    return img


def main() -> None:
    random.seed(SEED)
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    pairs = load_pairs()
    needed = {x for a, b, _ in pairs for x in (a, b)}
    print(f"[info] indexing manifest for {len(needed)} labels ...")
    by_label = index_manifest(needed)
    for lbl, lst in by_label.items():
        random.shuffle(lst)
    produced, skipped = [], []
    for a, b, count in pairs:
        ap = by_label.get(a, [])[:N_PER_ROW]
        bp = by_label.get(b, [])[:N_PER_ROW]
        if len(ap) < MIN_SAMPLES or len(bp) < MIN_SAMPLES:
            skipped.append((a, b, count, len(ap), len(bp)))
            continue
        img = build_collage(a, b, count, ap, bp)
        fname = f"pair_{safe_name(a)}_vs_{safe_name(b)}.png"
        img.save(OUT_DIR / fname)
        produced.append((a, b, count, fname))
        print(f"[ok] {a} vs {b} (n={count}) -> {fname}")
    lines = ["# Confusion Pair Collages", "", f"Source: `{CONFORMAL}`",
             f"Top-{TOP_K} unique pairs (by count).", "",
             "| true | pred | count | collage |", "|---|---|---|---|"]
    lines += [f"| {a} | {b} | {c} | `{fn}` |" for a, b, c, fn in produced]
    if skipped:
        lines += ["", "## Skipped (insufficient samples)",
                  "| true | pred | count | n_true | n_pred |", "|---|---|---|---|---|"]
        lines += [f"| {a} | {b} | {c} | {na} | {nb} |" for a, b, c, na, nb in skipped]
    (OUT_DIR / "README.md").write_text("\n".join(lines) + "\n")
    print(f"[done] {len(produced)} collages, {len(skipped)} skipped -> {OUT_DIR}")


if __name__ == "__main__":
    main()