#!/usr/bin/env python3
"""
benchmark_evaluation.py — Comparative evaluation suite for RL router vs baseline.

Evaluates two routing approaches on held-out evaluation sets:
  1. Baseline: backend's own LLM router (Qwen3-14B) — no search_sources passed
  2. GRPO: local RL router — explicit search_sources passed to backend

For each query the suite:
  - Routes via the assigned strategy
  - Executes the backend RAG pipeline
  - Compares predicted action to ground-truth action_label
  - Records accuracy, latency, and raw outputs

Outputs:
  - Per-strategy detailed CSV results
  - Aggregated comparison tables (CSV + LaTeX-ready)
  - Statistical summaries
"""

import os
import sys
import time
import json
import csv
import uuid
import logging
import argparse
import statistics
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Tuple, Optional

# ---------------------------------------------------------------------------
# Path / env setup
# ---------------------------------------------------------------------------
ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

REPO_ROOT = ROOT.parent
sys.path.append(str(REPO_ROOT))
sys.path.append(str(REPO_ROOT / "backend"))

import torch
import pandas as pd
import requests
from sentence_transformers import SentenceTransformer
from scipy import stats as scipy_stats

from model import RouterNetwork

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logger = logging.getLogger("benchmark_eval")
if not logger.handlers:
    logger.setLevel(logging.INFO)
    fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    ch = logging.StreamHandler()
    ch.setFormatter(fmt)
    logger.addHandler(ch)
    logs_dir = REPO_ROOT / "logs"
    try:
        os.makedirs(logs_dir, exist_ok=True)
        fh = logging.FileHandler(logs_dir / "benchmark_evaluation.log")
        fh.setFormatter(fmt)
        logger.addHandler(fh)
    except Exception:
        pass

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
NUM_ACTIONS = 3
ACTION_TO_TOOL = {0: "guidelines", 1: "web-search", 2: "course-registration"}
ACTION_TO_NAME = {0: "Guidelines", 1: "Web-Search", 2: "CBNU Database"}

BACKEND_CHAT_URL = os.getenv("BACKEND_CHAT_URL", "http://localhost:2000/api/chat")

BASELINE_LABEL = "Baseline-LLM"

TOP_GRPO_CHECKPOINTS = [
    {
        "id": "GRPO-E14",
        "label": "GRPO-E14",
        "path": ROOT / "adapters" / "router_grpo_epoch_14.pt",
        "epoch": 14,
        "train_avg_reward": 0.4262,
    },
    {
        "id": "GRPO-E12",
        "label": "GRPO-E12",
        "path": ROOT / "adapters" / "router_grpo_epoch_12.pt",
        "epoch": 12,
        "train_avg_reward": 0.4189,
    },
]

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def detect_language(text: str) -> str:
    if any("\uac00" <= char <= "\ud7a3" for char in text):
        return "ko"
    elif any("\u4e00" <= char <= "\u9fff" for char in text):
        return "zh"
    return "en"


def truncate(s: str, n: int = 500) -> str:
    if s is None:
        return ""
    s = str(s)
    return s if len(s) <= n else s[: n - 1] + "\u2026"


