#!/usr/bin/env python3
"""AI-classification accuracy harness for pdf_table_weave.list_tables().

Purely mechanical — no LLM call anywhere in this file. It has two jobs:
1. `--candidates`: sweep PDFs, call pdf_table_weave.list_tables(pdf) on each,
   and emit a flat manifest for a human (or an in-session agent, run
   separately from this script) to label with ground truth.
2. `--score`: compare a judgments file (the agent's per-table classification)
   against a labels file (human ground truth) and report accuracy metrics.

Usage:
    python judgment_eval.py --candidates PATH [PATH...] --out manifest.json \
        [--render-dir DIR]
    python judgment_eval.py --score --labels labels.jsonl \
        --judgments judgments.json [--json] [--seed 12345]
Exit: 0 on success, 2 on error (bad args, invalid label/judgment class,
empty labels file).
"""

from __future__ import annotations

import argparse
import json
import math
import random
import sys
from pathlib import Path

import pikepdf

SCRIPTS_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPTS_DIR))
REPO_ROOT = SCRIPTS_DIR.parent

import pdf_table_weave  # noqa: E402

TRUTH_CLASSES = {"data_table", "layout_table", "figure", "other", "layout_wrapper"}
JUDGMENT_CLASSES = TRUTH_CLASSES | {"unsure"}

# The action spec the engine consumes for a data_table: which rows/cols
# become /TH, and which rows get the alternating-band pattern. Meaningful
# only when truth == "data_table" — see score()'s docstring.
SPEC_KEYS = ("header_rows", "header_cols", "band_rows")


# ── candidates ───────────────────────────────────────────────────────────

def _repo_relative(p: Path) -> str:
    try:
        return str(p.relative_to(REPO_ROOT))
    except ValueError:
        return str(p)


def _resolve_pdfs(paths) -> list:
    found = set()
    for raw in paths:
        p = Path(raw)
        if p.is_dir():
            found.update(f.resolve() for f in p.rglob("*.pdf"))
        elif p.is_file():
            found.add(p.resolve())
    return sorted(found)


def build_candidates(paths, render_dir=None) -> tuple:
    """Returns (candidate_records, pdf_files_swept)."""
    pdf_files = _resolve_pdfs(paths)
    if render_dir:
        Path(render_dir).mkdir(parents=True, exist_ok=True)
    candidates = []
    for pdf_path in pdf_files:
        file_str = _repo_relative(pdf_path)
        try:
            with pikepdf.open(pdf_path) as pdf:
                tables = pdf_table_weave.list_tables(pdf, path=str(pdf_path))
        except Exception as exc:  # noqa: BLE001 - one bad PDF must not abort the sweep
            candidates.append({"file": file_str, "error": str(exc)})
            continue

        doc = None
        if render_dir:
            try:
                import fitz  # PyMuPDF
                doc = fitz.open(pdf_path)
            except Exception:  # noqa: BLE001 - render is best-effort
                doc = None

        for t in tables:
            rec = {"file": file_str, **t}
            pages = t.get("pages") or []
            if render_dir and doc is not None and pages:
                renders = []
                for page in pages:
                    try:
                        png_path = Path(render_dir) / f"{pdf_path.stem}_t{t['index']}_p{page}.png"
                        pix = doc[page - 1].get_pixmap(dpi=110)
                        pix.save(str(png_path))
                        renders.append(str(png_path))
                    except Exception:  # noqa: BLE001 - render is best-effort
                        pass
                if renders:
                    # "render" stays the first page's PNG for backward
                    # compatibility with existing consumers; "renders" is
                    # the full ordered list (additive, one PNG per page in
                    # `pages`). Both keys are absent when nothing rendered.
                    rec["render"] = renders[0]
                    rec["renders"] = renders
            candidates.append(rec)

        if doc is not None:
            doc.close()
    return candidates, pdf_files


# ── scoring ──────────────────────────────────────────────────────────────

