scripts/eval_pipeline.py

"""End-to-end evaluation harness for the Prompt Squirrel RAG pipeline.

Measures per-stage and overall metrics using ground-truth tagged samples
from the e621 evaluation dataset.

Metrics computed:
  - Stage 2 (Retrieval): Recall@k — what fraction of ground-truth tags
    appear among the retrieved candidates
  - Stage 3 (Selection): Precision, Recall, F1 — how well the final
    selected tags match the ground truth

Usage:
    # Full end-to-end (Stage 1 + 2 + 3), 20 random samples:
    python scripts/eval_pipeline.py --n 20

    # Reproducible run with specific seed:
    python scripts/eval_pipeline.py --n 50 --seed 123

    # Parallel processing with 4 workers (default):
    python scripts/eval_pipeline.py --n 50 --workers 4

    # Sequential mode (disable parallelism):
    python scripts/eval_pipeline.py --n 20 --workers 1

    # Skip Stage 1 LLM rewrite (cheaper, tests Stage 2+3 only):
    python scripts/eval_pipeline.py --n 20 --skip-rewrite

    # First N samples in file order (no shuffle):
    python scripts/eval_pipeline.py --n 20 --no-shuffle

Results are always saved as JSONL to data/eval_results/ (auto-named by timestamp)
or to a custom path with -o.

Requires:
    - OPENROUTER_API_KEY env var (for Stage 1 rewrite and Stage 3 selection)
    - fluffyrock_3m.csv and other retrieval assets in the project root
    - data/eval_samples/e621_sfw_sample_1000_seed123_buffer10000.jsonl
"""

from __future__ import annotations

import argparse
import json
import os
import random
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple

_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))
os.chdir(_REPO_ROOT)


def _ensure_utf8_stdio() -> None:
    try:
        if hasattr(sys.stdout, "reconfigure"):
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        if hasattr(sys.stderr, "reconfigure"):
            sys.stderr.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass

EVAL_DATA_PATH = _REPO_ROOT / "data" / "eval_samples" / "e621_sfw_sample_1000_seed123_buffer10000_expanded.jsonl"
EVAL_DATA_PATH_RAW = _REPO_ROOT / "data" / "eval_samples" / "e621_sfw_sample_1000_seed123_buffer10000.jsonl"

# Character tag types that go through the alias filter pipeline
_CHARACTER_TYPES = {"character"}
# Copyright tags are filtered out entirely
_COPYRIGHT_TYPES = {"copyright"}

# Tags excluded from evaluation metrics but NOT removed from the pipeline.
# These are tags that either: can't be inferred from a caption (resolution,
# art medium), describe structural properties better handled outside the
# retrieval pipeline (backgrounds), or are annotation artifacts.
_EVAL_EXCLUDED_TAGS = frozenset({
    # Annotation artifacts
    "invalid_tag", "invalid_background",
    # Resolution / file meta — not inferrable from caption
    "hi_res", "absurd_res", "low_res", "superabsurd_res",
    # Structural background tags — better recommended independently
    "simple_background", "abstract_background", "detailed_background",
    "gradient_background", "blurred_background", "textured_background",
    "transparent_background", "white_background",
})


def _classify_tags(tags: Set[str], get_type_fn) -> Tuple[Set[str], Set[str]]:
    """Split tags into (character_tags, general_tags).

    Copyright tags are excluded from both sets since they're filtered
    before any selection happens.
    """
    character = set()
    general = set()
    for tag in tags:
        ttype = get_type_fn(tag)
        if ttype in _CHARACTER_TYPES:
            character.add(tag)
        elif ttype not in _COPYRIGHT_TYPES:
            general.add(tag)
    return character, general


def _flatten_ground_truth_tags(tags_categorized_str: str) -> Set[str]:
    """Parse the categorized ground-truth JSON string into a flat set of tags."""
    if not tags_categorized_str:
        return set()
    try:
        cats = json.loads(tags_categorized_str)
    except json.JSONDecodeError:
        return set()
    tags = set()
    for tag_list in cats.values():
        if isinstance(tag_list, list):
            for t in tag_list:
                tags.add(t.strip())
    return tags


@dataclass
class SampleResult:
    sample_id: Any
    caption: str
    ground_truth_tags: Set[str]
    # Stage 1
    rewrite_phrases: List[str] = field(default_factory=list)
    # Stage 2
    retrieved_tags: Set[str] = field(default_factory=set)
    retrieval_recall: float = 0.0
    # Stage 3 — overall
    selected_tags: Set[str] = field(default_factory=set)
    stage3_selected_tags: Set[str] = field(default_factory=set)
    stage3_selected_scores: Dict[str, float] = field(default_factory=dict)
    stage3_selected_ranks: Dict[str, int] = field(default_factory=dict)
    stage3_selected_phrase_ranks: Dict[str, int] = field(default_factory=dict)
    selection_precision: float = 0.0
    selection_recall: float = 0.0
    selection_f1: float = 0.0
    # Stage 3 — character tags only
    gt_character_tags: Set[str] = field(default_factory=set)
    selected_character_tags: Set[str] = field(default_factory=set)
    retrieved_character_tags: Set[str] = field(default_factory=set)
    char_retrieval_recall: float = 0.0
    char_precision: float = 0.0
    char_recall: float = 0.0
    char_f1: float = 0.0
    # Stage 3 — general tags only (non-character, non-copyright)
    gt_general_tags: Set[str] = field(default_factory=set)
    selected_general_tags: Set[str] = field(default_factory=set)
    general_precision: float = 0.0
    general_recall: float = 0.0
    general_f1: float = 0.0
    # New diagnostic metrics
    retrieval_precision: float = 0.0       # |retrieved ∩ gt| / |retrieved|
    selection_given_retrieval: float = 0.0  # |selected ∩ gt| / |retrieved ∩ gt|
    over_selection_ratio: float = 0.0       # |selected| / |gt|
    # Why distribution (from Stage 3 LLM)
    why_counts: Dict[str, int] = field(default_factory=dict)
    stage3_diag: Dict[str, Any] = field(default_factory=dict)
    # Tag implications
    implied_tags: Set[str] = field(default_factory=set)  # tags added via implications (not LLM-selected)
    # Structural inference tags (solo/duo/male/female/anthro/biped etc.)
    structural_tags: List[str] = field(default_factory=list)
    # Simplified probe tags (reliability-gated fixed probe list)
    probe_tags: List[str] = field(default_factory=list)
    # Per-tag evidence: tag -> {"source": "stage3"|"structural"|"implied", "why": ..., "score": ...}
    tag_evidence: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    # Leaf-only metrics (strips implied ancestors from both sides)
    leaf_precision: float = 0.0
    leaf_recall: float = 0.0
    leaf_f1: float = 0.0
    leaf_selected_count: int = 0
    leaf_gt_count: int = 0
    # Timing
    stage1_time: float = 0.0
    stage2_time: float = 0.0
    stage3_time: float = 0.0
    stage3s_time: float = 0.0
    stage3p_time: float = 0.0
    # Categorized suggestions (for ranking metrics)
    categorized_suggestions: Dict[str, List[Tuple[str, float]]] = field(default_factory=dict)
    # Errors
    error: Optional[str] = None
    # Non-fatal issues/warnings captured from pipeline logs (fallbacks, retries, API errors)
    issues: List[str] = field(default_factory=list)