# ---------------------------------------------------------------------------
# Backend RAG pipeline
# ---------------------------------------------------------------------------
def simulate_rag_pipeline(query: str, search_sources: Optional[List[str]] = None, max_retries: int = 3) -> Tuple[str, str]:
    """Call the backend chat API.

    If search_sources is None, the backend's own LLM router decides which
    tool(s) to use (baseline behavior).
    If search_sources is provided, the backend is forced to use those tools
    (GRPO behavior).

    Includes automatic retry with exponential backoff on 429 responses.
    """
    url = BACKEND_CHAT_URL
    headers = {"accept": "application/json", "Content-Type": "application/json"}
    language = detect_language(query)

    payload = {
        "message": query,
        "language": language,
        "session_id": f"benchmark-{uuid.uuid4().hex}",
    }
    if search_sources is not None:
        payload["search_sources"] = search_sources

    last_exception = None
    for attempt in range(max_retries):
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=60)
            if response.status_code == 429:
                wait_time = (2 ** attempt) * 2  # 2s, 4s, 8s
                logger.warning(f"Rate limited (429). Retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
                time.sleep(wait_time)
                continue
            response.raise_for_status()
            data = response.json()
            answer = data.get("message", "")
            context_parts = []
            for src in data.get("sources", []):
                if isinstance(src, dict) and src.get("success"):
                    for res in src.get("data", {}).get("results", []):
                        if "content" in res:
                            context_parts.append(res["content"])
                        elif "snippet" in res:
                            context_parts.append(res["snippet"])
            context = "\n".join(context_parts)
            return context, answer
        except Exception as e:
            last_exception = e
            if attempt < max_retries - 1:
                wait_time = (2 ** attempt) * 2
                logger.warning(f"Request failed: {e}. Retrying in {wait_time}s...")
                time.sleep(wait_time)
            else:
                break

    logger.error(f"simulate_rag_pipeline failed after {max_retries} attempts: {last_exception}")
    return "", ""


# ---------------------------------------------------------------------------
# GRPO router inference
# ---------------------------------------------------------------------------
def load_grpo_checkpoint(path: Path, device: str = "cpu"):
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    embed_dim = embedder.get_sentence_embedding_dimension()
    model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)

    sd = torch.load(str(path), map_location="cpu")
    if isinstance(sd, dict) and "model_state_dict" in sd:
        model.load_state_dict(sd["model_state_dict"], strict=False)
    else:
        model.load_state_dict(sd, strict=False)
    model.eval()
    return model, embedder


def grpo_route_query(model: RouterNetwork, embedder: SentenceTransformer, query: str, device: str = "cpu") -> int:
    """Route query using GRPO model, return action_idx."""
    with torch.no_grad():
        emb = embedder.encode(query, convert_to_tensor=True).to(device)
        action, _log_prob, _value = model.get_action(emb)
    return int(action.item())


def action_to_search_sources(action_idx: int) -> List[str]:
    if action_idx == 0:
        return ["guidelines"]
    elif action_idx == 1:
        return ["web-search"]
    else:
        return ["course-registration", "department-info"]


# ---------------------------------------------------------------------------
# Evaluation runner
# ---------------------------------------------------------------------------
def evaluate_strategy(
    name: str,
    queries: List[str],
    ground_truths: List[int],
    route_fn,
    call_tool: bool = True,
    force_search_sources: bool = True,
    delay: float = 1.0,
) -> pd.DataFrame:
    """Run a routing strategy over all queries and return a results DataFrame.

    Args:
        force_search_sources: If True, pass explicit search_sources derived from
            the routed action_idx. If False, pass search_sources=None so the
            backend's own router decides (baseline behavior).
        delay: Delay in seconds between consecutive API calls to avoid rate limiting.
    """
    results = []
    for idx, (q, gt) in enumerate(zip(queries, ground_truths)):
        row = {"question": q, "ground_truth": gt}
        try:
            # Routing decision
            t0 = time.perf_counter()
            action_idx = route_fn(q)
            t1 = time.perf_counter()
            row["predicted_label"] = action_idx
            row["routing_time_sec"] = t1 - t0
            row["routing_error"] = ""

            # Tool execution
            context, answer = "", ""
            if call_tool:
                search_sources = action_to_search_sources(action_idx) if force_search_sources else None
                t2 = time.perf_counter()
                context, answer = simulate_rag_pipeline(q, search_sources=search_sources)
                t3 = time.perf_counter()
                row["tool_time_sec"] = t3 - t2
                row["end_to_end_time_sec"] = t3 - t0
                row["context"] = truncate(context, 1000)
                row["answer"] = truncate(answer, 1000)
            else:
                row["tool_time_sec"] = None
                row["end_to_end_time_sec"] = row["routing_time_sec"]
                row["context"] = ""
                row["answer"] = ""
        except Exception as e:
            logger.error(f"Error on query: {truncate(q, 80)} — {e}")
            row["routing_error"] = str(e)
            row["predicted_label"] = None
            row["routing_time_sec"] = None
            row["tool_time_sec"] = None
            row["end_to_end_time_sec"] = None
            row["context"] = ""
            row["answer"] = ""

        row["correct"] = 1 if row.get("predicted_label") == gt else 0
        results.append(row)

        # Delay between queries to avoid rate limiting
        if delay > 0 and idx < len(queries) - 1:
            time.sleep(delay)

    return pd.DataFrame(results)


