tools/annotate_dmhy_prefix_dag.py

"""Build annotation units from a DMHY prefix DAG.

The DAG producer keeps repeated suffix structure shared across many raw
prefixes. This tool turns those shared nodes into compact, traceable units for
LLM or human review without calling any remote service.
"""

from __future__ import annotations

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

from tools.annotate_dmhy_prefix_graph import (
    heuristic_patch,
    string_list,
)


DEFAULT_DAG = Path("datasets/AnimeName/dmhy_prefix_dag.json")
DEFAULT_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_dag.annotation_units.jsonl")


@dataclass(frozen=True)
class Args:
    dag: Path
    output: Path
    min_reachable_terminals: int
    min_incoming_count: int
    limit: int | None
    example_count: int
    include_terminal_sink_units: bool


def parse_args() -> Args:
    parser = argparse.ArgumentParser(
        description="Emit DAG-aware DMHY prefix annotation units as JSONL"
    )
    parser.add_argument("--dag", type=Path, default=DEFAULT_DAG, help="Input dmhy_prefix_dag.json")
    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT,
        help="Output annotation unit JSONL",
    )
    parser.add_argument(
        "--min-reachable-terminals",
        type=int,
        default=2,
        help="Select non-root nodes reaching at least this many terminals",
    )
    parser.add_argument(
        "--min-incoming-count",
        type=int,
        default=2,
        help="Select nodes with at least this many incoming DAG edges",
    )
    parser.add_argument("--limit", type=int, default=None, help="Maximum units to write")
    parser.add_argument(
        "--example-count",
        type=int,
        default=5,
        help="Maximum examples retained per example field",
    )
    parser.add_argument(
        "--include-terminal-sink-units",
        action="store_true",
        help="Include terminal nodes with no outgoing edges as annotation units",
    )
    ns = parser.parse_args()
    return Args(
        dag=ns.dag,
        output=ns.output,
        min_reachable_terminals=max(1, ns.min_reachable_terminals),
        min_incoming_count=max(1, ns.min_incoming_count),
        limit=ns.limit,
        example_count=max(1, ns.example_count),
        include_terminal_sink_units=ns.include_terminal_sink_units,
    )


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


def node_id(node: dict[str, Any], fallback: int) -> int:
    value = node.get("id", fallback)
    try:
        return int(value)
    except (TypeError, ValueError):
        raise SystemExit(f"invalid node id: {value!r}") from None


def terminal_id(terminal: dict[str, Any], fallback: int) -> str:
    value = terminal.get("terminal_id", terminal.get("id", fallback))
    return str(value)


def int_field(row: dict[str, Any], key: str, default: int = 0) -> int:
    try:
        return int(row.get(key, default) or default)
    except (TypeError, ValueError):
        return default


def build_indexes(dag: dict[str, Any]) -> tuple[dict[int, dict[str, Any]], dict[int, list[dict[str, Any]]]]:
    nodes: dict[int, dict[str, Any]] = {}
    for fallback, node in enumerate(dag["nodes"]):
        if not isinstance(node, dict):
            continue
        nodes[node_id(node, fallback)] = node

    terminals_by_node: dict[int, list[dict[str, Any]]] = {}
    for fallback, terminal in enumerate(dag["terminals"]):
        if not isinstance(terminal, dict):
            continue
        try:
            terminal_node_id = int(terminal.get("node_id"))
        except (TypeError, ValueError):
            continue
        terminal = dict(terminal)
        terminal["_terminal_id"] = terminal_id(terminal, fallback)
        terminal["_terminal_index"] = fallback
        terminals_by_node.setdefault(terminal_node_id, []).append(terminal)
    return nodes, terminals_by_node


