tools/classifier_confusion_matrix.py

"""
Generate the confusion matrix for the production classifier checkpoint.

Loads the saved model from `models/classifier/` and evaluates it on
`data_combined/combined_test_v2.json` (114 samples — same set that
produced the 87.72 % test accuracy in outputs/evaluation_report.json).

Outputs:
  - prints a markdown table to stdout
  - saves outputs/classifier_confusion_matrix.png (if matplotlib available)
  - saves outputs/classifier_confusion_matrix.json (raw counts + classes)
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import (
    LayoutLMv3ForSequenceClassification,
    LayoutLMv3Processor,
)

Image.MAX_IMAGE_PIXELS = None

# Anchored to repo root regardless of where this script is run from
REPO_ROOT = Path(__file__).resolve().parents[1]

TEST_JSON       = REPO_ROOT / "data_combined" / "combined_test_v2.json"
MAPPINGS        = REPO_ROOT / "assets" / "label_mappings.json"
CLASSIFIER_DIR  = REPO_ROOT / "models" / "classifier"
OUT_DIR         = REPO_ROOT / "outputs"

MAX_LENGTH     = 512
MAX_IMAGE_SIDE = 2048
MAX_WORDS      = 354
MIN_CONF       = 30


def load_image(image_path: str | None) -> Image.Image:
    if not image_path or not Path(image_path).exists():
        return Image.new("RGB", (1654, 2339), (255, 255, 255))
    img = Image.open(image_path).convert("RGB")
    if max(img.size) > MAX_IMAGE_SIDE:
        img.thumbnail((MAX_IMAGE_SIDE, MAX_IMAGE_SIDE))
    return img


def vertical_boxes(n: int, img_h: int) -> list[list[int]]:
    if n <= 0:
        return []
    h = max(img_h // n, 1)
    return [[0, int(i * h / img_h * 1000), 1000, int((i + 1) * h / img_h * 1000)] for i in range(n)]


def build_words_boxes(rec: dict) -> tuple[list[str], list[list[int]]]:
    img_h = rec.get("image_height", 2339)
    ocr_path = rec.get("ocr_path") or rec.get("ocr_json_path")
    if ocr_path and Path(ocr_path).exists():
        try:
            with open(ocr_path, encoding="utf-8") as f:
                ocr = json.load(f)
        except Exception:
            ocr = {}
        words_raw = ocr.get("words", [])[:MAX_WORDS]
        bnorm_raw = ocr.get("bboxes_norm", [])[:MAX_WORDS]
        confs_raw = ocr.get("confs", [])[:MAX_WORDS]
        words, bnorm = [], []
        for i, (w, bn) in enumerate(zip(words_raw, bnorm_raw)):
            try:
                conf = float(confs_raw[i] if i < len(confs_raw) else 100)
            except Exception:
                conf = 100
            if conf < MIN_CONF:
                continue
            words.append(w)
            bnorm.append(bn)
        if words:
            return words, bnorm

    words = (rec.get("ocr_text", "") or "").split()[:MAX_WORDS] or ["[PAD]"]
    return words, vertical_boxes(len(words), img_h)


def resolve_model_path(p: Path) -> Path:
    if (p / "model.safetensors").exists() or (p / "pytorch_model.bin").exists():
        return p
    ckpts = [c for c in p.glob("checkpoint-*") if c.is_dir()]
    if ckpts:
        return max(ckpts, key=lambda c: int(c.name.split("-")[-1]))
    raise FileNotFoundError(f"No model in {p}")


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    with open(MAPPINGS, encoding="utf-8") as f:
        mappings = json.load(f)
    doc_classes: list[str] = mappings["doc_classes"]
    n_classes = len(doc_classes)

    with open(TEST_JSON, encoding="utf-8") as f:
        test_data = json.load(f)
    print(f"Test set: {TEST_JSON}  ({len(test_data)} records)")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    processor = LayoutLMv3Processor.from_pretrained(
        "microsoft/layoutlmv3-base", apply_ocr=False,
    )
    model_path = resolve_model_path(CLASSIFIER_DIR)
    print(f"Classifier: {model_path}")
    model = LayoutLMv3ForSequenceClassification.from_pretrained(model_path).to(device).eval()

    matrix = np.zeros((n_classes, n_classes), dtype=int)
    correct = 0

    for i, rec in enumerate(test_data, 1):
        words, boxes = build_words_boxes(rec)
        image = load_image(rec.get("image_path"))
        enc = processor(
            image, words, boxes=boxes,
            max_length=MAX_LENGTH, padding="max_length",
            truncation=True, return_tensors="pt",
        ).to(device)
        with torch.no_grad():
            logits = model(**enc).logits
        pred_id = int(logits.argmax(dim=-1).item())
        true_id = int(rec["doc_class_id"])
        matrix[true_id, pred_id] += 1
        if pred_id == true_id:
            correct += 1
        if i % 20 == 0 or i == len(test_data):
            acc = correct / i
            print(f"  [{i:3d}/{len(test_data)}]  running acc = {acc:.4f}")

    acc = correct / len(test_data)
    print(f"\nFinal accuracy: {acc:.4f}  ({correct}/{len(test_data)})\n")

    # ── Print as markdown table ───────────────────────────────────────────
    name_w = max(len(c) for c in doc_classes)
    header = "true \\ pred".ljust(name_w + 2) + "".join(c.rjust(name_w + 2) for c in doc_classes) + "   total"
    print(header)
    print("-" * len(header))
    totals = matrix.sum(axis=1)
    for i, c in enumerate(doc_classes):
        row = c.ljust(name_w + 2) + "".join(str(int(matrix[i, j])).rjust(name_w + 2) for j in range(n_classes))
        row += f"   {totals[i]:5d}"
        print(row)
    # Per-class precision/recall/F1
    print()
    print("class".ljust(name_w + 2) + "support".rjust(8) + "precision".rjust(11) + "recall".rjust(9) + "f1".rjust(7))
    print("-" * (name_w + 2 + 8 + 11 + 9 + 7))
    for i, c in enumerate(doc_classes):
        tp = matrix[i, i]
        fp = matrix[:, i].sum() - tp
        fn = matrix[i, :].sum() - tp
        precision = tp / (tp + fp) if (tp + fp) else 0.0
        recall    = tp / (tp + fn) if (tp + fn) else 0.0
        f1        = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
        support   = int(matrix[i, :].sum())
        print(c.ljust(name_w + 2) + f"{support:8d}{precision:11.3f}{recall:9.3f}{f1:7.3f}")

    # ── Persist ──────────────────────────────────────────────────────────
    out_json = OUT_DIR / "classifier_confusion_matrix.json"
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump({
            "doc_classes": doc_classes,
            "matrix": matrix.tolist(),
            "test_samples": len(test_data),
            "accuracy": acc,
            "model_checkpoint": str(model_path),
            "test_file": str(TEST_JSON),
        }, f, indent=2)
    print(f"\nSaved: {out_json}")

    try:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(7, 6))
        im = ax.imshow(matrix, cmap="Oranges")
        ax.set_xticks(range(n_classes)); ax.set_yticks(range(n_classes))
        ax.set_xticklabels(doc_classes, rotation=35, ha="right")
        ax.set_yticklabels(doc_classes)
        ax.set_xlabel("Predicted"); ax.set_ylabel("True")
        ax.set_title(f"Classifier confusion matrix\nacc = {acc:.4f}  ({correct}/{len(test_data)})")
        # Cell counts
        thresh = matrix.max() / 2 if matrix.max() else 1
        for i in range(n_classes):
            for j in range(n_classes):
                v = int(matrix[i, j])
                if v == 0:
                    continue
                ax.text(j, i, str(v), ha="center", va="center",
                        color="white" if v > thresh else "black", fontsize=10)
        fig.colorbar(im, ax=ax, shrink=0.7)
        fig.tight_layout()
        out_png = OUT_DIR / "classifier_confusion_matrix.png"
        fig.savefig(out_png, dpi=150)
        print(f"Saved: {out_png}")
    except ImportError:
        print("matplotlib not installed — skipping PNG export.")


if __name__ == "__main__":
    main()