# ---------------------------------------------------------------------------
# Metrics computation
# ---------------------------------------------------------------------------
def compute_metrics(df: pd.DataFrame) -> Dict:
    """Compute accuracy, latency stats from a results DataFrame."""
    valid = df.dropna(subset=["predicted_label", "ground_truth"])
    accuracy = float((valid["predicted_label"] == valid["ground_truth"]).mean()) if len(valid) else 0.0

    e2e_valid = df.dropna(subset=["end_to_end_time_sec"])
    e2e_mean = float(e2e_valid["end_to_end_time_sec"].mean()) if len(e2e_valid) else 0.0
    e2e_std = float(e2e_valid["end_to_end_time_sec"].std()) if len(e2e_valid) > 1 else 0.0

    routing_valid = df.dropna(subset=["routing_time_sec"])
    routing_mean = float(routing_valid["routing_time_sec"].mean()) if len(routing_valid) else 0.0
    routing_std = float(routing_valid["routing_time_sec"].std()) if len(routing_valid) > 1 else 0.0

    tool_valid = df.dropna(subset=["tool_time_sec"])
    tool_mean = float(tool_valid["tool_time_sec"].mean()) if len(tool_valid) else 0.0
    tool_std = float(tool_valid["tool_time_sec"].std()) if len(tool_valid) > 1 else 0.0

    # Per-class precision / recall / F1
    classes = sorted(valid["ground_truth"].unique().tolist())
    per_class = {}
    for cls in classes:
        tp = int(((valid["predicted_label"] == cls) & (valid["ground_truth"] == cls)).sum())
        fp = int(((valid["predicted_label"] == cls) & (valid["ground_truth"] != cls)).sum())
        fn = int(((valid["predicted_label"] != cls) & (valid["ground_truth"] == cls)).sum())
        prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0.0
        per_class[int(cls)] = {"precision": prec, "recall": rec, "f1": f1, "support": int((valid["ground_truth"] == cls).sum())}

    # Confusion matrix
    conf_mat = {}
    for gt_cls in classes:
        for pred_cls in classes:
            conf_mat[(int(gt_cls), int(pred_cls))] = int(
                ((valid["ground_truth"] == gt_cls) & (valid["predicted_label"] == pred_cls)).sum()
            )

    # Action distribution
    action_dist = valid["predicted_label"].value_counts().to_dict()
    action_dist = {int(k): int(v) for k, v in action_dist.items()}

    return {
        "num_queries": int(len(df)),
        "num_valid": int(len(valid)),
        "num_errors": int(len(df) - len(valid)),
        "accuracy": accuracy,
        "e2e_latency_mean": e2e_mean,
        "e2e_latency_std": e2e_std,
        "routing_latency_mean": routing_mean,
        "routing_latency_std": routing_std,
        "tool_latency_mean": tool_mean,
        "tool_latency_std": tool_std,
        "per_class": per_class,
        "confusion_matrix": conf_mat,
        "action_distribution": action_dist,
    }


# ---------------------------------------------------------------------------
# Statistical tests
# ---------------------------------------------------------------------------
def paired_t_test_latency(baseline_latencies: List[float], treatment_latencies: List[float]) -> Dict:
    """Paired t-test between baseline and treatment latency vectors."""
    if len(baseline_latencies) != len(treatment_latencies) or len(baseline_latencies) < 2:
        return {"t_stat": None, "p_value": None, "significant": None}

    t_stat, p_value = scipy_stats.ttest_rel(treatment_latencies, baseline_latencies)
    return {
        "t_stat": float(t_stat),
        "p_value": float(p_value),
        "significant": p_value < 0.05,
    }


def mean_confidence_interval(data: List[float], confidence: float = 0.95) -> Tuple[float, float, float]:
    """Return (mean, lower, upper) for CI."""
    if not data:
        return 0.0, 0.0, 0.0
    n = len(data)
    mean = statistics.mean(data)
    if n < 2:
        return mean, mean, mean
    se = statistics.stdev(data) / (n ** 0.5)
    h = se * scipy_stats.t.ppf((1 + confidence) / 2.0, n - 1)
    return mean, mean - h, mean + h