def _compute_metrics(predicted: Set[str], ground_truth: Set[str]) -> Tuple[float, float, float]:
    """Compute precision, recall, F1."""
    if not predicted and not ground_truth:
        return 1.0, 1.0, 1.0
    if not predicted:
        return 0.0, 0.0, 0.0
    if not ground_truth:
        return 0.0, 0.0, 0.0

    tp = len(predicted & ground_truth)
    precision = tp / len(predicted)
    recall = tp / len(ground_truth)
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
    return precision, recall, f1


def _build_selection_query(
    caption: str,
    rewrite_phrases: Sequence[str],
    structural_tags: Sequence[str],
    probe_tags: Sequence[str],
) -> str:
    lines = [f"IMAGE DESCRIPTION: {caption.strip()}"]
    if rewrite_phrases:
        lines.append("REWRITE PHRASES: " + ", ".join(rewrite_phrases))
    hint_tags = list(structural_tags) + list(probe_tags)
    if hint_tags:
        lines.append("INFERRED TAG HINTS (context only): " + ", ".join(sorted(set(hint_tags))))
    return "\n".join(lines)


def _process_one_sample(
    sample: Dict[str, Any],
    index: int,
    total: int,
    skip_rewrite: bool,
    allow_nsfw: bool,
    mode: str,
    chunk_size: int,
    per_phrase_k: int,
    per_phrase_final_k: int,
    temperature: float,
    max_tokens: int,
    verbose: bool,
    print_lock: threading.Lock,
    min_why: Optional[str] = None,
    expand_implications: bool = False,
    infer_structural: bool = False,
    infer_probe: bool = False,
) -> SampleResult:
    """Process a single eval sample through the full pipeline. Thread-safe."""
    from psq_rag.llm.rewrite import llm_rewrite_prompt
    from psq_rag.retrieval.psq_retrieval import psq_candidates_from_rewrite_phrases
    from psq_rag.llm.select import llm_select_indices, llm_infer_structural_tags, llm_infer_probe_tags
    from psq_rag.retrieval.state import get_tag_type_name, expand_tags_via_implications, get_leaf_tags

    def log(msg: str) -> None:
        msg_str = str(msg)
        msg_l = msg_str.lower()
        if any(k in msg_l for k in ("error", "fallback", "gave up", "warning", "filtered", "refusal")):
            result.issues.append(msg_str)
        if verbose:
            with print_lock:
                print(f"  [{index+1}] {msg_str}")

    sid = sample["id"]
    caption = sample["caption"]
    gt_tags = sample["gt_tags"]

    result = SampleResult(
        sample_id=sid,
        caption=caption[:120] + ("..." if len(caption) > 120 else ""),
        ground_truth_tags=gt_tags,
    )

    with print_lock:
        print(f"[{index+1}/{total}] id={sid} gt_tags={len(gt_tags)}")

    try:
        # --- Stage 1x: Start independent LLM calls concurrently ---
        def _run_stage1_rewrite() -> Tuple[str, float]:
            t0 = time.time()
            rewritten_local = llm_rewrite_prompt(caption, log)
            return rewritten_local or "", (time.time() - t0)

        def _run_stage3_structural() -> Tuple[List[str], float]:
            t0s = time.time()
            structural = llm_infer_structural_tags(
                caption, log=log, temperature=temperature,
            )
            return structural, (time.time() - t0s)

        def _run_stage3_probe() -> Tuple[List[str], float]:
            t0p = time.time()
            probed = llm_infer_probe_tags(
                caption, log=log, temperature=temperature,
            )
            return probed, (time.time() - t0p)

        pre_workers = (0 if skip_rewrite else 1) + int(infer_structural) + int(infer_probe)
        if pre_workers > 0:
            with ThreadPoolExecutor(max_workers=pre_workers) as pre_ex:
                fut_rewrite = pre_ex.submit(_run_stage1_rewrite) if not skip_rewrite else None
                fut_struct = pre_ex.submit(_run_stage3_structural) if infer_structural else None
                fut_probe = pre_ex.submit(_run_stage3_probe) if infer_probe else None

                # --- Stage 1: LLM Rewrite (result consumed first for retrieval) ---
                if skip_rewrite:
                    phrases = [p.strip() for p in caption.split(",") if p.strip()]
                    if len(phrases) <= 1:
                        phrases = [p.strip() for p in caption.replace(".", ",").split(",") if p.strip()]
                    result.rewrite_phrases = phrases
                    result.stage1_time = 0.0
                else:
                    rewritten, t1 = fut_rewrite.result() if fut_rewrite is not None else ("", 0.0)
                    result.stage1_time = t1
                    if rewritten:
                        result.rewrite_phrases = [p.strip() for p in rewritten.split(",") if p.strip()]
                    else:
                        result.rewrite_phrases = [p.strip() for p in caption.split(",") if p.strip()]
                        if len(result.rewrite_phrases) <= 1:
                            result.rewrite_phrases = [p.strip() for p in caption.replace(".", ",").split(",") if p.strip()]

                log(f"Phrases ({len(result.rewrite_phrases)}): {result.rewrite_phrases[:5]}")

                # Wait for side-channel calls before retrieval so their tags can
                # influence TF-IDF context scoring in Stage 2.
                if fut_struct is not None:
                    structural, stage3s_time = fut_struct.result()
                    result.stage3s_time = stage3s_time
                    result.structural_tags = structural
                else:
                    result.stage3s_time = 0.0
                    result.structural_tags = []

                if fut_probe is not None:
                    probe_tags, stage3p_time = fut_probe.result()
                    result.stage3p_time = stage3p_time
                    result.probe_tags = probe_tags
                else:
                    result.stage3p_time = 0.0
                    result.probe_tags = []

                # --- Stage 2: Retrieval ---
                t0 = time.time()
                retrieval_context_tags = list(dict.fromkeys(result.structural_tags + result.probe_tags))
                retrieval_result = psq_candidates_from_rewrite_phrases(
                    rewrite_phrases=result.rewrite_phrases,
                    allow_nsfw_tags=allow_nsfw,
                    context_tags=retrieval_context_tags,
                    per_phrase_final_k=per_phrase_final_k,
                    global_k=300,
                    return_phrase_ranks=True,
                    verbose=False,
                )
                result.stage2_time = time.time() - t0
        else:
            # Should not happen with current defaults, but keep safe behavior.
            result.rewrite_phrases = [p.strip() for p in caption.split(",") if p.strip()]
            result.stage1_time = 0.0
            result.stage2_time = 0.0
            result.stage3s_time = 0.0
            result.structural_tags = []
            result.stage3p_time = 0.0
            result.probe_tags = []
            retrieval_result = []

        phrase_rank_by_tag = {}
        if isinstance(retrieval_result, tuple):
            if len(retrieval_result) == 2:
                candidates, phrase_rank_by_tag = retrieval_result
            else:
                candidates = retrieval_result[0]
        else:
            candidates = retrieval_result

        result.retrieved_tags = {c.tag for c in candidates}
        if gt_tags:
            result.retrieval_recall = len(result.retrieved_tags & gt_tags) / len(gt_tags)

        log(f"Retrieved {len(candidates)} candidates, recall={result.retrieval_recall:.3f}")

        # --- Stage 3: LLM Selection (uses rewrite + structural/probe context) ---
        def _run_stage3_selection() -> Tuple[List[int], Dict[str, str], Dict[str, Any], float]:
            t0 = time.time()
            selection_query = _build_selection_query(
                caption=caption,
                rewrite_phrases=result.rewrite_phrases,
                structural_tags=result.structural_tags,
                probe_tags=result.probe_tags,
            )
            picked, why_map, diag = llm_select_indices(
                query_text=selection_query,
                candidates=candidates,
                max_pick=0,
                log=log,
                mode=mode,
                chunk_size=chunk_size,
                per_phrase_k=per_phrase_k,
                temperature=temperature,
                max_tokens=max_tokens,
                return_metadata=True,
                return_diagnostics=True,
                min_why=min_why,
            )
            return picked, why_map, diag or {}, (time.time() - t0)

        picked_indices, tag_why, stage3_diag, stage3_time = _run_stage3_selection()

        result.stage3_time = stage3_time
        result.stage3_diag = stage3_diag or {}

        result.selected_tags = {candidates[idx].tag for idx in picked_indices} if picked_indices else set()
        result.stage3_selected_tags = set(result.selected_tags)

        # Build per-tag evidence from Stage 3 selection
        rank_by_tag = {c.tag: i + 1 for i, c in enumerate(candidates)}
        for idx in picked_indices:
            tag = candidates[idx].tag
            result.stage3_selected_scores[tag] = round(candidates[idx].score_combined, 4)
            result.stage3_selected_ranks[tag] = rank_by_tag.get(tag, len(candidates) + 1)
            if phrase_rank_by_tag:
                result.stage3_selected_phrase_ranks[tag] = phrase_rank_by_tag.get(tag, len(candidates) + 1)
            result.tag_evidence[tag] = {
                "source": "stage3",
                "why": tag_why.get(tag, "unknown"),
                "retrieval_score": round(candidates[idx].score_combined, 4),
            }

        # Why distribution
        why_counts: Dict[str, int] = {}
        for w in tag_why.values():
            why_counts[w] = why_counts.get(w, 0) + 1
        result.why_counts = why_counts

        # Structural tag inference (solo/duo/male/female/anthro/biped etc.)
        if infer_structural:
            # Add structural tags not already selected
            for st in result.structural_tags:
                if st not in result.selected_tags:
                    result.tag_evidence[st] = {"source": "structural"}
                result.selected_tags.add(st)
            log(f"Structural: {result.structural_tags}")

        if infer_probe:
            for pt in result.probe_tags:
                if pt not in result.selected_tags:
                    result.tag_evidence[pt] = {"source": "probe"}
                result.selected_tags.add(pt)
            log(f"Probe: {result.probe_tags}")

        # Tag implication expansion (post-Stage 3)
        if expand_implications and result.selected_tags:
            expanded, implied_only = expand_tags_via_implications(result.selected_tags)
            result.implied_tags = implied_only
            for imp_tag in implied_only:
                result.tag_evidence[imp_tag] = {"source": "implied"}
            result.selected_tags = expanded
            log(f"Implications: +{len(implied_only)} tags")

        # Generate categorized suggestions (for ranking metrics)
        try:
            from psq_rag.tagging.categorized_suggestions import (
                generate_categorized_suggestions,
            )

            # Use selected tags to generate category-wise ranked suggestions
            categorized = generate_categorized_suggestions(
                selected_tags=list(result.selected_tags),
                allow_nsfw_tags=allow_nsfw,
                top_n_per_category=20,  # Get top 20 per category for eval
                top_n_other=50,
            )

            # Convert to simple dict format: category -> [(tag, score), ...]
            result.categorized_suggestions = {}
            for cat_name, cat_sugg in categorized.by_category.items():
                result.categorized_suggestions[cat_name] = cat_sugg.suggestions

            # Also store "other" suggestions
            result.categorized_suggestions['other'] = categorized.other_suggestions

            log(f"Categorized: {len(result.categorized_suggestions)} categories")
        except Exception as e:
            log(f"Warning: Failed to generate categorized suggestions: {e}")

        # Remove eval-excluded tags from predictions before scoring
        result.selected_tags -= _EVAL_EXCLUDED_TAGS
        result.retrieved_tags -= _EVAL_EXCLUDED_TAGS

        # Overall selection metrics (expanded — both sides have full implication chains)
        p, r, f1 = _compute_metrics(result.selected_tags, gt_tags)
        result.selection_precision = p
        result.selection_recall = r
        result.selection_f1 = f1

        # Leaf-only metrics (strips implied ancestors from both sides)
        leaf_sel = get_leaf_tags(result.selected_tags)
        leaf_gt = get_leaf_tags(gt_tags)
        lp, lr, lf1 = _compute_metrics(leaf_sel, leaf_gt)
        result.leaf_precision = lp
        result.leaf_recall = lr
        result.leaf_f1 = lf1
        result.leaf_selected_count = len(leaf_sel)
        result.leaf_gt_count = len(leaf_gt)

        # Diagnostic metrics
        retrieved_and_gt = result.retrieved_tags & gt_tags
        selected_and_gt = result.selected_tags & gt_tags
        if result.retrieved_tags:
            result.retrieval_precision = len(retrieved_and_gt) / len(result.retrieved_tags)
        if retrieved_and_gt:
            result.selection_given_retrieval = len(selected_and_gt) / len(retrieved_and_gt)
        if gt_tags:
            result.over_selection_ratio = len(result.selected_tags) / len(gt_tags)

        # Split ground-truth and selected tags by type
        gt_char, gt_gen = _classify_tags(gt_tags, get_tag_type_name)
        sel_char, sel_gen = _classify_tags(result.selected_tags, get_tag_type_name)
        ret_char, _ = _classify_tags(result.retrieved_tags, get_tag_type_name)

        result.gt_character_tags = gt_char
        result.selected_character_tags = sel_char
        result.retrieved_character_tags = ret_char
        result.gt_general_tags = gt_gen
        result.selected_general_tags = sel_gen

        # Character-specific metrics
        if gt_char:
            result.char_retrieval_recall = len(ret_char & gt_char) / len(gt_char)
        cp, cr, cf1 = _compute_metrics(sel_char, gt_char)
        result.char_precision = cp
        result.char_recall = cr
        result.char_f1 = cf1

        # General-tag metrics
        gp, gr, gf1 = _compute_metrics(sel_gen, gt_gen)
        result.general_precision = gp
        result.general_recall = gr
        result.general_f1 = gf1

        # Per-sample output line
        char_info = ""
        if gt_char:
            char_info = f" char[gt={len(gt_char)} sel={len(sel_char)} P={cp:.2f} R={cr:.2f}]"
        impl_info = f" (+{len(result.implied_tags)} implied)" if result.implied_tags else ""
        struct_info = f" (+{len(result.structural_tags)} structural)" if result.structural_tags else ""
        probe_info = f" (+{len(result.probe_tags)} probe)" if result.probe_tags else ""
        with print_lock:
            print(
                f"  [{index+1}] retrieval_recall={result.retrieval_recall:.3f} "
                f"sel_P={p:.3f} sel_R={r:.3f} sel_F1={f1:.3f} "
                f"selected={len(result.selected_tags)}{impl_info}{struct_info}{probe_info}{char_info} "
                f"t1={result.stage1_time:.1f}s t2={result.stage2_time:.1f}s t3={result.stage3_time:.1f}s"
            )

    except Exception as e:
        result.error = str(e)
        result.issues.append(f"fatal_exception: {e}")
        with print_lock:
            print(f"  [{index+1}] ERROR: {e}")

    return result