def _validate_spec_list(value, context, key):
    """A label/judgment header key must be a list of plain ints (bool is a
    subclass of int in Python but is never a valid row/col index, so it's
    rejected too)."""
    if not isinstance(value, list):
        raise ValueError(f"{key} must be a list for {context}, got {type(value).__name__}")
    for v in value:
        if not isinstance(v, int) or isinstance(v, bool):
            raise ValueError(f"{key} must contain only integers for {context}, got {v!r}")


def load_labels(path) -> list:
    labels = []
    seen = set()
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rec = json.loads(line)
            truth = rec.get("truth")
            if truth not in TRUTH_CLASSES:
                raise ValueError(f"invalid truth value {truth!r} in {path}")
            key = (rec["file"], int(rec["index"]))
            if key in seen:
                raise ValueError(
                    f"duplicate label for {key[0]}:{key[1]} in {path} — "
                    "remove or reconcile the repeated line, ground truth must be unambiguous"
                )
            seen.add(key)
            label = {"file": key[0], "index": key[1], "truth": truth}
            context = f"{key[0]}:{key[1]} in {path}"
            for spec_key in SPEC_KEYS:
                if spec_key in rec:
                    _validate_spec_list(rec[spec_key], context, spec_key)
                    label[spec_key] = rec[spec_key]
            labels.append(label)
    if not labels:
        raise ValueError(f"empty labels file: {path} — no labeled items, nothing to score")
    return labels


def load_judgments(path) -> dict:
    """A judgment value is either the legacy bare class string, or an
    object `{"class": ..., "header_rows": [...], "header_cols": [...],
    "band_rows": [...]}` carrying the action spec for a data_table. The
    header keys are optional even inside the object form; only "class" is
    required."""
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    for file, idx_map in data.items():
        for idx, value in idx_map.items():
            try:
                canonical = int(idx)
            except ValueError:
                raise ValueError(f"non-integer index key {idx!r} for {file} in {path}") from None
            # score() looks up str(index) against these keys and
            # unlabeled-detection normalizes via int() — a non-canonical
            # key (" 0", "+0", "00") passes int() but round-trips to a
            # different string, so it matches neither lookup and silently
            # becomes an invisible abstention (see PR review, round 2 FIX3).
            if str(canonical) != idx:
                raise ValueError(
                    f"non-canonical index key {idx!r} for {file} in {path} "
                    f"(expected {canonical!r})"
                )
            context = f"{file}:{idx} in {path}"
            if isinstance(value, str):
                if value not in JUDGMENT_CLASSES:
                    raise ValueError(f"invalid judgment class {value!r} for {file}:{idx}")
            elif isinstance(value, dict):
                if "class" not in value:
                    raise ValueError(f"judgment object missing 'class' for {context}")
                extra = set(value) - {"class", *SPEC_KEYS}
                if extra:
                    raise ValueError(
                        f"unknown key(s) {sorted(extra)!r} in judgment object for {context}"
                    )
                if value["class"] not in JUDGMENT_CLASSES:
                    raise ValueError(f"invalid judgment class {value['class']!r} for {file}:{idx}")
                for spec_key in SPEC_KEYS:
                    if spec_key in value:
                        _validate_spec_list(value[spec_key], context, spec_key)
            else:
                raise ValueError(
                    f"invalid judgment value type {type(value).__name__} for {context} "
                    "(expected a class string or an object)"
                )
    return data


def wilson95(phat, n) -> list:
    """95% Wilson score interval for a proportion phat over n trials."""
    if n <= 0:
        return [0.0, 0.0]
    z = 1.959963984540054
    z2 = z * z
    denom = 1 + z2 / n
    center = phat + z2 / (2 * n)
    margin = z * math.sqrt((phat * (1 - phat) / n) + (z2 / (4 * n * n)))
    lo = (center - margin) / denom
    hi = (center + margin) / denom
    return [max(0.0, lo), min(1.0, hi)]


def _percentile(sorted_vals, pct) -> float:
    """Linear-interpolation percentile (numpy's default 'linear' method),
    no numpy dependency."""
    k = (len(sorted_vals) - 1) * (pct / 100)
    f, c = math.floor(k), math.ceil(k)
    if f == c:
        return sorted_vals[int(k)]
    return sorted_vals[f] * (c - k) + sorted_vals[c] * (k - f)