# ---------------------------------------------------------------------------
# Table / report generation
# ---------------------------------------------------------------------------
def build_comparison_table(all_metrics: Dict[str, Dict], split_name: str) -> str:
    """Build a Markdown comparison table for a dataset split."""
    rows = []
    rows.append(f"### {split_name} Evaluation Results\n")
    rows.append("| Metric | Baseline-LLM | GRPO-E14 | GRPO-E12 |")
    rows.append("|--------|-------------|---------|---------|")

    rows.append(
        f"| Routing Accuracy | {all_metrics['Baseline-LLM']['accuracy']:.2%} | "
        f"{all_metrics['GRPO-E14']['accuracy']:.2%} | {all_metrics['GRPO-E12']['accuracy']:.2%} |"
    )
    rows.append(
        f"| E2E Latency (mean ± std, s) | {all_metrics['Baseline-LLM']['e2e_latency_mean']:.4f} ± {all_metrics['Baseline-LLM']['e2e_latency_std']:.4f} | "
        f"{all_metrics['GRPO-E14']['e2e_latency_mean']:.4f} ± {all_metrics['GRPO-E14']['e2e_latency_std']:.4f} | "
        f"{all_metrics['GRPO-E12']['e2e_latency_mean']:.4f} ± {all_metrics['GRPO-E12']['e2e_latency_std']:.4f} |"
    )
    rows.append(
        f"| Routing Latency (mean ± std, s) | {all_metrics['Baseline-LLM']['routing_latency_mean']:.4f} ± {all_metrics['Baseline-LLM']['routing_latency_std']:.4f} | "
        f"{all_metrics['GRPO-E14']['routing_latency_mean']:.4f} ± {all_metrics['GRPO-E14']['routing_latency_std']:.4f} | "
        f"{all_metrics['GRPO-E12']['routing_latency_mean']:.4f} ± {all_metrics['GRPO-E12']['routing_latency_std']:.4f} |"
    )
    rows.append(
        f"| Tool Latency (mean ± std, s) | {all_metrics['Baseline-LLM']['tool_latency_mean']:.4f} ± {all_metrics['Baseline-LLM']['tool_latency_std']:.4f} | "
        f"{all_metrics['GRPO-E14']['tool_latency_mean']:.4f} ± {all_metrics['GRPO-E14']['tool_latency_std']:.4f} | "
        f"{all_metrics['GRPO-E12']['tool_latency_mean']:.4f} ± {all_metrics['GRPO-E12']['tool_latency_std']:.4f} |"
    )
    rows.append("")
    return "\n".join(rows)