def _prewarm_retrieval_assets() -> None:
    """Force-load all lazy retrieval assets so threads don't race on init."""
    from psq_rag.retrieval.state import (
        get_tfidf_components,
        get_tag2aliases,
        get_tag_type_name,
        get_tag_implications,
    )
    print("Pre-warming retrieval assets (TF-IDF, FastText, HNSW, aliases)...")
    t0 = time.time()
    get_tfidf_components()  # loads joblib, HNSW indexes, FastText model
    get_tag2aliases()       # loads CSV alias dict
    get_tag_type_name("_warmup_")  # ensures tag type dict is built
    get_tag_implications()  # loads implication graph
    print(f"  Assets loaded in {time.time() - t0:.1f}s")


def run_eval(
    n_samples: int = 20,
    caption_field: str = "caption_cogvlm",
    skip_rewrite: bool = False,
    allow_nsfw: bool = False,
    mode: str = "chunked_map_union",
    chunk_size: int = 60,
    per_phrase_k: int = 2,
    per_phrase_final_k: int = 1,
    temperature: float = 0.0,
    max_tokens: int = 512,
    verbose: bool = False,
    shuffle: bool = True,
    seed: int = 42,
    workers: int = 1,
    min_why: Optional[str] = "strong_implied",
    eval_path: Optional[str] = None,
    expand_implications: bool = False,
    infer_structural: bool = False,
    infer_probe: bool = False,
) -> List[SampleResult]:
    expand_gt = expand_implications
    if expand_gt:
        from psq_rag.retrieval.state import expand_tags_via_implications as _expand_gt_tags

    # Load eval samples — prefer expanded file, fall back to raw
    eval_path_obj = Path(eval_path) if eval_path else EVAL_DATA_PATH
    if not eval_path_obj.is_absolute():
        eval_path_obj = (_REPO_ROOT / eval_path_obj).resolve()

    if not eval_path_obj.is_file() and eval_path is None:
        eval_path_obj = EVAL_DATA_PATH_RAW
        if not eval_path_obj.is_file():
            print(f"ERROR: Eval data not found: {EVAL_DATA_PATH}")
            sys.exit(1)
        print(f"WARNING: Expanded eval data not found, falling back to raw: {eval_path_obj}")
        print("  Run: python scripts/preprocess_eval_data.py")
    elif not eval_path_obj.is_file():
        print(f"ERROR: Eval data not found: {eval_path_obj}")
        sys.exit(1)

    all_samples = []
    using_expanded = False
    with eval_path_obj.open("r", encoding="utf-8") as f:
        for line in f:
            row = json.loads(line)
            caption = row.get(caption_field, "")
            if not caption or not caption.strip():
                continue
            # Prefer pre-expanded GT; fall back to flattening categorized
            if "tags_ground_truth_expanded" in row:
                gt_tags = set(row["tags_ground_truth_expanded"])
                using_expanded = True
            else:
                gt_tags = _flatten_ground_truth_tags(row.get("tags_ground_truth_categorized", ""))
            if not gt_tags:
                continue
            # Remove eval-excluded tags from GT
            gt_tags -= _EVAL_EXCLUDED_TAGS
            if expand_gt:
                gt_tags, _ = _expand_gt_tags(gt_tags)
                gt_tags -= _EVAL_EXCLUDED_TAGS
            all_samples.append({
                "id": row.get("id", row.get("row_id", len(all_samples))),
                "caption": caption.strip(),
                "gt_tags": gt_tags,
            })
    if using_expanded:
        print("Using implication-expanded ground truth")

    if shuffle:
        rng = random.Random(seed)
        rng.shuffle(all_samples)

    samples = all_samples[:n_samples]

    print(f"Loaded {len(samples)}/{len(all_samples)} samples (caption_field={caption_field})")
    print(f"eval_path={eval_path_obj}")
    print(f"shuffle={shuffle}, seed={seed}, skip_rewrite={skip_rewrite}, allow_nsfw={allow_nsfw}, mode={mode}")
    print(f"workers={workers}")
    print()

    # Pre-warm shared retrieval assets before spawning threads
    _prewarm_retrieval_assets()
    print()

    print_lock = threading.Lock()
    total = len(samples)

    if workers <= 1:
        # Sequential mode (original behavior)
        results: List[SampleResult] = []
        for i, sample in enumerate(samples):
            result = _process_one_sample(
                sample, i, total,
                skip_rewrite, allow_nsfw, mode, chunk_size,
                per_phrase_k, per_phrase_final_k, temperature, max_tokens, verbose,
                print_lock, min_why,
                expand_implications,
                infer_structural,
                infer_probe,
            )
            results.append(result)
    else:
        # Parallel mode
        print(f"Processing {total} samples with {workers} parallel workers...")
        print()
        # Submit all samples; use index to preserve original ordering
        results_by_index: Dict[int, SampleResult] = {}
        with ThreadPoolExecutor(max_workers=workers) as executor:
            futures = {
                executor.submit(
                    _process_one_sample,
                    sample, i, total,
                    skip_rewrite, allow_nsfw, mode, chunk_size,
                    per_phrase_k, per_phrase_final_k, temperature, max_tokens, verbose,
                    print_lock, min_why,
                    expand_implications,
                    infer_structural,
                    infer_probe,
                ): i
                for i, sample in enumerate(samples)
            }
            for future in as_completed(futures):
                idx = futures[future]
                try:
                    results_by_index[idx] = future.result()
                except Exception as e:
                    # Should not happen since _process_one_sample catches exceptions,
                    # but guard against unexpected errors
                    with print_lock:
                        print(f"  [{idx+1}] WORKER ERROR: {e}")
                    result = SampleResult(
                        sample_id=samples[idx]["id"],
                        caption=samples[idx]["caption"][:120],
                        ground_truth_tags=samples[idx]["gt_tags"],
                        error=f"Worker error: {e}",
                    )
                    results_by_index[idx] = result

        # Reassemble in original order
        results = [results_by_index[i] for i in range(total)]

    return results