class ReachableTerminalIndex:
    def __init__(
        self,
        nodes: dict[int, dict[str, Any]],
        terminals_by_node: dict[int, list[dict[str, Any]]],
    ) -> None:
        self.nodes = nodes
        self.terminals_by_node = terminals_by_node
        self.memo: dict[int, list[dict[str, Any]]] = {}
        self.visiting: set[int] = set()

    def get(self, start_node_id: int) -> list[dict[str, Any]]:
        if start_node_id in self.memo:
            return self.memo[start_node_id]
        if start_node_id in self.visiting:
            raise SystemExit(f"cycle detected while traversing DAG at node {start_node_id}")
        self.visiting.add(start_node_id)
        found = list(self.terminals_by_node.get(start_node_id, []))
        node = self.nodes.get(start_node_id, {})
        for edge in node.get("children") or []:
            if not isinstance(edge, dict):
                continue
            try:
                target = int(edge.get("target"))
            except (TypeError, ValueError):
                continue
            found.extend(self.get(target))
        self.visiting.remove(start_node_id)
        deduped = dedupe_terminals(found)
        self.memo[start_node_id] = deduped
        return deduped


def dedupe_terminals(terminals: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
    seen: set[str] = set()
    result: list[dict[str, Any]] = []
    for terminal in terminals:
        tid = str(terminal.get("_terminal_id") or terminal.get("terminal_id") or "")
        if not tid or tid in seen:
            continue
        seen.add(tid)
        result.append(terminal)
    return result


def limited_unique(values: Iterable[str], limit: int) -> list[str]:
    seen: set[str] = set()
    result: list[str] = []
    for value in values:
        if not value or not value.strip() or value in seen:
            continue
        seen.add(value)
        result.append(value)
        if len(result) >= limit:
            break
    return result


def edge_labels(node: dict[str, Any], limit: int) -> list[str]:
    labels = []
    for edge in node.get("children") or []:
        if isinstance(edge, dict) and edge.get("label") is not None:
            labels.append(str(edge["label"]))
    return limited_unique(labels, limit)


def has_outgoing_edges(node: dict[str, Any]) -> bool:
    return any(isinstance(edge, dict) for edge in node.get("children") or [])


def aggregate_examples(terminals: list[dict[str, Any]], key: str, limit: int) -> list[str]:
    values: list[str] = []
    for terminal in terminals:
        if key == "prefix":
            value = terminal.get("prefix")
            if value is not None:
                values.append(str(value))
        else:
            values.extend(string_list(terminal.get(key)))
    return limited_unique(values, limit)


def aggregate_needs_review(terminals: list[dict[str, Any]]) -> bool:
    for index, terminal in enumerate(terminals):
        annotations = terminal.get("annotations")
        if isinstance(annotations, dict) and annotations.get("needs_llm_review") is not None:
            if bool(annotations["needs_llm_review"]):
                return True
            continue
        if heuristic_patch(terminal, index)["needs_llm_review"]:
            return True
    return False


def annotation_template() -> dict[str, Any]:
    return {
        "episode_title_suffixes": [],
        "media_suffixes": [],
        "title_candidates": [],
        "notes": None,
    }


def recommended_prompt(kind: str, terminal_count: int) -> str:
    if kind == "shared_suffix":
        return (
            "Review the shared DAG suffix examples and mark episode-title text, media metadata, "
            f"and possible title candidates for {terminal_count} linked terminals."
        )
    return "Review this terminal cluster and mark episode-title text, media metadata, and title candidates."


def make_unit(
    node: dict[str, Any],
    node_id_value: int,
    terminals: list[dict[str, Any]],
    kind: str,
    example_count: int,
) -> dict[str, Any]:
    terminal_ids = [str(terminal["_terminal_id"]) for terminal in terminals]
    reachable_weight = int_field(node, "reachable_weight")
    if reachable_weight <= 0:
        reachable_weight = sum(int_field(terminal, "weight", int_field(terminal, "count", 1)) for terminal in terminals)
    return {
        "unit_id": f"dag-node-{node_id_value}",
        "node_id": node_id_value,
        "kind": kind,
        "incoming_count": int_field(node, "incoming_count"),
        "reachable_terminals": len(terminals),
        "reachable_weight": reachable_weight,
        "terminal_ids": terminal_ids,
        "prefix_examples": aggregate_examples(terminals, "prefix", example_count),
        "value_examples": aggregate_examples(terminals, "value_examples", example_count),
        "suffix_examples": aggregate_examples(terminals, "suffix_examples", example_count),
        "common_edge_labels": edge_labels(node, example_count),
        "needs_llm_review": aggregate_needs_review(terminals),
        "recommended_prompt": recommended_prompt(kind, len(terminals)),
        "annotations": annotation_template(),
    }


def node_reachable_count(node: dict[str, Any]) -> int:
    return int_field(node, "reachable_terminals")


def node_reachable_weight(node: dict[str, Any]) -> int:
    return int_field(node, "reachable_weight")


def selected_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]:
    nodes, terminals_by_node = build_indexes(dag)
    reachable_index = ReachableTerminalIndex(nodes, terminals_by_node)
    root = int(dag.get("root", 0) or 0)
    shared_candidates: list[tuple[tuple[int, int, int, int, int], int, dict[str, Any]]] = []

    for current_id, node in nodes.items():
        if not args.include_terminal_sink_units and not has_outgoing_edges(node):
            continue
        incoming_count = int_field(node, "incoming_count")
        reachable_count = node_reachable_count(node)
        by_shared_incoming = incoming_count >= args.min_incoming_count
        by_reachable = current_id != root and reachable_count >= args.min_reachable_terminals
        if by_shared_incoming or by_reachable:
            sort_key = (
                0,
                -node_reachable_weight(node),
                -reachable_count,
                -incoming_count,
                current_id,
            )
            shared_candidates.append((sort_key, current_id, node))

    shared_candidates.sort(key=lambda item: item[0])
    candidates: list[tuple[tuple[int, int, int, int, int], dict[str, Any]]] = []
    for sort_key, current_id, node in shared_candidates:
        if args.limit is not None and len(candidates) >= args.limit:
            break
        terminals = reachable_index.get(current_id)
        if not terminals:
            continue
        unit = make_unit(node, current_id, terminals, "shared_suffix", args.example_count)
        candidates.append((sort_key, unit))

    if args.limit is not None and len(candidates) >= args.limit:
        return [unit for _sort_key, unit in candidates[: max(0, args.limit)]]

    covered_terminal_ids = {
        terminal_id
        for _sort_key, unit in candidates
        for terminal_id in unit["terminal_ids"]
    }
    for current_id, terminals in terminals_by_node.items():
        node = nodes.get(current_id, {"id": current_id})
        if not args.include_terminal_sink_units and not has_outgoing_edges(node):
            continue
        uncovered = [terminal for terminal in terminals if terminal["_terminal_id"] not in covered_terminal_ids]
        if not uncovered:
            continue
        unit = make_unit(node, current_id, uncovered, "terminal_cluster", args.example_count)
        sort_key = (
            1,
            -unit["reachable_weight"],
            -unit["reachable_terminals"],
            -unit["incoming_count"],
            current_id,
        )
        candidates.append((sort_key, unit))

    candidates.sort(key=lambda item: item[0])
    units = [unit for _sort_key, unit in candidates]
    if args.limit is not None:
        units = units[: max(0, args.limit)]
    return units


def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with path.open("w", encoding="utf-8", newline="\n") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
            count += 1
    return count


def main() -> None:
    args = parse_args()
    dag = load_dag(args.dag)
    units = selected_units(dag, args)
    count = write_jsonl(args.output, units)
    summary = {
        "dag": str(args.dag),
        "output": str(args.output),
        "annotation_units": count,
        "min_reachable_terminals": args.min_reachable_terminals,
        "min_incoming_count": args.min_incoming_count,
        "example_count": args.example_count,
    }
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()