def build_latex_table(all_metrics: Dict[str, Dict], split_name: str) -> str:
    """Build a LaTeX table suitable for direct inclusion in a journal paper."""
    def bold_if_best(vals: List[float], idx: int, fmt_fn, higher_is_better: bool = True) -> str:
        best_idx = 0
        best_val = vals[0]
        for i, v in enumerate(vals[1:], 1):
            if (higher_is_better and v > best_val) or (not higher_is_better and v < best_val):
                best_val = v
                best_idx = i
        if idx == best_idx:
            return "\\textbf{" + fmt_fn(vals[idx]) + "}"
        return fmt_fn(vals[idx])

    b = all_metrics["Baseline-LLM"]
    e14 = all_metrics["GRPO-E14"]
    e12 = all_metrics["GRPO-E12"]

    accs = [b["accuracy"], e14["accuracy"], e12["accuracy"]]
    e2es = [b["e2e_latency_mean"], e14["e2e_latency_mean"], e12["e2e_latency_mean"]]
    routings = [b["routing_latency_mean"], e14["routing_latency_mean"], e12["routing_latency_mean"]]
    tools = [b["tool_latency_mean"], e14["tool_latency_mean"], e12["tool_latency_mean"]]

    rows = []
    rows.append("% Auto-generated by benchmark_evaluation.py")
    rows.append("\\begin{table}[ht]")
    rows.append("\\centering")
    rows.append(f"\\caption{{Performance comparison on {split_name} evaluation set. "
                f"Results reported as mean $\\pm$ standard deviation. "
                f"Bold indicates best result.}}")
    rows.append("\\label{tab:comparison_" + split_name.lower() + "}")
    rows.append("\\begin{tabular}{lccc}")
    rows.append("\\toprule")
    rows.append("Metric & Baseline-LLM & GRPO-E14 & GRPO-E12 \\\\")
    rows.append("\\midrule")

    def acc_fmt(v): return f"{v:.2%}"
    def lat_fmt(v): return f"{v:.3f}"

    rows.append(
        f"Routing Accuracy & {bold_if_best(accs, 0, acc_fmt, higher_is_better=True)} & "
        f"{bold_if_best(accs, 1, acc_fmt, higher_is_better=True)} & "
        f"{bold_if_best(accs, 2, acc_fmt, higher_is_better=True)} \\\\"
    )
    rows.append(
        f"E2E Latency (s) & {bold_if_best(e2es, 0, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(b['e2e_latency_std'])} & "
        f"{bold_if_best(e2es, 1, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e14['e2e_latency_std'])} & "
        f"{bold_if_best(e2es, 2, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e12['e2e_latency_std'])} \\\\"
    )
    rows.append(
        f"Routing Latency (s) & {bold_if_best(routings, 0, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(b['routing_latency_std'])} & "
        f"{bold_if_best(routings, 1, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e14['routing_latency_std'])} & "
        f"{bold_if_best(routings, 2, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e12['routing_latency_std'])} \\\\"
    )
    rows.append(
        f"Tool Latency (s) & {bold_if_best(tools, 0, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(b['tool_latency_std'])} & "
        f"{bold_if_best(tools, 1, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e14['tool_latency_std'])} & "
        f"{bold_if_best(tools, 2, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e12['tool_latency_std'])} \\\\"
    )

    rows.append("\\bottomrule")
    rows.append("\\end{tabular}")
    rows.append("\\end{table}")
    rows.append("")
    return "\n".join(rows)


def build_confusion_table(conf_mat: Dict, split_name: str) -> str:
    """Build a confusion matrix table in Markdown."""
    classes = sorted({k[0] for k in conf_mat.keys()})
    rows = []
    rows.append(f"#### Confusion Matrix — {split_name}\n")
    header = "| Pred \\ GT | " + " | ".join(f"Class {c}" for c in classes) + " |"
    rows.append(header)
    sep = "|" + "|".join(["---"] * (len(classes) + 1)) + "|"
    rows.append(sep)
    for gt in classes:
        vals = [str(conf_mat.get((gt, pred), 0)) for pred in classes]
        rows.append(f"| Class {gt} | " + " | ".join(vals) + " |")
    rows.append("")
    return "\n".join(rows)


def build_per_class_table(per_class: Dict, split_name: str) -> str:
    """Build per-class precision/recall/F1 table."""
    rows = []
    rows.append(f"#### Per-Class Metrics — {split_name}\n")
    rows.append("| Class | Support | Precision | Recall | F1 |")
    rows.append("|-------|---------|-----------|--------|----|")
    for cls_id in sorted(per_class.keys()):
        m = per_class[cls_id]
        rows.append(
            f"| Class {cls_id} | {m['support']} | {m['precision']:.4f} | {m['recall']:.4f} | {m['f1']:.4f} |"
        )
    rows.append("")
    return "\n".join(rows)


def build_action_distribution_table(action_dist: Dict, split_name: str) -> str:
    """Build action distribution table."""
    rows = []
    rows.append(f"#### Action Distribution — {split_name}\n")
    rows.append("| Action | Count | Percentage |")
    rows.append("|--------|-------|------------|")
    total = sum(action_dist.values()) if action_dist else 0
    for action_id in sorted(action_dist.keys()):
        count = action_dist[action_id]
        pct = count / total if total > 0 else 0.0
        rows.append(f"| {ACTION_TO_NAME.get(action_id, f'Action {action_id}')} | {count} | {pct:.2%} |")
    rows.append("")
    return "\n".join(rows)