def _cluster_bootstrap(file_items: dict, seed=12345, n_boot=2000):
    """95% CI for accuracy via cluster bootstrap over FILES, not items.

    Wilson (and any per-item binomial CI) assumes n independent trials, but
    the unit here is a table and tables cluster by file — one PDF's tagger,
    document, and table idiom correlate its tables' errors. A per-item
    interval on clustered data is falsely narrow (verified by independent
    simulation in PR review, round 2: a nominal-95% Wilson interval covered
    the true rate only ~78% of the time on data shaped like the real corpus,
    81 tables / 22 files).

    Resamples distinct files WITH replacement to the original file count,
    pools the acted items of the drawn files (a file drawn twice
    contributes its items twice), recomputes accuracy, repeats `n_boot`
    times, and returns the [2.5th, 97.5th] percentile. Deterministic under
    a fixed `seed`. Returns None when fewer than 2 distinct files have
    acted items — a bootstrap over 0 or 1 clusters is meaningless.
    """
    files = sorted(file_items)
    if len(files) < 2:
        return None
    rng = random.Random(seed)
    accs = []
    for _ in range(n_boot):
        drawn = rng.choices(files, k=len(files))
        total = correct = 0
        for f in drawn:
            items = file_items[f]
            total += len(items)
            correct += sum(items)
        if total:
            accs.append(correct / total)
    if not accs:
        return None
    accs.sort()
    return [_percentile(accs, 2.5), _percentile(accs, 97.5)]


def _parse_judgment(raw):
    """Split a raw judgment value (bare class string, spec-carrying object,
    or None/absent) into `(cls_or_None, spec_or_None)`.

    `spec_or_None` is `None` exactly when no action spec was given at all —
    a bare string, or nothing — as opposed to an object that supplies the
    class but happens to omit one or more header keys (there `spec` is a
    dict, just missing keys, which the caller treats as `[]` per key)."""
    if raw is None:
        return None, None
    if isinstance(raw, str):
        return raw, None
    return raw.get("class"), {k: raw[k] for k in SPEC_KEYS if k in raw}