def _safe_avg(values: List[float]) -> float:
    return sum(values) / len(values) if values else 0.0


def print_summary(results: List[SampleResult]) -> None:
    """Print aggregate metrics across all samples."""
    valid = [r for r in results if r.error is None]
    errored = [r for r in results if r.error is not None]

    if not valid:
        print("\nNo valid results to summarize.")
        return

    n = len(valid)

    avg_retrieval_recall = sum(r.retrieval_recall for r in valid) / n
    avg_sel_precision = sum(r.selection_precision for r in valid) / n
    avg_sel_recall = sum(r.selection_recall for r in valid) / n
    avg_sel_f1 = sum(r.selection_f1 for r in valid) / n

    avg_retrieved = sum(len(r.retrieved_tags) for r in valid) / n
    avg_selected = sum(len(r.selected_tags) for r in valid) / n
    avg_gt = sum(len(r.ground_truth_tags) for r in valid) / n

    avg_t1 = sum(r.stage1_time for r in valid) / n
    avg_t2 = sum(r.stage2_time for r in valid) / n
    avg_t3 = sum(r.stage3_time for r in valid) / n

    print()
    print("=" * 70)
    print(f"EVALUATION SUMMARY ({n} samples, {len(errored)} errors)")
    print("=" * 70)
    print()
    print("Stage 2 - Retrieval:")
    print(f"  Avg recall@300:       {avg_retrieval_recall:.4f}")
    print(f"  Avg candidates:       {avg_retrieved:.1f}")
    avg_retrieval_precision = _safe_avg([r.retrieval_precision for r in valid])
    avg_sel_given_ret = _safe_avg([r.selection_given_retrieval for r in valid
                                   if (r.retrieved_tags & r.ground_truth_tags)])
    avg_over_sel = _safe_avg([r.over_selection_ratio for r in valid])

    avg_implied = sum(len(r.implied_tags) for r in valid) / n
    avg_structural = sum(len(r.structural_tags) for r in valid) / n
    avg_probe = sum(len(r.probe_tags) for r in valid) / n

    print()
    print("Stage 3 - Selection (ALL tags):")
    print(f"  Avg precision:        {avg_sel_precision:.4f}")
    print(f"  Avg recall:           {avg_sel_recall:.4f}")
    print(f"  Avg F1:               {avg_sel_f1:.4f}")
    print(f"  Avg selected tags:    {avg_selected:.1f}")
    if avg_implied > 0:
        print(f"  Avg implied tags:     {avg_implied:.1f}  (added via tag implications)")
    if avg_structural > 0:
        print(f"  Avg structural tags:  {avg_structural:.1f}  (inferred via statement agreement)")
    if avg_probe > 0:
        print(f"  Avg probe tags:       {avg_probe:.1f}  (inferred via simplified probe query)")
    print(f"  Avg ground-truth tags:{avg_gt:.1f}")

    # Leaf-only metrics
    avg_leaf_p = _safe_avg([r.leaf_precision for r in valid])
    avg_leaf_r = _safe_avg([r.leaf_recall for r in valid])
    avg_leaf_f1 = _safe_avg([r.leaf_f1 for r in valid])
    avg_leaf_sel = _safe_avg([r.leaf_selected_count for r in valid])
    avg_leaf_gt = _safe_avg([r.leaf_gt_count for r in valid])
    print()
    print("Stage 3 - Selection (LEAF tags only — implied ancestors stripped):")
    print(f"  Avg precision:        {avg_leaf_p:.4f}")
    print(f"  Avg recall:           {avg_leaf_r:.4f}")
    print(f"  Avg F1:               {avg_leaf_f1:.4f}")
    print(f"  Avg leaf selected:    {avg_leaf_sel:.1f}")
    print(f"  Avg leaf ground-truth:{avg_leaf_gt:.1f}")

    print()
    print("Diagnostic Metrics:")
    print(f"  Retrieval precision:  {avg_retrieval_precision:.4f}  (|ret∩gt|/|ret|, noise level fed to Stage 3)")
    print(f"  Sel-given-retrieval:  {avg_sel_given_ret:.4f}  (of gt tags retrieved, fraction kept by Stage 3)")
    print(f"  Over-selection ratio: {avg_over_sel:.2f}x  (|selected|/|gt|, ideal ~1.0)")

    stage3_diag_rows = [r.stage3_diag for r in valid if r.stage3_diag]
    if stage3_diag_rows:
        calls_total = sum(int(d.get("calls_total", 0)) for d in stage3_diag_rows)
        calls_exhausted = sum(int(d.get("calls_exhausted_retries", 0)) for d in stage3_diag_rows)
        attempts_total = sum(int(d.get("attempts_total", 0)) for d in stage3_diag_rows)
        attempts_parse_fail = sum(int(d.get("attempt_parse_fail", 0)) for d in stage3_diag_rows)
        attempts_errors = sum(int(d.get("attempt_errors", 0)) for d in stage3_diag_rows)

        print()
        print("Stage 3 Structured Output Reliability:")
        print(f"  Calls total:          {calls_total}")
        print(f"  Calls exhausted:      {calls_exhausted} ({(100 * calls_exhausted / calls_total) if calls_total else 0:.1f}%)")
        print(f"  Attempts total:       {attempts_total}")
        print(f"  Parse/schema failures:{attempts_parse_fail} ({(100 * attempts_parse_fail / attempts_total) if attempts_total else 0:.1f}%)")
        print(f"  Call errors/exc:      {attempts_errors} ({(100 * attempts_errors / attempts_total) if attempts_total else 0:.1f}%)")

        by_n_agg: Dict[int, Dict[str, int]] = {}
        for d in stage3_diag_rows:
            for n_str, n_stats in d.get("attempts_by_n_local", {}).items():
                try:
                    n_local = int(n_str)
                except Exception:
                    continue
                cur = by_n_agg.setdefault(n_local, {"attempts": 0, "parse_fail": 0, "errors": 0})
                cur["attempts"] += int(n_stats.get("attempts", 0))
                cur["parse_fail"] += int(n_stats.get("parse_fail", 0))
                cur["errors"] += int(n_stats.get("errors", 0))

        if by_n_agg:
            print("  Failure by call size (N_local):")
            for n_local in sorted(by_n_agg.keys()):
                s = by_n_agg[n_local]
                fail = s["parse_fail"] + s["errors"]
                rate = (100 * fail / s["attempts"]) if s["attempts"] else 0.0
                print(
                    f"    N={n_local:3d} attempts={s['attempts']:4d} "
                    f"fail={fail:4d} ({rate:5.1f}%)"
                )

    # Why distribution across all samples
    total_why: Dict[str, int] = {}
    for r in valid:
        for w, cnt in r.why_counts.items():
            total_why[w] = total_why.get(w, 0) + cnt
    if total_why:
        total_selections = sum(total_why.values())
        print()
        print("Why Distribution (Stage 3 LLM rationale):")
        for w in ["explicit", "strong_implied", "weak_implied", "style_or_meta", "other"]:
            cnt = total_why.get(w, 0)
            pct = 100 * cnt / total_selections if total_selections else 0
            print(f"  {w:20s} {cnt:4d}  ({pct:5.1f}%)")

    # --- Character tag breakdown ---
    # Only include samples that actually have character tags in ground truth
    samples_with_chars = [r for r in valid if r.gt_character_tags]
    # Samples where the system selected character tags (true or false positive)
    samples_selecting_chars = [r for r in valid if r.selected_character_tags]

    print()
    print("-" * 70)
    print(f"CHARACTER TAGS ({len(samples_with_chars)}/{n} samples have character ground-truth)")
    print("-" * 70)

    if samples_with_chars:
        avg_char_retrieval_recall = _safe_avg([r.char_retrieval_recall for r in samples_with_chars])
        avg_char_p = _safe_avg([r.char_precision for r in samples_with_chars])
        avg_char_r = _safe_avg([r.char_recall for r in samples_with_chars])
        avg_char_f1 = _safe_avg([r.char_f1 for r in samples_with_chars])
        avg_gt_char = _safe_avg([len(r.gt_character_tags) for r in samples_with_chars])
        avg_sel_char = _safe_avg([len(r.selected_character_tags) for r in samples_with_chars])

        print(f"  Retrieval recall:     {avg_char_retrieval_recall:.4f}")
        print(f"  Selection precision:  {avg_char_p:.4f}")
        print(f"  Selection recall:     {avg_char_r:.4f}")
        print(f"  Selection F1:         {avg_char_f1:.4f}")
        print(f"  Avg gt char tags:     {avg_gt_char:.1f}")
        print(f"  Avg selected chars:   {avg_sel_char:.1f}")

        # Show character-specific failures
        char_misses = []
        char_false_pos = []
        for r in samples_with_chars:
            missed = r.gt_character_tags - r.selected_character_tags
            for m in missed:
                char_misses.append((r.sample_id, m))
            extra = r.selected_character_tags - r.gt_character_tags
            for e in extra:
                char_false_pos.append((r.sample_id, e))

        if char_misses:
            print(f"\n  Missed characters ({len(char_misses)} total):")
            for sid, tag in char_misses[:10]:
                print(f"    id={sid}: missed {tag}")

        if char_false_pos:
            print(f"\n  False positive characters ({len(char_false_pos)} total):")
            for sid, tag in char_false_pos[:10]:
                print(f"    id={sid}: wrongly selected {tag}")
    else:
        print("  (no samples had character tags in ground truth)")

    # False positive characters in samples WITHOUT character ground-truth
    no_char_gt_but_selected = [r for r in valid if not r.gt_character_tags and r.selected_character_tags]
    if no_char_gt_but_selected:
        print(f"\n  Spurious character selections ({len(no_char_gt_but_selected)} samples):")
        print("  (These samples had NO character in ground truth but system selected one)")
        for r in no_char_gt_but_selected[:5]:
            print(f"    id={r.sample_id}: selected {sorted(r.selected_character_tags)}")

    # --- General tag breakdown ---
    print()
    print("-" * 70)
    print("GENERAL TAGS (non-character, non-copyright)")
    print("-" * 70)
    avg_gen_p = _safe_avg([r.general_precision for r in valid])
    avg_gen_r = _safe_avg([r.general_recall for r in valid])
    avg_gen_f1 = _safe_avg([r.general_f1 for r in valid])
    avg_gt_gen = _safe_avg([len(r.gt_general_tags) for r in valid])
    avg_sel_gen = _safe_avg([len(r.selected_general_tags) for r in valid])
    print(f"  Selection precision:  {avg_gen_p:.4f}")
    print(f"  Selection recall:     {avg_gen_r:.4f}")
    print(f"  Selection F1:         {avg_gen_f1:.4f}")
    print(f"  Avg gt general tags:  {avg_gt_gen:.1f}")
    print(f"  Avg selected general: {avg_sel_gen:.1f}")

    print()
    print("-" * 70)
    avg_t3s = sum(r.stage3s_time for r in valid) / n
    avg_t3p = sum(r.stage3p_time for r in valid) / n
    print("Timing (avg per sample):")
    print(f"  Stage 1 (rewrite):    {avg_t1:.2f}s")
    print(f"  Stage 2 (retrieval):  {avg_t2:.2f}s")
    print(f"  Stage 3 (selection):  {avg_t3:.2f}s")
    if avg_t3s > 0:
        print(f"  Stage 3s (structural):{avg_t3s:.2f}s")
    if avg_t3p > 0:
        print(f"  Stage 3p (probe):     {avg_t3p:.2f}s")
    print(f"  Total:                {avg_t1 + avg_t2 + avg_t3 + avg_t3s + avg_t3p:.2f}s")
    print()

    # Show worst and best F1 samples
    by_f1 = sorted(valid, key=lambda r: r.selection_f1)
    print("Lowest F1 samples (overall):")
    for r in by_f1[:3]:
        print(f"  id={r.sample_id} F1={r.selection_f1:.3f} P={r.selection_precision:.3f} R={r.selection_recall:.3f}")
        missed = r.ground_truth_tags - r.selected_tags
        extra = r.selected_tags - r.ground_truth_tags
        if missed:
            print(f"    missed: {sorted(missed)[:10]}")
        if extra:
            print(f"    extra:  {sorted(extra)[:10]}")

    print()
    print("Highest F1 samples (overall):")
    for r in by_f1[-3:]:
        print(f"  id={r.sample_id} F1={r.selection_f1:.3f} P={r.selection_precision:.3f} R={r.selection_recall:.3f}")

    if errored:
        print()
        print(f"Errors ({len(errored)}):")
        for r in errored[:5]:
            print(f"  id={r.sample_id}: {r.error}")

    print("=" * 70)