def _stringify_keys(obj):
    """Recursively convert dict keys to strings for JSON serialization."""
    if isinstance(obj, dict):
        return {str(k): _stringify_keys(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_stringify_keys(item) for item in obj]
    return obj


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Benchmark evaluation for RL router")
    parser.add_argument("--splits", nargs="+", default=["eng", "ko"], choices=["eng", "ko", "both"],
                        help="Which evaluation splits to run")
    parser.add_argument("--output-dir", type=str, default=str(ROOT / "benchmark_results"),
                        help="Directory for outputs")
    parser.add_argument("--call-tool", action="store_true", default=True,
                        help="Whether to call RAG tools (required for latency/accuracy measurement)")
    parser.add_argument("--device", type=str, default="cpu",
                        help="Device for GRPO model inference")
    parser.add_argument("--max-queries", type=int, default=None,
                        help="Limit queries per split (for smoke testing)")
    parser.add_argument("--delay", type=float, default=1.0,
                        help="Delay in seconds between API calls to avoid rate limiting")
    parser.add_argument("--backoff-on-429", action="store_true", default=True,
                        help="Enable exponential backoff when receiving 429 responses")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    os.makedirs(output_dir, exist_ok=True)

    split_files = {
        "eng": REPO_ROOT / "routing_agent" / "grpo_benchmark_en.csv",
        "ko": REPO_ROOT / "routing_agent" / "grpo_benchmark_ko.csv",
    }

    splits_to_run = args.splits if "both" not in args.splits else ["eng", "ko"]
    all_split_metrics = {}
    all_detailed_dfs = {}

    strategies = [BASELINE_LABEL] + [ckpt["label"] for ckpt in TOP_GRPO_CHECKPOINTS]

    for split in splits_to_run:
        csv_path = split_files[split]
        if not csv_path.exists():
            logger.error(f"Evaluation file not found: {csv_path}")
            continue

        df_eval = pd.read_csv(csv_path)
        queries = df_eval["question"].tolist()
        ground_truths = df_eval["action_label"].astype(int).tolist()

        if args.max_queries:
            queries = queries[: args.max_queries]
            ground_truths = ground_truths[: args.max_queries]

        logger.info(f"=== Evaluating {split} split: {len(queries)} queries ===")

        # Pre-load GRPO models
        grpo_models = {}
        for ckpt in TOP_GRPO_CHECKPOINTS:
            logger.info(f"Loading GRPO checkpoint: {ckpt['label']} from {ckpt['path']}")
            model, embedder = load_grpo_checkpoint(ckpt["path"], device=args.device)
            grpo_models[ckpt["label"]] = (model, embedder)

        split_results = {}

        # ---- Baseline: backend LLM router, no search_sources ----
        logger.info(f"Running {BASELINE_LABEL} on {split}...")
        def baseline_route(q):
            # For baseline, we don't need local routing — backend decides.
            # Return a dummy action; the actual routing happens inside backend.
            return 0

        df_base = evaluate_strategy(
            name=BASELINE_LABEL,
            queries=queries,
            ground_truths=ground_truths,
            route_fn=baseline_route,
            call_tool=args.call_tool,
            force_search_sources=False,
            delay=args.delay,
        )
        split_results[BASELINE_LABEL] = df_base

        # ---- GRPO checkpoints ----
        for ckpt in TOP_GRPO_CHECKPOINTS:
            label = ckpt["label"]
            model, embedder = grpo_models[label]
            logger.info(f"Running {label} on {split}...")

            def make_route_fn(m, e):
                def fn(q):
                    return grpo_route_query(m, e, q, device=args.device)
                return fn

            df_grpo = evaluate_strategy(
                name=label,
                queries=queries,
                ground_truths=ground_truths,
                route_fn=make_route_fn(model, embedder),
                call_tool=args.call_tool,
                force_search_sources=True,
                delay=args.delay,
            )
            split_results[label] = df_grpo

        # Save detailed CSV per strategy
        for strategy, df_res in split_results.items():
            out_csv = output_dir / f"{split}_{strategy}_detailed.csv"
            df_res.to_csv(out_csv, index=False)
            logger.info(f"Saved detailed results: {out_csv}")

        # Compute metrics
        metrics = {}
        for strategy, df_res in split_results.items():
            metrics[strategy] = compute_metrics(df_res)

        # Statistical tests (baseline vs each GRPO) on E2E latency
        base_latencies = split_results[BASELINE_LABEL]["end_to_end_time_sec"].dropna().tolist()
        stat_tests = {}
        for ckpt in TOP_GRPO_CHECKPOINTS:
            label = ckpt["label"]
            treat_latencies = split_results[label]["end_to_end_time_sec"].dropna().tolist()
            stat_tests[label] = {
                "paired_t": paired_t_test_latency(base_latencies, treat_latencies),
                "baseline_latency_ci": mean_confidence_interval(base_latencies),
                "treatment_latency_ci": mean_confidence_interval(treat_latencies),
            }

        all_split_metrics[split] = {"metrics": metrics, "stat_tests": stat_tests}
        all_detailed_dfs[split] = split_results

        # Print tables to console and save to markdown
        md_tables = []
        md_tables.append(build_comparison_table(metrics, split))
        for strategy in [BASELINE_LABEL] + [ckpt["label"] for ckpt in TOP_GRPO_CHECKPOINTS]:
            md_tables.append(build_confusion_table(metrics[strategy]["confusion_matrix"], split))
            md_tables.append(build_per_class_table(metrics[strategy]["per_class"], split))
            md_tables.append(build_action_distribution_table(metrics[strategy]["action_distribution"], split))

        md_report = "\n".join(md_tables)
        report_path = output_dir / f"{split}_report.md"
        with open(report_path, "w", encoding="utf-8") as f:
            f.write(md_report)
        logger.info(f"Saved report: {report_path}")
        print(md_report)

        # Save LaTeX tables
        latex_path = output_dir / f"{split}_tables.tex"
        with open(latex_path, "w", encoding="utf-8") as f:
            f.write(build_latex_table(metrics, split))
            f.write("\n")
        logger.info(f"Saved LaTeX tables: {latex_path}")

        # Save JSON metrics for programmatic access
        json_path = output_dir / f"{split}_metrics.json"
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(
                _stringify_keys({
                    "metrics": metrics,
                    "stat_tests": {
                        k: {
                            t_label: (
                                {tk: tv for tk, tv in t_val.items()}
                                if isinstance(t_val, dict)
                                else list(t_val)
                            )
                            for t_label, t_val in stat_tests[k].items()
                        }
                        for k, t_val in stat_tests.items()
                    },
                }),
                f,
                indent=2,
                default=str,
            )
        logger.info(f"Saved JSON metrics: {json_path}")

    # ---- Combined (both splits) aggregate ----
    if len(splits_to_run) == 2:
        logger.info("=== Building combined aggregate report ===")
        combined_results = {}
        for strategy in strategies:
            combined_results[strategy] = pd.concat(
                [all_detailed_dfs["eng"][strategy], all_detailed_dfs["ko"][strategy]], ignore_index=True
            )
            out_csv = output_dir / f"combined_{strategy}_detailed.csv"
            combined_results[strategy].to_csv(out_csv, index=False)

        combined_metrics = {s: compute_metrics(df) for s, df in combined_results.items()}

        base_latencies_combined = combined_results[BASELINE_LABEL]["end_to_end_time_sec"].dropna().tolist()
        combined_stat = {}
        for ckpt in TOP_GRPO_CHECKPOINTS:
            label = ckpt["label"]
            treat = combined_results[label]["end_to_end_time_sec"].dropna().tolist()
            combined_stat[label] = {
                "paired_t": paired_t_test_latency(base_latencies_combined, treat),
            }

        md_combined = []
        md_combined.append(build_comparison_table(combined_metrics, "Combined"))
        report_path = output_dir / "combined_report.md"
        with open(report_path, "w", encoding="utf-8") as f:
            f.write("\n".join(md_combined))
        logger.info(f"Saved combined report: {report_path}")
        print("\n".join(md_combined))

        latex_combined = output_dir / "combined_tables.tex"
        with open(latex_combined, "w", encoding="utf-8") as f:
            f.write(build_latex_table(combined_metrics, "Combined"))
        logger.info(f"Saved combined LaTeX: {latex_combined}")

    logger.info("Benchmark evaluation complete.")


if __name__ == "__main__":
    main()