def score(labels, judgments, seed=12345) -> dict:
    """Score judgments against labels. See module docstring for the two
    files' roles.

    `harmful` = acted items where judged != truth AND judged not in
    ("other", "layout_wrapper"). `other` means no structural action was
    taken (a wrong `other` is a miss). `layout_wrapper` (a single-cell
    bordered box used purely to emphasize a paragraph — not a data table,
    and not a list either since there's nothing to list) is ALSO
    non-actionable today: no unwrap-to-paragraph engine exists yet in
    pdf_table_weave, so a `layout_wrapper` judgment triggers no structural
    rewrite and a wrong one is a miss, not destruction.
    # ponytail: once an unwrap action ships for layout_wrapper, move it back
    # into the actionable set below (drop it from this exclusion tuple).
    Every other class drives a real rewrite — promote_headers (invents /TH
    on a table that may have none, which "would have a screen reader
    announce ordinary data as column/row headers... which actively
    misleads", per pdf_table_weave's own module docstring),
    convert_tables_to_lists, or convert_tables_to_figures (collapses the
    descendant struct hierarchy to a single /Figure + /Alt — the original
    content structure does not survive). A wrong rewrite is destructive
    regardless of what the element actually was, so harm is a property of
    the ACTION taken, not of the truth class — this intentionally covers
    layout_table->figure, other->figure, layout_table->data_table, and
    other->layout_table, not just data_table mischaracterized.

    `acted_accuracy`/`wilson95`/`cluster95` are None when acted == 0 (no
    evidence to state a confident number). `cluster95` is additionally
    None when fewer than 2 distinct files contributed acted items — see
    `_cluster_bootstrap`. `wilson95` is kept for reference but assumes
    per-item independence, which clustered table data violates; `cluster95`
    is the corrected, headline interval.

    Everything above scores the CLASS. The `spec` block scores the ACTION
    SPEC — for a data_table, the `header_rows`/`header_cols`/`band_rows`
    the engine actually consumes to decide which cells become `/TH` and how
    `/Headers` associations get written. That decision is measured
    separately because it has been wrong far more often than the class:
    getting "this is a data_table" right does not mean the header structure
    written into the PDF was right.

    `spec.scorable` = labeled items with truth == "data_table" whose label
    carries at least one of the three header keys, whose class was judged
    correctly (acted, not abstained), AND whose judgment actually supplied
    an action spec (an object, not a bare class string). You cannot grade
    the header structure of a table the judge didn't think was a table —
    those land in `spec.skipped_wrong_class` instead, not in
    `spec.scorable`.

    A judgment given as a bare class string (no spec object at all)
    supplied no structural decision to grade at all — that is closer to an
    abstention than to a wrong answer, exactly parallel to how `unsure`/
    absent is excluded from `acted`/`acted_accuracy` on the classification
    side above. So it is EXCLUDED from `spec.scorable` and from the
    `spec.exact_accuracy` denominator, not folded in as an implicit miss —
    otherwise a judge that appropriately declines to specify would show a
    flattering accuracy over an artificially tiny (or zero) denominator
    while looking identical to one that specifies and gets it wrong. It is
    tallied and reported separately as `spec.judgment_missing_spec` so
    "declined to specify" stays visible rather than disappearing into
    either the numerator or the denominator.

    `spec.by_key` isolates which of the three keys tends to be wrong
    (order-insensitive; a key missing from either the label or the
    judgment object is treated as `[]`), which is the more useful
    diagnostic in practice: the axis (which rows/cols are headers) is
    typically right before the depth (how many) is."""
    total = len(labels)
    confusion = {}
    per_class_acc = {}  # truth -> [n, correct]
    correct = 0
    abstained = 0
    harmful = 0
    labeled_keys = set()
    labeled_files = set()
    file_items = {}  # file -> [bool, ...] (judged == truth), acted items only

    spec_scorable = 0
    spec_exact = 0
    spec_skipped_wrong_class = 0
    spec_judgment_missing = 0
    spec_by_key = {k: {"n": 0, "correct": 0} for k in SPEC_KEYS}
    spec_file_items = {}  # file -> [bool, ...] (exact match), spec-scorable items only

    for item in labels:
        file, index, truth = item["file"], item["index"], item["truth"]
        labeled_keys.add((file, index))
        labeled_files.add(file)
        judged_cls, judged_spec = _parse_judgment(judgments.get(file, {}).get(str(index)))
        row = confusion.setdefault(truth, {})

        if judged_cls is None or judged_cls == "unsure":
            abstained += 1
            row["abstain"] = row.get("abstain", 0) + 1
            continue

        is_correct = judged_cls == truth
        row[judged_cls] = row.get(judged_cls, 0) + 1
        n, c = per_class_acc.setdefault(truth, [0, 0])
        if is_correct:
            c += 1
            correct += 1
        per_class_acc[truth] = [n + 1, c]
        file_items.setdefault(file, []).append(is_correct)

        if not is_correct and judged_cls not in ("other", "layout_wrapper"):
            harmful += 1

        has_spec_label = any(k in item for k in SPEC_KEYS)
        if truth == "data_table" and has_spec_label:
            if not is_correct:
                spec_skipped_wrong_class += 1
            elif judged_spec is None:
                # Declined to specify — closer to an abstention than a
                # wrong answer. Reported on its own, never folded into
                # spec_scorable/spec_exact/spec_by_key/spec_file_items.
                spec_judgment_missing += 1
            else:
                spec_scorable += 1
                exact_match = True
                for k in SPEC_KEYS:
                    label_vals = tuple(sorted(item.get(k, [])))
                    judged_vals = tuple(sorted(judged_spec.get(k, [])))
                    spec_by_key[k]["n"] += 1
                    if label_vals == judged_vals:
                        spec_by_key[k]["correct"] += 1
                    else:
                        exact_match = False
                if exact_match:
                    spec_exact += 1
                spec_file_items.setdefault(file, []).append(exact_match)

    acted = total - abstained
    acted_accuracy = (correct / acted) if acted else None
    coverage = acted / total if total else 0.0
    per_class = {
        t: {"n": n, "correct": c, "accuracy": (c / n if n else 0.0)}
        for t, (n, c) in per_class_acc.items()
    }

    unlabeled = []
    for file, idx_map in judgments.items():
        for idx_str in idx_map:
            idx = int(idx_str)
            if (file, idx) not in labeled_keys:
                unlabeled.append({"file": file, "index": idx})
    unlabeled.sort(key=lambda r: (r["file"], r["index"]))

    acted_files = sorted(file_items)
    mean_tables_per_file = (total / len(labeled_files)) if labeled_files else 0.0
    wilson = wilson95(acted_accuracy, acted) if acted else None
    cluster = _cluster_bootstrap(file_items, seed=seed) if len(acted_files) >= 2 else None

    spec_exact_accuracy = (spec_exact / spec_scorable) if spec_scorable else None
    spec_cluster = (
        _cluster_bootstrap(spec_file_items, seed=seed)
        if len(spec_file_items) >= 2 else None
    )
    spec_by_key_out = {
        k: {"n": v["n"], "correct": v["correct"],
            "accuracy": (v["correct"] / v["n"] if v["n"] else 0.0)}
        for k, v in spec_by_key.items()
    }

    return {
        "total": total,
        "files": len(labeled_files),
        "abstained": abstained,
        "acted": acted,
        "acted_files": len(acted_files),
        "mean_tables_per_file": mean_tables_per_file,
        "correct": correct,
        "acted_accuracy": acted_accuracy,
        "coverage": coverage,
        "harmful": harmful,
        "confusion": confusion,
        "per_class": per_class,
        "wilson95": wilson,
        "cluster95": cluster,
        "unlabeled_judgments": unlabeled,
        "spec": {
            "scorable": spec_scorable,
            "exact": spec_exact,
            "exact_accuracy": spec_exact_accuracy,
            "skipped_wrong_class": spec_skipped_wrong_class,
            "judgment_missing_spec": spec_judgment_missing,
            "by_key": spec_by_key_out,
            "cluster95": spec_cluster,
        },
    }


