#!/usr/bin/env python3
"""Score an automated audiobook cast against a hand-corrected one, by word.

  python3 scripts/score_attribution.py <manual_cast.json> <auto_cast.json>
         [--characters <manual_characters.json>] [--confusions N] [--json]

The older scripts/compare_casts.py matches whole segments by exact text, so any
difference in where a segment starts or ends counts as an error even when every
word is spoken by the right voice. This scorer aligns the two casts word by
word instead, which makes it insensitive to how narration is chunked and lets
runs with different segmentation be compared fairly.

Reported (all as % of the manual cast's words unless stated):
  coverage        words of the manual cast found in the auto cast
  type            narration/dialogue agreement over aligned words
  dialogue        of the words the human labelled dialogue: correct speaker,
                  strict and alias-tolerant (a variant such as "Baronin" for
                  "Baronin Ira von Seewiesen", or a listed alias, counts as right)
  speaker (both)  speaker accuracy over words BOTH casts label dialogue - the
                  pure attribution quality, separate from segmentation quality
  overall         every word right (type and, for dialogue, speaker), tolerant
"""
from __future__ import annotations

import argparse
import collections
import difflib
import json
import re
import sys
from pathlib import Path

_QUOTES = re.compile(r"[»«„“”\"'`‚‘’›‹]")
_ARTICLE = re.compile(r"^(?:der|die|das|den|dem|ein|eine)\s+")
_GLUE = {"von", "der", "die", "und", "zu", "vom"}


def _clean_name(name: str) -> str:
    return _ARTICLE.sub("", str(name or "").strip().lower())


def _tokens(name: str) -> set[str]:
    return set(re.findall(r"\w+", _clean_name(name)))


def load_segments(path: str | Path) -> list[dict]:
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    return (data.get("segments") if isinstance(data, dict) else data) or []


def load_aliases(path: str | Path | None) -> dict[str, set[str]]:
    """{canonical name: {cleaned alias, ...}} from a characters export."""
    aliases: dict[str, set[str]] = collections.defaultdict(set)
    if not path:
        return aliases
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    for c in data if isinstance(data, list) else data.get("characters", []):
        name = c.get("name")
        if not name:
            continue
        aliases[name].add(_clean_name(name))
        raw = (c.get("sheet") or {}).get("aliases") or ""
        for a in str(raw).split(","):
            if a.strip():
                aliases[name].add(_clean_name(a))
    return aliases


def speaker_matches(human: str, auto: str, aliases: dict[str, set[str]]) -> bool:
    """Strict equality, a listed alias, or a name that is a token subset of the other."""
    if human == auto:
        return True
    if _clean_name(auto) in aliases.get(human, ()):
        return True
    ta, th = _tokens(auto), _tokens(human)
    if ta and th and (ta <= th or th <= ta) and not ((ta & th) <= _GLUE):
        return True
    return False


def _words(segments: list[dict]) -> tuple[list[str], list[tuple[bool, str]]]:
    words, labels = [], []
    for seg in segments:
        dialogue = seg.get("type") == "dialogue"
        speaker = str(seg.get("speaker") or "Narrator") if dialogue else "Narrator"
        for w in re.findall(r"\w+", _QUOTES.sub(" ", str(seg.get("text") or "")).lower()):
            words.append(w)
            labels.append((dialogue, speaker))
    return words, labels


def score(manual: list[dict], auto: list[dict], aliases: dict[str, set[str]] | None = None, top: int = 25) -> dict:
    aliases = aliases or {}
    mw, ml = _words(manual)
    aw, al = _words(auto)
    blocks = difflib.SequenceMatcher(None, mw, aw, autojunk=False).get_matching_blocks()
    aligned = type_ok = 0
    d_total = d_strict = d_tol = 0            # human-dialogue words
    both = both_strict = both_tol = 0         # dialogue in both casts
    overall_tol = 0
    confusions: collections.Counter = collections.Counter()
    per_speaker: dict[str, list[int]] = collections.defaultdict(lambda: [0, 0])
    unknown_words = 0
    for b in blocks:
        for k in range(b.size):
            (hd, hs), (ad, as_) = ml[b.a + k], al[b.b + k]
            aligned += 1
            if hd == ad:
                type_ok += 1
            if ad and str(as_).lower().startswith(("unknown", "unbekannt")):
                unknown_words += 1
            tol = strict = False
            if hd and ad:
                strict, tol = hs == as_, speaker_matches(hs, as_, aliases)
                both += 1
                both_strict += strict
                both_tol += tol
            elif not hd and not ad:
                strict = tol = True
            if hd:
                d_total += 1
                d_strict += strict
                d_tol += tol
                per_speaker[hs][0] += 1
                per_speaker[hs][1] += tol
            overall_tol += tol
            if not tol:
                confusions[(hs, as_)] += 1
    n = len(mw) or 1
    pct = lambda x, d: round(100.0 * x / d, 2) if d else 0.0
    return {
        "manual_words": len(mw), "auto_words": len(aw),
        "coverage": pct(aligned, n),
        "type": pct(type_ok, aligned),
        "dialogue_strict": pct(d_strict, d_total), "dialogue_tolerant": pct(d_tol, d_total),
        "dialogue_words": d_total,
        "speaker_both_strict": pct(both_strict, both), "speaker_both_tolerant": pct(both_tol, both),
        "overall_tolerant": pct(overall_tol, n),
        "auto_unknown_words": unknown_words,
        "auto_speakers": len({s.get("speaker") for s in auto if s.get("type") == "dialogue"}),
        "manual_speakers": len({s.get("speaker") for s in manual if s.get("type") == "dialogue"}),
        "top_confusions": [
            {"human": h, "auto": a, "words": c} for (h, a), c in confusions.most_common(top)
        ],
        "weakest_speakers": [
            {"speaker": s, "words": t[0], "accuracy": pct(t[1], t[0])}
            for s, t in sorted(per_speaker.items(), key=lambda kv: -kv[1][0])[:14]
        ],
    }


def render(res: dict) -> str:
    lines = [
        f"words: manual={res['manual_words']}  auto={res['auto_words']}  coverage={res['coverage']}%",
        f"narration/dialogue type agreement : {res['type']}%",
        f"dialogue words ({res['dialogue_words']}) right speaker : strict {res['dialogue_strict']}%  tolerant {res['dialogue_tolerant']}%",
        f"speaker accuracy where both say dialogue : strict {res['speaker_both_strict']}%  tolerant {res['speaker_both_tolerant']}%",
        f"overall every-word-right (tolerant) : {res['overall_tolerant']}%",
        f"speakers: manual={res['manual_speakers']} auto={res['auto_speakers']}  auto Unknown words={res['auto_unknown_words']}",
        "", "top remaining confusions (human -> auto, words):",
    ]
    lines += [f"  {c['words']:6d}  {c['human']:26s} -> {c['auto']}" for c in res["top_confusions"]]
    lines += ["", "biggest human speakers (tolerant accuracy):"]
    lines += [f"  {s['speaker']:26s} {s['words']:6d} words  {s['accuracy']:5.1f}%" for s in res["weakest_speakers"]]
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("manual")
    ap.add_argument("auto")
    ap.add_argument("--characters", help="manual characters export, for alias tolerance")
    ap.add_argument("--confusions", type=int, default=25)
    ap.add_argument("--json", action="store_true", help="print machine-readable JSON")
    args = ap.parse_args(argv)
    res = score(load_segments(args.manual), load_segments(args.auto), load_aliases(args.characters), args.confusions)
    print(json.dumps(res, ensure_ascii=False, indent=1) if args.json else render(res))
    return 0


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