def main(argv=None) -> int:
    _ensure_utf8_stdio()
    ap = argparse.ArgumentParser(description="End-to-end pipeline evaluation")
    ap.add_argument("--n", type=int, default=20, help="Number of samples to evaluate")
    ap.add_argument("--caption-field", default="caption_cogvlm",
                    choices=["caption_cogvlm", "caption_llm_0", "caption_llm_1",
                             "caption_llm_2", "caption_llm_3", "caption_llm_4",
                             "caption_llm_5", "caption_llm_6", "caption_llm_7"],
                    help="Which caption field to use as input")
    ap.add_argument("--skip-rewrite", action="store_true",
                    help="Skip Stage 1 LLM rewrite; split caption directly into phrases")
    ap.add_argument("--allow-nsfw", action="store_true", help="Allow NSFW tags")
    ap.add_argument("--mode", default="chunked_map_union",
                    choices=["single_shot", "chunked_map_union"])
    ap.add_argument("--chunk-size", type=int, default=60)
    ap.add_argument("--per-phrase-k", type=int, default=2)
    ap.add_argument("--per-phrase-final-k", type=int, default=1,
                    help="Top-K candidates per phrase after scoring (retrieval cap)")
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--max-tokens", type=int, default=512)
    ap.add_argument("--verbose", "-v", action="store_true", help="Show per-call Stage 3 logs")
    ap.add_argument("--output", "-o", type=str, default=None,
                    help="Save detailed results as JSONL (default: auto-generated in data/eval_results/)")
    ap.add_argument("--shuffle", action="store_true", default=True,
                    help="Randomly shuffle samples before selecting (default: True)")
    ap.add_argument("--no-shuffle", dest="shuffle", action="store_false",
                    help="Use samples in file order (first N)")
    ap.add_argument("--seed", type=int, default=42,
                    help="Random seed for shuffle (default: 42)")
    ap.add_argument("--workers", "-w", type=int, default=4,
                    help="Number of parallel workers (default: 4, use 1 for sequential)")
    ap.add_argument("--eval-path", type=str, default=None,
                    help="Optional path to eval JSONL (defaults to expanded 1000-sample set).")
    ap.add_argument("--min-why", default="strong_implied",
                    choices=["explicit", "strong_implied", "weak_implied", "style_or_meta", "other", "none"],
                    help="Minimum 'why' confidence to keep (default: strong_implied). Use 'none' to disable filtering.")
    ap.add_argument("--expand-implications", action="store_true", default=False,
                    help="Expand selected tags via tag implication chains (e.g. fox→canine→canid→mammal)")
    ap.add_argument("--infer-structural", action="store_true", default=False,
                    help="Infer structural tags (solo/duo/male/female/anthro/biped) via LLM statement agreement")
    ap.add_argument("--infer-probe", action="store_true", default=True,
                    help="Infer simplified reliability-gated probe tags via LLM (default: on)")
    ap.add_argument("--no-infer-probe", dest="infer_probe", action="store_false",
                    help="Disable simplified probe inference")

    args = ap.parse_args(list(argv) if argv is not None else None)

    # Convert "none" string to actual None for disabling the filter
    min_why_val = None if args.min_why == "none" else args.min_why

    results = run_eval(
        n_samples=args.n,
        caption_field=args.caption_field,
        skip_rewrite=args.skip_rewrite,
        allow_nsfw=args.allow_nsfw,
        mode=args.mode,
        chunk_size=args.chunk_size,
        per_phrase_k=args.per_phrase_k,
        per_phrase_final_k=args.per_phrase_final_k,
        temperature=args.temperature,
        max_tokens=args.max_tokens,
        verbose=args.verbose,
        shuffle=args.shuffle,
        seed=args.seed,
        workers=args.workers,
        min_why=min_why_val,
        eval_path=args.eval_path,
        expand_implications=args.expand_implications,
        infer_structural=args.infer_structural,
        infer_probe=args.infer_probe,
    )

    print_summary(results)

    # Save results in two formats:
    # 1. Compact metrics JSONL (small, for git / LLM reading)
    # 2. Full detail JSONL (large, for analysis scripts, gitignored)
    results_dir = _REPO_ROOT / "data" / "eval_results"
    results_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_name = f"eval_{args.caption_field}_n{args.n}_seed{args.seed}_{timestamp}"

    if args.output:
        out_path = Path(args.output)
    else:
        out_path = results_dir / f"{base_name}.jsonl"
    detail_path = results_dir / f"{base_name}_detail.jsonl"

    out_path.parent.mkdir(parents=True, exist_ok=True)

    # Write run metadata as first line
    meta = {
        "_meta": True,
        "timestamp": datetime.now().isoformat(),
        "n_samples": len(results),
        "caption_field": args.caption_field,
        "skip_rewrite": args.skip_rewrite,
        "allow_nsfw": args.allow_nsfw,
        "mode": args.mode,
        "chunk_size": args.chunk_size,
        "eval_path": args.eval_path,
        "per_phrase_k": args.per_phrase_k,
        "per_phrase_final_k": args.per_phrase_final_k,
        "temperature": args.temperature,
        "shuffle": args.shuffle,
        "seed": args.seed,
        "workers": args.workers,
        "min_why": args.min_why,
        "expand_implications": args.expand_implications,
        "infer_structural": args.infer_structural,
        "infer_probe": args.infer_probe,
        "n_errors": sum(1 for r in results if r.error),
        "n_issue_samples": sum(1 for r in results if r.issues),
        "n_issues_total": sum(len(r.issues) for r in results),
    }

    with out_path.open("w", encoding="utf-8") as f:
        f.write(json.dumps(meta, ensure_ascii=False) + "\n")
        for r in results:
            # Compact format: metrics + counts + small diff sets (not full tag lists)
            missed_tags = sorted(r.ground_truth_tags - r.selected_tags)
            extra_tags = sorted(r.selected_tags - r.ground_truth_tags)
            row = {
                "id": r.sample_id,
                # Counts (not full lists)
                "n_gt": len(r.ground_truth_tags),
                "n_retrieved": len(r.retrieved_tags),
                "n_selected": len(r.selected_tags),
                "n_implied": len(r.implied_tags),
                "n_structural": len(r.structural_tags),
                "n_probe": len(r.probe_tags),
                # Overall metrics
                "ret_R": round(r.retrieval_recall, 4),
                "P": round(r.selection_precision, 4),
                "R": round(r.selection_recall, 4),
                "F1": round(r.selection_f1, 4),
                # Leaf metrics
                "leaf_P": round(r.leaf_precision, 4),
                "leaf_R": round(r.leaf_recall, 4),
                "leaf_F1": round(r.leaf_f1, 4),
                "n_leaf_sel": r.leaf_selected_count,
                "n_leaf_gt": r.leaf_gt_count,
                # Diagnostic
                "ret_P": round(r.retrieval_precision, 4),
                "sel_given_ret": round(r.selection_given_retrieval, 4),
                "over_sel": round(r.over_selection_ratio, 2),
                "why": r.why_counts,
                "stage3_diag": r.stage3_diag,
                # Character metrics (compact)
                "n_gt_char": len(r.gt_character_tags),
                "n_sel_char": len(r.selected_character_tags),
                "char_F1": round(r.char_f1, 4),
                # General metrics (compact)
                "gen_P": round(r.general_precision, 4),
                "gen_R": round(r.general_recall, 4),
                "gen_F1": round(r.general_f1, 4),
                # Diff sets (small — only the errors, not the full lists)
                "missed": missed_tags,
                "extra": extra_tags,
                # Full tag lists (needed for categorized evaluation)
                "ground_truth_tags": sorted(r.ground_truth_tags),
                "selected_tags": sorted(r.selected_tags),
                "stage3_selected": sorted(r.stage3_selected_tags),
                "stage3_selected_scores": r.stage3_selected_scores,
                "stage3_selected_ranks": r.stage3_selected_ranks,
                "stage3_selected_phrase_ranks": r.stage3_selected_phrase_ranks,
                # Evidence for extra tags (why did these false positives get through?)
                "extra_evidence": {t: r.tag_evidence.get(t, {}) for t in extra_tags},
                # Structural tags inferred
                "structural": r.structural_tags,
                "probe": r.probe_tags,
                # Timing
                "t1": round(r.stage1_time, 2),
                "t2": round(r.stage2_time, 2),
                "t3": round(r.stage3_time, 2),
                "t3s": round(r.stage3s_time, 2),
                "t3p": round(r.stage3p_time, 2),
                "err": r.error,
                "issues": r.issues,
            }
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(f"\nCompact results saved to: {out_path}")

    # Write full detail file (for analysis scripts)
    with detail_path.open("w", encoding="utf-8") as f:
        f.write(json.dumps(meta, ensure_ascii=False) + "\n")
        for r in results:
            row = {
                "sample_id": r.sample_id,
                "caption": r.caption,
                "ground_truth_tags": sorted(r.ground_truth_tags),
                "rewrite_phrases": r.rewrite_phrases,
                "retrieved_tags": sorted(r.retrieved_tags),
                "selected_tags": sorted(r.selected_tags),
                "stage3_selected": sorted(r.stage3_selected_tags),
                "stage3_selected_scores": r.stage3_selected_scores,
                "stage3_selected_ranks": r.stage3_selected_ranks,
                "stage3_selected_phrase_ranks": r.stage3_selected_phrase_ranks,
                "implied_tags": sorted(r.implied_tags),
                "structural_tags": r.structural_tags,
                "probe_tags": r.probe_tags,
                "categorized_suggestions": r.categorized_suggestions,
                "why_counts": r.why_counts,
                "stage3_diag": r.stage3_diag,
                "tag_evidence": r.tag_evidence,
                "gt_character_tags": sorted(r.gt_character_tags),
                "selected_character_tags": sorted(r.selected_character_tags),
                "gt_general_tags": sorted(r.gt_general_tags),
                "selected_general_tags": sorted(r.selected_general_tags),
                "error": r.error,
                "issues": r.issues,
            }
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(f"Detail results saved to: {detail_path}")

    return 0


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