def format_report(m: dict) -> str:
    lines = [
        "=" * 64,
        "CLASSIFICATION — which /Table class (data_table / layout_table / "
        "figure / other / layout_wrapper)",
        "=" * 64,
        f"total={m['total']}  files={m['files']}  "
        f"mean_tables_per_file={m['mean_tables_per_file']:.2f}",
        f"acted={m['acted']}  acted_files={m['acted_files']}  abstained={m['abstained']}  "
        f"coverage={m['coverage']:.1%}",
    ]
    if m["acted_accuracy"] is None:
        lines.append("acted_accuracy=n/a (no acted items)")
    else:
        lines.append(
            f"acted_accuracy={m['acted_accuracy']:.1%}  correct={m['correct']}/{m['acted']}"
        )
    if m["cluster95"] is None:
        lines.append("cluster95=n/a (needs >=2 files)  <- headline CI (cluster bootstrap by file)")
    else:
        lo, hi = m["cluster95"]
        lines.append(f"cluster95=[{lo:.1%}, {hi:.1%}]  <- headline CI (cluster bootstrap by file)")
    if m["wilson95"] is None:
        lines.append("wilson95=n/a  (uncorrected, assumes independence)")
    else:
        lo, hi = m["wilson95"]
        lines.append(f"wilson95=[{lo:.1%}, {hi:.1%}]  (uncorrected, assumes independence)")
    lines.append(
        f"harmful={m['harmful']}  (judged != truth and judged not in "
        "('other', 'layout_wrapper') — a wrong classification that drives a structural "
        "PDF rewrite; wrong 'other'/'layout_wrapper' is a miss, not a rewrite)"
    )
    lines.append("")
    lines.append("confusion (truth -> judged/abstain counts):")
    for truth in sorted(m["confusion"]):
        row = m["confusion"][truth]
        row_str = ", ".join(f"{k}={v}" for k, v in sorted(row.items()))
        lines.append(f"  {truth}: {row_str}")
    lines.append("")
    lines.append("per_class accuracy (acted items only):")
    for truth in sorted(m["per_class"]):
        pc = m["per_class"][truth]
        lines.append(f"  {truth}: {pc['correct']}/{pc['n']} = {pc['accuracy']:.1%}")
    if m["unlabeled_judgments"]:
        lines.append("")
        lines.append(f"unlabeled_judgments ({len(m['unlabeled_judgments'])}, not scored):")
        for u in m["unlabeled_judgments"]:
            lines.append(f"  {u['file']}:{u['index']}")

    spec = m["spec"]
    lines.append("")
    lines.append("=" * 64)
    lines.append("ACTION SPEC — data_table header structure "
                  "(header_rows / header_cols / band_rows)")
    lines.append("=" * 64)
    lines.append(
        f"scorable={spec['scorable']}  skipped_wrong_class={spec['skipped_wrong_class']}  "
        f"judgment_missing_spec={spec['judgment_missing_spec']}"
    )
    if spec["exact_accuracy"] is None:
        lines.append("spec_exact_accuracy=n/a (no scorable items)")
    else:
        lines.append(
            f"spec_exact_accuracy={spec['exact_accuracy']:.1%}  "
            f"exact={spec['exact']}/{spec['scorable']}"
        )
    if spec["cluster95"] is None:
        lines.append(
            "spec_cluster95=n/a (needs >=2 files)  <- headline CI (cluster bootstrap by file)"
        )
    else:
        lo, hi = spec["cluster95"]
        lines.append(
            f"spec_cluster95=[{lo:.1%}, {hi:.1%}]  <- headline CI (cluster bootstrap by file)"
        )
    lines.append("")
    lines.append("by_key (which part of the structure is wrong, over scorable items):")
    for k in SPEC_KEYS:
        bk = spec["by_key"][k]
        lines.append(f"  {k}: {bk['correct']}/{bk['n']} = {bk['accuracy']:.1%}")

    return "\n".join(lines)


