tools/expand_dmhy_dag_annotations.py

"""Expand DAG annotation patches over the full DMHY source list.

The DAG annotation runner samples terminal.value_examples. This tool keeps the
LLM patch layer intact, but applies it to every dmhy_list.jsonl value by
deterministic terminal prefix matching.
"""

from __future__ import annotations

import argparse
import json
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

from anifilebert.tokenizer import AnimeTokenizer
from tools.annotate_dmhy_prefix_graph import (
    build_prefix_trie,
    matching_terminal_ordinal,
    normalize_generated_tokens,
    prefix_boundary_ok as prefix_boundary_ok,
)
from tools.dmhy_dataset import weak_label_filename


DEFAULT_DAG = Path("datasets/AnimeName/dmhy_prefix_dag.json")
DEFAULT_SOURCE_LIST = Path("datasets/AnimeName/dmhy_list.jsonl")
DEFAULT_PATCHES = Path("datasets/AnimeName/dmhy_prefix_dag.annotations.dag_full.jsonl")
DEFAULT_OUTPUT = Path("datasets/AnimeName/dmhy_weak.dag_full_expanded.jsonl")


@dataclass(frozen=True)
class Args:
    dag: Path
    source_list: Path
    patches: Path
    output: Path
    manifest_output: Path
    limit: int | None
    preserve_i_labels: bool
    strict_boundary: bool


def parse_args() -> Args:
    parser = argparse.ArgumentParser(
        description="Expand DMHY DAG terminal annotations across the full source-list"
    )
    parser.add_argument("--dag", type=Path, default=DEFAULT_DAG)
    parser.add_argument("--source-list", type=Path, default=DEFAULT_SOURCE_LIST)
    parser.add_argument("--patches", type=Path, default=DEFAULT_PATCHES)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument(
        "--manifest-output",
        type=Path,
        default=None,
        help="Output manifest JSON; defaults to <output stem>.manifest.json",
    )
    parser.add_argument("--limit", type=int, default=None, help="Maximum source-list rows to scan")
    parser.add_argument("--preserve-i-labels", action="store_true")
    parser.add_argument(
        "--strict-boundary",
        action="store_true",
        help="Require a separator boundary after the matched terminal prefix",
    )
    ns = parser.parse_args()
    manifest_output = ns.manifest_output or ns.output.with_suffix(".manifest.json")
    return Args(
        dag=ns.dag,
        source_list=ns.source_list,
        patches=ns.patches,
        output=ns.output,
        manifest_output=manifest_output,
        limit=None if ns.limit is None else max(0, ns.limit),
        preserve_i_labels=ns.preserve_i_labels,
        strict_boundary=ns.strict_boundary,
    )


def load_json(path: Path) -> dict[str, Any]:
    if not path.exists():
        raise SystemExit(f"file not found: {path}")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise SystemExit(f"invalid JSON in {path}: {exc}") from exc
    if not isinstance(value, dict):
        raise SystemExit(f"invalid JSON schema in {path}: root must be object")
    return value


def iter_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
    if not path.exists():
        raise SystemExit(f"file not found: {path}")
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise SystemExit(f"invalid JSON in {path}:{line_number}: {exc}") from exc
            if isinstance(row, dict):
                yield line_number, row


def terminal_id(terminal: dict[str, Any], fallback: int) -> str:
    for key in ("terminal_id", "id", "node_id"):
        value = terminal.get(key)
        if value is not None:
            return str(value)
    return str(fallback)


def load_terminals(dag_path: Path) -> list[tuple[int, dict[str, Any], dict[str, Any]]]:
    dag = load_json(dag_path)
    terminals = dag.get("terminals")
    if not isinstance(terminals, list):
        raise SystemExit(f"invalid DAG schema in {dag_path}: missing terminals list")

    selected: list[tuple[int, dict[str, Any], dict[str, Any]]] = []
    for index, terminal in enumerate(terminals):
        if not isinstance(terminal, dict):
            continue
        row = dict(terminal)
        row["_terminal_id"] = terminal_id(row, index)
        row["_terminal_index"] = index
        if not str(row.get("prefix") or ""):
            continue
        selected.append((index, row, {}))
    if not selected:
        raise SystemExit(f"DAG has no prefix-bearing terminals: {dag_path}")
    return selected


def load_patches(path: Path) -> tuple[dict[str, dict[str, Any]], Counter[str]]:
    patches_by_terminal: dict[str, dict[str, Any]] = {}
    patch_statuses: Counter[str] = Counter()
    for _line_number, patch in iter_jsonl(path):
        terminal_ids = patch.get("terminal_ids")
        if not isinstance(terminal_ids, list):
            terminal_ids = [patch.get("terminal_id")] if patch.get("terminal_id") is not None else []
        clean_ids = [str(item) for item in terminal_ids if str(item).strip()]
        status = str(patch.get("status") or ("fallback" if patch.get("fallback") else "ok"))
        patch_statuses[status] += 1
        for tid in clean_ids:
            patches_by_terminal[tid] = patch
    if not patches_by_terminal:
        raise SystemExit(f"no terminal patches found in {path}")
    return patches_by_terminal, patch_statuses