# ── CLI ──────────────────────────────────────────────────────────────────

def main() -> int:
    parser = argparse.ArgumentParser(description="AI table-classification accuracy harness")
    parser.add_argument("--candidates", nargs="+", metavar="PATH")
    parser.add_argument("--out", help="output manifest path for --candidates")
    parser.add_argument("--render-dir", help="render each candidate's page to a PNG here")
    parser.add_argument("--score", action="store_true")
    parser.add_argument("--labels", help="labels.jsonl for --score")
    parser.add_argument("--judgments", help="judgments.json for --score")
    parser.add_argument("--json", action="store_true", help="emit metrics as JSON (--score only)")
    parser.add_argument(
        "--seed", type=int, default=12345,
        help="cluster-bootstrap RNG seed, for a deterministic cluster95 (default 12345)",
    )
    args = parser.parse_args()

    if args.candidates:
        if not args.out:
            print("ERROR: --out required with --candidates", file=sys.stderr)
            return 2
        candidates, pdf_files = build_candidates(args.candidates, args.render_dir)
        manifest = {"generated_from": args.candidates, "candidates": candidates}
        Path(args.out).write_text(json.dumps(manifest, indent=2), encoding="utf-8")
        unreadable = sum(1 for c in candidates if "error" in c)
        files_with_candidates = len({c["file"] for c in candidates})
        print(
            f"{len(candidates)} candidates across {files_with_candidates} files "
            f"({len(pdf_files)} PDFs swept, {unreadable} unreadable) -> {args.out}"
        )
        return 0

    if args.score:
        if not args.labels or not args.judgments:
            print("ERROR: --labels and --judgments required with --score", file=sys.stderr)
            return 2
        try:
            labels = load_labels(args.labels)
            judgments = load_judgments(args.judgments)
        except (ValueError, OSError, json.JSONDecodeError) as exc:
            print(f"ERROR: {exc}", file=sys.stderr)
            return 2
        metrics = score(labels, judgments, seed=args.seed)
        print(json.dumps(metrics, indent=2) if args.json else format_report(metrics))
        return 0

    print("ERROR: one of --candidates or --score is required", file=sys.stderr)
    return 2


if __name__ == "__main__":
    sys.exit(main())