def matching_terminal_ordinal_relaxed(
    value: str,
    trie: Any,
    selected: list[tuple[int, dict[str, Any], dict[str, Any]]],
    *,
    strict_boundary: bool,
) -> int | None:
    node = trie
    best: int | None = None
    for char in value:
        node = node.children.get(char)
        if node is None:
            break
        for ordinal in node.terminal_ordinals:
            if not strict_boundary:
                best = ordinal
                continue
            prefix = str(selected[ordinal][1].get("prefix") or "")
            if prefix_boundary_ok(value, prefix):
                best = ordinal
    return best


def row_from_match(
    *,
    filename: str,
    source_line: int,
    terminal: dict[str, Any],
    patch: dict[str, Any],
    tokenizer: AnimeTokenizer,
    preserve_i_labels: bool,
) -> dict[str, Any] | None:
    sample = weak_label_filename(filename, tokenizer)
    if sample is None:
        return None
    tokens, labels = normalize_generated_tokens(
        sample["tokens"],
        sample["labels"],
        preserve_i_labels=preserve_i_labels,
    )
    terminal_id_value = str(terminal["_terminal_id"])
    unit_id = str(patch.get("unit_id") or terminal_id_value)
    terminal_index = terminal["_terminal_index"]
    source = str(patch.get("source") or "unknown")
    fallback = bool(patch.get("fallback", False))
    status = str(patch.get("status") or ("fallback" if fallback else "ok"))
    episode_title_suffixes = patch.get("episode_title_suffixes") or []
    media_suffixes = patch.get("media_suffixes") or []
    title_candidates = patch.get("title_candidates") or []
    return {
        "file_id": f"prefix-dag-expanded:{unit_id}:{terminal_id_value}:{source_line}",
        "filename": filename,
        "tokens": tokens,
        "labels": labels,
        "terminal_id": terminal_id_value,
        "terminal_index": terminal_index,
        "unit_id": unit_id,
        "source": source,
        "needs_llm_review": fallback,
        "episode_title_suffixes": episode_title_suffixes,
        "media_suffixes": media_suffixes,
        "title_candidates": title_candidates,
        "annotations": {
            "unit_id": unit_id,
            "terminal_id": terminal_id_value,
            "terminal_ids": [str(item) for item in patch.get("terminal_ids") or [terminal_id_value]],
            "terminal_index": terminal_index,
            "source": source,
            "fallback": fallback,
            "status": status,
            "episode_title_suffixes": episode_title_suffixes,
            "media_suffixes": media_suffixes,
            "title_candidates": title_candidates,
            "llm_label": patch.get("llm_label"),
            "notes": patch.get("notes"),
        },
    }


def expand(args: Args) -> dict[str, Any]:
    selected = load_terminals(args.dag)
    patches_by_terminal, patch_statuses = load_patches(args.patches)
    trie = build_prefix_trie(selected)
    tokenizer = AnimeTokenizer()

    input_rows = 0
    matched_source_rows = 0
    emitted_rows = 0
    unmatched_rows = 0
    weak_label_failed_rows = 0
    patch_covered_rows = 0
    unique_terminals_matched: set[str] = set()
    status_counts: Counter[str] = Counter()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8", newline="\n") as output_handle:
        for source_line, source_row in iter_jsonl(args.source_list):
            if args.limit is not None and input_rows >= args.limit:
                break
            input_rows += 1
            value = source_row.get("value")
            if not isinstance(value, str) or not value.strip():
                unmatched_rows += 1
                status_counts["unmatched"] += 1
                continue

            ordinal = matching_terminal_ordinal_relaxed(
                value,
                trie,
                selected,
                strict_boundary=args.strict_boundary,
            )
            if ordinal is None:
                unmatched_rows += 1
                status_counts["unmatched"] += 1
                continue

            matched_source_rows += 1
            terminal = selected[ordinal][1]
            terminal_id_value = str(terminal["_terminal_id"])
            unique_terminals_matched.add(terminal_id_value)
            patch = patches_by_terminal.get(terminal_id_value)
            if patch is None:
                status_counts["missing_patch"] += 1
                continue

            row = row_from_match(
                filename=value,
                source_line=source_line,
                terminal=terminal,
                patch=patch,
                tokenizer=tokenizer,
                preserve_i_labels=args.preserve_i_labels,
            )
            if row is None:
                weak_label_failed_rows += 1
                status_counts["weak_label_failed"] += 1
                continue

            output_handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
            emitted_rows += 1
            patch_covered_rows += 1
            status_counts[str(row["annotations"]["status"])] += 1

    manifest = {
        "dag": str(args.dag),
        "source_list": str(args.source_list),
        "patches": str(args.patches),
        "output": str(args.output),
        "input_rows": input_rows,
        "matched_rows": matched_source_rows,
        "output_rows": emitted_rows,
        "unmatched_rows": unmatched_rows,
        "weak_label_failed_rows": weak_label_failed_rows,
        "patch_covered_rows": patch_covered_rows,
        "unique_terminals_matched": len(unique_terminals_matched),
        "status_counts": dict(sorted(status_counts.items())),
        "patch_status_counts": dict(sorted(patch_statuses.items())),
    }
    args.manifest_output.parent.mkdir(parents=True, exist_ok=True)
    args.manifest_output.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return manifest


def main() -> None:
    manifest = expand(parse_args())
    print(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()