eval_rag.py

"""RAG evaluation: compare retrieval quality across chunking methods.

Chunking methods (internal):
  (a) Graph-based -- chunks from chunk_builder.extract_chunks
  (b) Naive row-based -- one chunk per spreadsheet row
  (c) Single-cell -- each non-empty cell = one chunk (no context)

  Oracle methods (no GNN; chunk from whatever labels exist in the per-sheet JSON):
  graph_gold, graph_row_gold, graph_kg_gold — same queries and metrics as model runs.

  **Drive + service account:** Structure JSONs live under ``rag_eval/labeled_structure/``.
  For batch eval, run ``python eval_rag.py --pull-drive`` first (or use the app) so files
  mirror to ``data/rag_eval_labeled/``. Share the ``rag_eval`` folder with the labeler
  service account email (read-only is enough for download).

  **When this is a true "human structure" upper bound:** sheet JSON must contain real
  table-structure labels (``header``, ``value``, ``metadata``, ``table`` ids, etc.). The
  usual ``rag_eval_cache`` from raw xlsx download is often **unlabeled**, so oracle
  then does *not* represent gold structure—only "chunk without running the GNN." To
  measure oracle vs models, use labeled exports for the same ``sheet_id`` keys as your
  RAG eval held-out sheets (replace or augment cache files, or use ``data/labeled/``).

External baseline methods (via rag_baselines.py):
  pandas, BeautifulSoup, lxml, unstructured, LlamaSheets, Calamari

Evaluation dimensions:
  1. Retrieval quality: Recall@K, MRR (answer value in retrieved chunks)
  2. Source cell accuracy: Source Recall@K, Source MRR (correct cell retrieved)
  3. LLM-as-a-judge: answer generation + correctness scoring
  4. (Learned graph methods) Structure vs human labels: levels 1–4 via ``evaluate_sheet``,
     saved to ``data/rag_eval_structure_metrics.json`` (see ``rag_structure_rag_analysis.py`` for joins).
  5. (Learned graph methods) ``predict_sheet`` wall time per held-out sheet; aggregates
     (mean/p95 ms per graph node, throughput) in ``inference_timing`` on the same JSON entry.
     First sheet may include one-off CPU/GPU warmup — interpret percentiles accordingly.
     When structure metrics are recomputed, the new entry is merged with the previous one so
     ``inference_timing_per_sheet`` (and other extra keys) are not dropped unless this run
     saves a fresh per-sheet list with ``--save-structure-details``.

Usage:
    python eval_rag.py                          # run on rag_eval_dataset.json
    python eval_rag.py --queries path/to.json   # custom query file
    python eval_rag.py --methods graph,pandas   # specific methods only
"""

from __future__ import annotations

import json
import logging
import os
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Optional

import numpy as np

import config
import rag_eval_storage
from chunk_builder import extract_chunks, extract_row_chunks, chunk_to_text
from evaluate import evaluate_sheet, naive_flat_chunks
from label_utils import is_data_label

logger = logging.getLogger("eval_rag")


def _load_embed_model():
    from embed_text import _get_model
    return _get_model()


def _embed_texts(model, texts: list[str], batch_size: int = 64) -> np.ndarray:
    """Embed a list of texts using E5.  Returns (N, 768) float32 array."""
    prefixed = [f"query: {t}" if t else "" for t in texts]
    vecs = model.encode(
        prefixed, batch_size=batch_size,
        show_progress_bar=False, normalize_embeddings=True,
    )
    return np.array(vecs, dtype=np.float32)


# ═══════════════════════════════════════════════════════════════════════════
# Internal chunking methods (return chunks with positional metadata)
# ═══════════════════════════════════════════════════════════════════════════

# Sheets above this cell count use naive chunking directly (avoids OS OOM kill
# which cannot be caught by Python try/except).  Sheets below this threshold
# attempt GNN prediction with a fallback to naive on RuntimeError/MemoryError.
GNN_CELL_LIMIT = 3000


def _rag_grid_cell_counts(grid: dict) -> tuple[int, int]:
    """Return (n_json_cells, n_graph_nodes).

    ``n_graph_nodes`` matches ``predict_sheet`` nodes: cells with non-empty ``value``.
    ``n_json_cells`` is ``len(cells)`` on the grid dict.
    """
    cells = grid.get("cells") or []
    n_json = len(cells)
    n_graph = sum(1 for c in cells if (c.get("value") or "") != "")
    return n_json, n_graph


def _aggregate_predict_timing(
    per_sheet: list[dict[str, Any]],
    *,
    n_skipped_limit: int = 0,
    stats_rows: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
    """Aggregate wall-clock ``predict_sheet`` timings into JSON-serialisable metrics.

    ``per_sheet`` is the full list of attempts (for counts). ``stats_rows``, if set,
    is the subset used for mean / percentiles (e.g. exclude the first sheet for warmup).
    """
    rows_stats = stats_rows if stats_rows is not None else per_sheet
    ok_rows = [r for r in rows_stats if r.get("success")]
    n_failed = sum(1 for r in per_sheet if not r.get("skipped_limit") and not r.get("success"))

    def _ms_per_node(r: dict) -> float:
        n = max(int(r.get("n_graph_nodes") or 0), 1)
        return float(r["elapsed_s"]) * 1000.0 / n

    def _ms_per_json(r: dict) -> float:
        n = max(int(r.get("n_json_cells") or 0), 1)
        return float(r["elapsed_s"]) * 1000.0 / n

    ms_nodes = [_ms_per_node(r) for r in ok_rows]
    ms_json = [_ms_per_json(r) for r in ok_rows]
    total_s = sum(float(r["elapsed_s"]) for r in ok_rows)
    total_nodes = sum(int(r.get("n_graph_nodes") or 0) for r in ok_rows)

    def _pctl(arr: list[float], q: float) -> float | None:
        if not arr:
            return None
        return float(np.percentile(np.array(arr, dtype=np.float64), q))

    out: dict[str, Any] = {
        "n_sheets_attempted": len(per_sheet),
        "n_sheets_used_for_percentiles": len(rows_stats),
        "n_sheets_predict_ok": len(ok_rows),
        "n_sheets_predict_failed": n_failed,
        "n_sheets_skipped_gnn_cell_limit": int(n_skipped_limit),
        "predict_total_s": round(total_s, 4) if ok_rows else None,
        "total_graph_nodes": int(total_nodes) if ok_rows else None,
        "throughput_nodes_per_s": round(total_nodes / total_s, 2) if ok_rows and total_s > 0 else None,
        "mean_ms_per_graph_node": round(float(np.mean(ms_nodes)), 3) if ms_nodes else None,
        "p50_ms_per_graph_node": (round(v, 3) if (v := _pctl(ms_nodes, 50)) is not None else None),
        "p95_ms_per_graph_node": (round(v, 3) if (v := _pctl(ms_nodes, 95)) is not None else None),
        "mean_ms_per_json_cell": round(float(np.mean(ms_json)), 3) if ms_json else None,
        "p95_ms_per_json_cell": (round(v, 3) if (v := _pctl(ms_json, 95)) is not None else None),
    }
    return out


def _predict_and_label(
    grid_data: dict,
    model=None,
    *,
    sheet_id: str | None = None,
    pred_cache: dict[str, list[dict] | None] | None = None,
) -> dict:
    """Run GNN prediction on raw grid data and return a labeled copy.

    Merges predicted labels + table IDs back into the cell dicts so that
    extract_chunks() / build_knowledge_graph() can operate on them.
    Returns grid_data unchanged if prediction fails (OOM, etc.).

    If ``pred_cache`` is provided with ``sheet_id``, uses cached predictions from
    :func:`_fill_prediction_cache` (single inference per sheet per architecture).
    """
    n_cells = len(grid_data.get("cells", []))

    from predict import predict_sheet

    if pred_cache is not None and sheet_id is not None and sheet_id in pred_cache:
        predictions = pred_cache[sheet_id]
    else:
        try:
            predictions = predict_sheet(grid_data, model_override=model)
        except (RuntimeError, MemoryError) as e:
            logger.warning("GNN prediction failed on sheet with %d cells: %s", n_cells, e)
            import gc; gc.collect()
            return grid_data

    if not predictions:
        return grid_data

    pred_map = {(p["row"], p["col"]): p for p in predictions}
    labeled_cells = []
    for c in grid_data["cells"]:
        cell = dict(c)
        pred = pred_map.get((c["row"], c["col"]))
        if pred:
            cell["label"] = pred["label"]
            cell["table"] = pred.get("table", 1)
            cell["confidence"] = pred.get("confidence", 0.0)
        else:
            cell["label"] = "empty"
            cell["table"] = 0
        labeled_cells.append(cell)

    labeled = dict(grid_data)
    labeled["cells"] = labeled_cells
    return labeled


def _fill_prediction_cache(
    sheet_data: dict[str, dict],
    model,
    log: Callable[[str], None] | None = None,
    *,
    timing_exclude_first_sheet: bool = False,
) -> tuple[dict[str, list[dict] | None], dict[str, Any]]:
    """One ``predict_sheet`` per held-out sheet (under GNN_CELL_LIMIT).

    Returns ``(pred_cache, timing_bundle)`` where ``timing_bundle`` has keys
    ``aggregate`` (from :func:`_aggregate_predict_timing`) and ``per_sheet`` (list of
    per-attempt records for optional persistence / debugging).
    """
    from predict import predict_sheet

    cache: dict[str, list[dict] | None] = {}
    per_sheet: list[dict[str, Any]] = []
    n_skipped_limit = 0

    for sheet_id, data in sheet_data.items():
        n_json, n_graph = _rag_grid_cell_counts(data)
        n_cells = len(data.get("cells", []))
        if n_cells > GNN_CELL_LIMIT:
            cache[sheet_id] = None
            n_skipped_limit += 1
            continue
        t0 = time.perf_counter()
        try:
            p = predict_sheet(data, model_override=model)
            elapsed = time.perf_counter() - t0
            cache[sheet_id] = p if p else None
            success = bool(p)
            per_sheet.append({
                "sheet_id": sheet_id,
                "elapsed_s": round(elapsed, 6),
                "n_json_cells": n_json,
                "n_graph_nodes": n_graph,
                "success": success,
                "skipped_limit": False,
                "ms_per_graph_node": round(elapsed * 1000.0 / max(n_graph, 1), 4),
                "ms_per_json_cell": round(elapsed * 1000.0 / max(n_json, 1), 4),
            })
        except (RuntimeError, MemoryError) as e:
            elapsed = time.perf_counter() - t0
            if log:
                log(f"  predict_sheet failed {sheet_id[:40]}... ({n_cells} cells): {e}")
            logger.warning("predict_sheet failed on %s: %s", sheet_id, e)
            import gc; gc.collect()
            cache[sheet_id] = None
            per_sheet.append({
                "sheet_id": sheet_id,
                "elapsed_s": round(elapsed, 6),
                "n_json_cells": n_json,
                "n_graph_nodes": n_graph,
                "success": False,
                "skipped_limit": False,
                "error": str(e),
                "ms_per_graph_node": round(elapsed * 1000.0 / max(n_graph, 1), 4),
                "ms_per_json_cell": round(elapsed * 1000.0 / max(n_json, 1), 4),
            })

    stats_rows = per_sheet
    if timing_exclude_first_sheet and len(per_sheet) > 1:
        stats_rows = per_sheet[1:]
    agg = _aggregate_predict_timing(
        per_sheet, n_skipped_limit=n_skipped_limit, stats_rows=stats_rows,
    )
    bundle = {"aggregate": agg, "per_sheet": per_sheet}
    if log:
        if per_sheet:
            a = agg
            log(
                f"  predict_sheet timing: total={a.get('predict_total_s')}s "
                f"ok={a.get('n_sheets_predict_ok')}/{a.get('n_sheets_attempted')} "
                f"mean_ms/node={a.get('mean_ms_per_graph_node')} "
                f"p95_ms/node={a.get('p95_ms_per_graph_node')} "
                f"nodes/s={a.get('throughput_nodes_per_s')} "
                f"skipped_limit={a.get('n_sheets_skipped_gnn_cell_limit')}"
            )
        else:
            log(
                f"  predict_sheet timing: no predict attempts "
                f"(skipped GNN_CELL_LIMIT={n_skipped_limit} / {len(sheet_data)} sheets)"
            )

    return cache, bundle


def _aggregate_structure_sheet_results(rows: list[dict]) -> dict[str, float | None]:
    """Macro mean (equal weight per sheet) over ``evaluate_sheet`` outputs."""

    def _mean(getter: Callable[[dict], float | None]) -> float | None:
        vals = [getter(r) for r in rows]
        nums = [float(v) for v in vals if v is not None and isinstance(v, (int, float))]
        if not nums:
            return None
        return sum(nums) / len(nums)

    if not rows:
        return {}

    out: dict[str, float | None] = {}
    out["level1_macro_f1"] = _mean(lambda r: r.get("level1", {}).get("macro_f1"))
    out["level1_weighted_f1"] = _mean(lambda r: r.get("level1", {}).get("weighted_f1"))
    out["level2_macro_f1"] = _mean(lambda r: r.get("level2", {}).get("macro_f1"))
    out["level3_joint_accuracy"] = _mean(lambda r: r.get("level3", {}).get("joint_accuracy"))
    out["level3_ari"] = _mean(lambda r: r.get("level3", {}).get("ari"))
    l4k = (
        "chunk_f1", "chunk_precision", "chunk_recall", "soft_chunk_f1",
        "chunk_answerability", "soft_chunk_answerability", "soft_mean_iou",
    )
    for k in l4k:
        out[f"level4_{k}"] = _mean(lambda r, kk=k: r.get("level4", {}).get(kk))

    return out


def _compute_structure_vs_gold_for_arch(
    sheet_data: dict[str, dict],
    pred_cache: dict[str, list[dict] | None],
    *,
    save_per_sheet: bool = False,
) -> dict[str, Any]:
    """Run ``evaluate_sheet`` on sheets that have human RAG structure labels."""
    per_eval: list[dict] = []
    sheet_ids_used: list[str] = []
    skipped_no_gold: list[str] = []
    skipped_no_pred: list[str] = []

    for sheet_id, data in sheet_data.items():
        if not _grid_has_rag_structure_labels(data):
            skipped_no_gold.append(sheet_id)
            continue
        preds = pred_cache.get(sheet_id)
        if not preds:
            skipped_no_pred.append(sheet_id)
            continue
        ev = evaluate_sheet(data, preds)
        if not ev:
            skipped_no_pred.append(sheet_id)
            continue
        sheet_ids_used.append(sheet_id)
        per_eval.append(ev)

    agg = _aggregate_structure_sheet_results(per_eval) if per_eval else {}

    meta: dict[str, Any] = {
        "n_sheets_with_gold": len(sheet_ids_used),
        "sheet_ids_evaluated": sheet_ids_used,
        "skipped_no_structure_labels": skipped_no_gold,
        "skipped_no_predictions_or_empty_eval": skipped_no_pred,
        "aggregates": agg,
    }
    if save_per_sheet and per_eval:
        meta["_per_sheet"] = [
            {"sheet_id": sid, "levels": row} for sid, row in zip(sheet_ids_used, per_eval)
        ]
    return meta


def _load_rag_structure_metrics_doc() -> dict[str, Any]:
    p = config.RAG_EVAL_STRUCTURE_METRICS
    if not p.exists():
        return {}
    try:
        with open(p, encoding="utf-8") as f:
            d = json.load(f)
        return d if isinstance(d, dict) else {}
    except Exception:
        return {}


def _write_rag_structure_metrics_doc(doc: dict[str, Any]) -> None:
    path = config.RAG_EVAL_STRUCTURE_METRICS
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(doc, f, indent=2, default=str)


def _should_skip_structure_resume(
    method_key_base: str,
    resume: bool,
    existing: dict[str, Any],
    force: bool,
) -> bool:
    if force or not resume:
        return False
    entry = existing.get(method_key_base)
    if not isinstance(entry, dict):
        return False
    if entry.get("error"):
        return False
    agg = entry.get("aggregates")
    if isinstance(agg, dict) and agg.get("level4_chunk_f1") is not None:
        return True
    if isinstance(agg, dict) and entry.get("n_sheets_with_gold", 0) == 0:
        return True
    return False


def _graph_chunks(
    grid_data: dict,
    model=None,
    *,
    sheet_id: str = "",
    pred_cache: dict[str, list[dict] | None] | None = None,
) -> list[dict]:
    """Run GNN prediction → structured chunking (value + header paths)."""
    labeled = _predict_and_label(
        grid_data, model=model,
        sheet_id=sheet_id or None, pred_cache=pred_cache,
    )
    chunks = extract_chunks(labeled)
    result = []
    for c in chunks:
        text = chunk_to_text(c)
        if text:
            pos = c.get("position", {})
            result.append({
                "text": text,
                "method": "graph",
                "row": pos.get("row"),
                "col": pos.get("col"),
            })
    if not result:
        return _naive_chunks(labeled)
    return result


def _graph_row_chunks(
    grid_data: dict,
    model=None,
    *,
    sheet_id: str = "",
    pred_cache: dict[str, list[dict] | None] | None = None,
) -> list[dict]:
    """Run GNN prediction → row-based structured chunking."""
    labeled = _predict_and_label(
        grid_data, model=model,
        sheet_id=sheet_id or None, pred_cache=pred_cache,
    )
    chunks = extract_row_chunks(labeled)
    result = []
    for c in chunks:
        text = c.get("text", "")
        if text:
            pos = c.get("position", {})
            result.append({
                "text": text,
                "method": "graph_row",
                "row": pos.get("row"),
                "col": pos.get("col"),
            })
    if not result:
        return _naive_chunks(labeled)
    return result


def _graph_kg_chunks(
    grid_data: dict,
    model=None,
    *,
    sheet_id: str = "",
    pred_cache: dict[str, list[dict] | None] | None = None,
) -> list[dict]:
    """Run GNN prediction → knowledge graph → graph-walk chunking."""
    from chunk_builder import build_knowledge_graph, generate_chunks_from_graph

    labeled = _predict_and_label(
        grid_data, model=model,
        sheet_id=sheet_id or None, pred_cache=pred_cache,
    )
    sheet_name = labeled.get("sheet_name", "")
    G = build_knowledge_graph(labeled["cells"], sheet_name)
    kg_chunks = generate_chunks_from_graph(G, sheet_name)
    result = []
    for c in kg_chunks:
        if c.get("text"):
            result.append({
                "text": c["text"],
                "method": "graph_kg",
                "row": None,
                "col": None,
            })
    if not result:
        return _naive_chunks(labeled)
    return result


# ── Oracle chunking (human labels from cached JSON; no _predict_and_label) ──


def _graph_chunks_gold(grid_data: dict) -> list[dict]:
    """Structured graph chunks using labels in ``grid_data`` as-is (oracle / upper bound)."""
    chunks = extract_chunks(dict(grid_data))
    result = []
    for c in chunks:
        text = chunk_to_text(c)
        if text:
            pos = c.get("position", {})
            result.append({
                "text": text,
                "method": "graph_gold",
                "row": pos.get("row"),
                "col": pos.get("col"),
            })
    if not result:
        return _naive_chunks(dict(grid_data))
    return result


def _graph_row_chunks_gold(grid_data: dict) -> list[dict]:
    """Row chunks from human labels only (oracle)."""
    chunks = extract_row_chunks(dict(grid_data))
    result = []
    for c in chunks:
        text = c.get("text", "")
        if text:
            pos = c.get("position", {})
            result.append({
                "text": text,
                "method": "graph_row_gold",
                "row": pos.get("row"),
                "col": pos.get("col"),
            })
    if not result:
        return _naive_chunks(dict(grid_data))
    return result


def _graph_kg_chunks_gold(grid_data: dict) -> list[dict]:
    """KG walk chunks from human labels only (oracle)."""
    from chunk_builder import build_knowledge_graph, generate_chunks_from_graph

    labeled = dict(grid_data)
    sheet_name = labeled.get("sheet_name", "") or labeled.get("source", {}).get("sheet_name", "")
    G = build_knowledge_graph(labeled["cells"], sheet_name)
    kg_chunks = generate_chunks_from_graph(G, sheet_name)
    result = []
    for c in kg_chunks:
        if c.get("text"):
            result.append({
                "text": c["text"],
                "method": "graph_kg_gold",
                "row": None,
                "col": None,
            })
    if not result:
        return _naive_chunks(labeled)
    return result


def _naive_chunks(labeled_data: dict) -> list[dict]:
    cells = labeled_data["cells"]
    sheet = labeled_data.get("sheet_name", labeled_data.get("source", {}).get("sheet_name", ""))
    return naive_flat_chunks(cells, sheet)


def _single_cell_chunks(labeled_data: dict) -> list[dict]:
    cells = labeled_data["cells"]
    chunks = []
    for c in cells:
        val = c.get("value", "").strip()
        if val:
            chunks.append({"text": val, "method": "single_cell",
                           "row": c["row"], "col": c["col"]})
    return chunks


# ═══════════════════════════════════════════════════════════════════════════
# Retrieval strategies
# ═══════════════════════════════════════════════════════════════════════════

RETRIEVAL_STRATEGIES = ["vector", "bm25", "hybrid_rrf", "hybrid_rrf_rerank"]
DEFAULT_RETRIEVAL_STRATEGIES = ["vector", "bm25", "hybrid_rrf"]


def _build_faiss_index(embeddings: np.ndarray):
    """Build a FAISS inner-product index."""
    import faiss
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings)
    return index


def _retrieve_vector(index, query_vec: np.ndarray, k: int = 10) -> list[int]:
    """Return top-K indices via dense vector similarity."""
    D, I = index.search(query_vec.reshape(1, -1), k)
    return [int(idx) for idx in I[0] if idx >= 0]


def _tokenize(text: str) -> list[str]:
    """Simple whitespace + lowercasing tokenizer for BM25."""
    import re
    return re.findall(r"\w+", text.lower())


def _build_bm25_index(texts: list[str]):
    """Build a BM25 index over tokenized chunk texts."""
    from rank_bm25 import BM25Okapi
    tokenized = [_tokenize(t) for t in texts]
    return BM25Okapi(tokenized)


def _retrieve_bm25(bm25_index, query: str, k: int = 10) -> list[int]:
    """Return top-K indices via BM25 keyword matching."""
    scores = bm25_index.get_scores(_tokenize(query))
    top_k = np.argsort(scores)[::-1][:k]
    return [int(idx) for idx in top_k if scores[idx] > 0]


def _retrieve_hybrid_rrf(
    faiss_index, query_vec: np.ndarray,
    bm25_index, query_text: str,
    k: int = 10, rrf_k: int = 60,
) -> list[int]:
    """Reciprocal Rank Fusion of vector + BM25 results."""
    vec_ids = _retrieve_vector(faiss_index, query_vec, k=k * 2)
    bm25_ids = _retrieve_bm25(bm25_index, query_text, k=k * 2)

    rrf_scores: dict[int, float] = {}
    for rank, idx in enumerate(vec_ids):
        rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1)
    for rank, idx in enumerate(bm25_ids):
        rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1)

    ranked = sorted(rrf_scores, key=rrf_scores.get, reverse=True)
    return ranked[:k]


def _retrieve_hybrid_rrf_rerank(
    faiss_index,
    query_vec: np.ndarray,
    bm25_index,
    query_text: str,
    chunk_embs: np.ndarray,
    k: int = 10,
    pool_factor: int = 4,
    rrf_k: int = 60,
) -> list[int]:
    """RRF fusion for a wide candidate pool, then re-order by dense similarity."""
    pool = max(k * pool_factor, k)
    vec_ids = _retrieve_vector(faiss_index, query_vec, k=pool)
    bm25_ids = _retrieve_bm25(bm25_index, query_text, k=pool)
    rrf_scores: dict[int, float] = {}
    for rank, idx in enumerate(vec_ids):
        rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1)
    for rank, idx in enumerate(bm25_ids):
        rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1)
    candidates = sorted(rrf_scores.keys(), key=lambda i: rrf_scores[i], reverse=True)[:pool]
    if not candidates:
        return []
    q = query_vec.reshape(-1)
    scored = [(float(np.dot(chunk_embs[idx], q)), idx) for idx in candidates]
    scored.sort(reverse=True)
    return [idx for _, idx in scored[:k]]


# ═══════════════════════════════════════════════════════════════════════════
# Source cell matching
# ═══════════════════════════════════════════════════════════════════════════

def _chunk_contains_source(chunk: dict, source_cells: list[dict]) -> bool:
    """Check if a chunk covers any of the expected source cells.

    Two matching strategies (tried in order):
      1. **Value match** — if source cells carry a ``value`` field, check
         whether that value string appears anywhere in the chunk text.
         This works for *all* chunking methods, even those without positional
         metadata (pandas, unstructured, beautifulsoup …).
      2. **Position match** — fall back to row/col matching when chunk has
         positional metadata.
    """
    chunk_text = chunk.get("text", "")
    c_row = chunk.get("row")
    c_col = chunk.get("col")

    for src in source_cells:
        src_val = src.get("value", "")
        if src_val and src_val in chunk_text:
            return True
        s_row, s_col = src.get("row"), src.get("col")
        if c_row is not None and c_row == s_row:
            if c_col is None or c_col == s_col:
                return True
    return False


# ═══════════════════════════════════════════════════════════════════════════
# LLM-as-a-Judge
# ═══════════════════════════════════════════════════════════════════════════

_ANSWER_GEN_PROMPT = """You are answering questions about data in spreadsheets.
Given the following spreadsheet excerpts, answer the question.
Include the specific values and numbers from the data that support your answer.
If you cannot find the answer in the provided excerpts, say "NOT FOUND".

Spreadsheet excerpts:
{chunks}

Question: {question}

Answer:"""

_JUDGE_PROMPT = """You are evaluating whether a generated answer correctly answers the same question as an expected answer.
Both answers refer to data extracted from a spreadsheet.

Expected answer: {expected}
Generated answer: {generated}

Rate the generated answer on a scale of 1-5:
1 = Completely wrong, contradicts the expected answer, or unrelated
2 = Partially related but key facts are wrong or misleading
3 = Answers a different aspect of the question, or gives only tangentially related information
4 = Correctly answers the core question but with less detail or supporting data than the expected answer
5 = Correctly answers the question with equivalent or sufficient detail

Important: if the question is yes/no or asks "which one", and the generated answer gives the correct yes/no/choice, that is at least a 4 even if it omits supporting numbers. The core answer matters most.

Respond with ONLY a single integer (1-5)."""

JUDGE_MODELS = [
    ("vertex_google", "gemini-2.0-flash"),
    ("vertex_google", "gemini-2.0-flash-lite"),
    ("vertex_google", "gemini-2.5-flash"),
    ("openai",        "gpt-4o-mini"),
    ("openai",        "gpt-4.1-nano"),
]

_VERTEX_PROJECT = os.environ.get("GOOGLE_CLOUD_PROJECT", "acoustic-atom-386613")
_VERTEX_LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")

_openai_client = None
_vertex_google_client = None
_anthropic_client = None


def _has_vertex_credentials() -> bool:
    """Check if Google Cloud / Vertex AI credentials are available."""
    if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
        return True
    try:
        import google.auth
        google.auth.default()
        return True
    except Exception:
        return False


def _get_available_judges() -> list[tuple[str, str]]:
    """Return the subset of JUDGE_MODELS whose credentials are available."""
    has_vertex = _has_vertex_credentials()
    has_openai = bool(os.environ.get("OPENAI_API_KEY"))
    has_anthropic = bool(os.environ.get("ANTHROPIC_API_KEY"))

    available = []
    for provider, model in JUDGE_MODELS:
        if provider == "vertex_google" and has_vertex:
            available.append((provider, model))
        elif provider == "anthropic" and has_anthropic:
            available.append((provider, model))
        elif provider == "openai" and has_openai:
            available.append((provider, model))
    return available


def _call_openai(prompt: str, model: str) -> str:
    global _openai_client
    import openai
    if _openai_client is None:
        _openai_client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    resp = _openai_client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=256,
        temperature=0,
    )
    return resp.choices[0].message.content.strip()


def _call_vertex_google(prompt: str, model: str) -> str:
    global _vertex_google_client
    from google import genai
    from google.genai.types import GenerateContentConfig
    if _vertex_google_client is None:
        _vertex_google_client = genai.Client(
            vertexai=True,
            project=_VERTEX_PROJECT,
            location=_VERTEX_LOCATION,
        )
    resp = _vertex_google_client.models.generate_content(
        model=model,
        contents=prompt,
        config=GenerateContentConfig(max_output_tokens=256, temperature=0),
    )
    text = resp.text
    if text is None:
        raise RuntimeError(f"Empty response from {model}")
    return text.strip()


def _call_anthropic(prompt: str, model: str) -> str:
    global _anthropic_client
    import anthropic
    if _anthropic_client is None:
        _anthropic_client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    resp = _anthropic_client.messages.create(
        model=model,
        max_tokens=256,
        messages=[{"role": "user", "content": prompt}],
    )
    return resp.content[0].text.strip()


_CALL_DISPATCH = {
    "openai": _call_openai,
    "vertex_google": _call_vertex_google,
    "anthropic": _call_anthropic,
}


def _call_llm(prompt: str, provider: str, model: str) -> str:
    """Call a single LLM by provider and model name."""
    fn = _CALL_DISPATCH.get(provider)
    if fn is None:
        raise ValueError(f"Unknown provider: {provider}")
    return fn(prompt, model)


def _generate_answer(question: str, chunks: list[str]) -> str:
    """Generate an answer using the first available judge model."""
    chunk_text = "\n---\n".join(chunks[:5])
    prompt = _ANSWER_GEN_PROMPT.format(chunks=chunk_text, question=question)
    judges = _get_available_judges()
    if not judges:
        raise RuntimeError(
            "No LLM credentials found. Set GOOGLE_APPLICATION_CREDENTIALS "
            "(for Vertex AI), ANTHROPIC_API_KEY, or OPENAI_API_KEY."
        )
    provider, model = judges[0]
    return _call_llm(prompt, provider, model)


def _parse_judge_score(response: str) -> int:
    """Extract a 1-5 integer from a judge response."""
    try:
        score = int(response.strip())
        return max(1, min(5, score))
    except (ValueError, TypeError):
        import re
        match = re.search(r"[1-5]", response)
        return int(match.group()) if match else 1


def _judge_answer(expected: str, generated: str) -> dict:
    """Score answer correctness using all available judge models.

    Returns dict with per-model scores and the averaged score.
    """
    if not generated or generated == "NOT FOUND":
        judges = _get_available_judges()
        per_model = {f"{p}/{m}": 1 for p, m in judges} if judges else {}
        return {"score": 1, "per_model": per_model}

    prompt = _JUDGE_PROMPT.format(expected=expected, generated=generated)
    judges = _get_available_judges()
    if not judges:
        raise RuntimeError("No LLM credentials found.")

    scores: dict[str, int] = {}
    for provider, model in judges:
        key = f"{provider}/{model}"
        try:
            response = _call_llm(prompt, provider, model)
            scores[key] = _parse_judge_score(response)
        except Exception as e:
            logger.warning("Judge %s failed: %s", key, e)

    if not scores:
        return {"score": 1, "per_model": {}}

    avg = sum(scores.values()) / len(scores)
    return {"score": round(avg, 2), "per_model": scores}


# ═══════════════════════════════════════════════════════════════════════════
# Unified evaluation
# ═══════════════════════════════════════════════════════════════════════════

def evaluate_method_v2(
    chunks: list[dict],
    queries: list[dict],
    embed_model,
    retrieval: str = "vector",
    ks: tuple[int, ...] = (1, 5, 10),
    use_llm_judge: bool = False,
    save_details: bool = False,
    _prebuilt_faiss=None,
    _prebuilt_bm25=None,
    _prebuilt_chunk_embs=None,
) -> dict:
    """Evaluate retrieval quality with source accuracy and optional LLM judge.

    Args:
        retrieval: one of "vector", "bm25", "hybrid_rrf", "hybrid_rrf_rerank"
        save_details: if True, ``result["_details"]`` will contain per-query
            data (retrieved chunks, generated answer, per-judge scores).
        _prebuilt_*: optional pre-computed indices to avoid rebuilding across
                     retrieval strategies for the same set of chunks.
    """
    texts = [c.get("text", "") for c in chunks]
    if not texts:
        result = {f"recall@{k}": 0.0 for k in ks}
        result.update({f"source_recall@{k}": 0.0 for k in ks})
        result.update({"mrr": 0.0, "source_mrr": 0.0, "n_chunks": 0})
        if use_llm_judge:
            result.update({"judge_score": 0.0, "judge_binary_acc": 0.0})
        return result

    needs_vector = retrieval in ("vector", "hybrid_rrf", "hybrid_rrf_rerank")
    needs_bm25 = retrieval in ("bm25", "hybrid_rrf", "hybrid_rrf_rerank")

    chunk_embs = _prebuilt_chunk_embs
    faiss_index = _prebuilt_faiss
    bm25_index = _prebuilt_bm25

    if needs_vector and chunk_embs is None:
        chunk_embs = _embed_texts(embed_model, texts)
    if needs_vector and faiss_index is None:
        faiss_index = _build_faiss_index(chunk_embs)
    if needs_bm25 and bm25_index is None:
        bm25_index = _build_bm25_index(texts)

    q_texts = [q["question"] for q in queries]
    q_embs = None
    if needs_vector:
        q_embs = _embed_texts(embed_model, q_texts)

    recalls = {k: 0 for k in ks}
    source_recalls = {k: 0 for k in ks}
    rr_sum = 0.0
    source_rr_sum = 0.0
    judge_scores = []
    per_query_details: list[dict] = []

    n_total = len(queries)
    for qi, q in enumerate(queries):
        expected_answer = q.get("expected_answer", q.get("expected_value", ""))
        source_cells = q.get("source_cells", [])

        if retrieval == "vector":
            valid_indices = _retrieve_vector(faiss_index, q_embs[qi], k=max(ks))
        elif retrieval == "bm25":
            valid_indices = _retrieve_bm25(bm25_index, q["question"], k=max(ks))
        elif retrieval == "hybrid_rrf":
            valid_indices = _retrieve_hybrid_rrf(
                faiss_index, q_embs[qi], bm25_index, q["question"], k=max(ks),
            )
        elif retrieval == "hybrid_rrf_rerank":
            if chunk_embs is None:
                raise ValueError("hybrid_rrf_rerank requires chunk embeddings")
            valid_indices = _retrieve_hybrid_rrf_rerank(
                faiss_index, q_embs[qi], bm25_index, q["question"], chunk_embs, k=max(ks),
            )
        else:
            raise ValueError(f"Unknown retrieval strategy: {retrieval}")

        retrieved_texts = [texts[idx] for idx in valid_indices]
        retrieved_chunks = [chunks[idx] for idx in valid_indices]

        first_answer_hit = None
        for rank, text in enumerate(retrieved_texts, start=1):
            if expected_answer and expected_answer in text:
                if first_answer_hit is None:
                    first_answer_hit = rank
                break
        for k in ks:
            if any(expected_answer and expected_answer in t for t in retrieved_texts[:k]):
                recalls[k] += 1
        if first_answer_hit is not None:
            rr_sum += 1.0 / first_answer_hit

        source_hit_rank = None
        if source_cells:
            first_source_hit = None
            for rank, chunk in enumerate(retrieved_chunks, start=1):
                if _chunk_contains_source(chunk, source_cells):
                    if first_source_hit is None:
                        first_source_hit = rank
                    break
            source_hit_rank = first_source_hit
            for k in ks:
                if any(_chunk_contains_source(c, source_cells) for c in retrieved_chunks[:k]):
                    source_recalls[k] += 1
            if first_source_hit is not None:
                source_rr_sum += 1.0 / first_source_hit

        generated_answer = None
        verdict = None
        if use_llm_judge:
            try:
                generated_answer = _generate_answer(q["question"], retrieved_texts[:5])
                verdict = _judge_answer(expected_answer, generated_answer)
                judge_scores.append(verdict)
                if (qi + 1) % 10 == 0 or qi == n_total - 1:
                    running_avg = sum(v["score"] for v in judge_scores) / len(judge_scores)
                    print(f"    [judge] {qi+1}/{n_total} queries scored  (running avg: {running_avg:.2f})")
            except Exception as e:
                logger.warning("LLM judge failed for query %d: %s", qi, e)
                judge_scores.append({"score": 1, "per_model": {}})

        if save_details:
            detail = {
                "query_id": q.get("id", f"q_{qi}"),
                "question": q["question"],
                "expected_answer": expected_answer,
                "sheet_id": q.get("sheet_id", ""),
                "query_difficulty": q.get("query_difficulty"),
                "sheet_complexity": q.get("sheet_complexity"),
                "retrieved_chunks": [
                    {"rank": r + 1, "text": retrieved_texts[r],
                     "row": retrieved_chunks[r].get("row"),
                     "col": retrieved_chunks[r].get("col"),
                     "contains_answer": bool(expected_answer and expected_answer in retrieved_texts[r]),
                     "contains_source": _chunk_contains_source(retrieved_chunks[r], source_cells) if source_cells else None}
                    for r in range(len(retrieved_texts))
                ],
                "answer_hit_rank": first_answer_hit,
                "source_hit_rank": source_hit_rank,
                "generated_answer": generated_answer,
                "judge_verdict": verdict,
            }
            per_query_details.append(detail)

    n_q = max(len(queries), 1)
    result = {f"recall@{k}": recalls[k] / n_q for k in ks}
    result["mrr"] = rr_sum / n_q
    result.update({f"source_recall@{k}": source_recalls[k] / n_q for k in ks})
    result["source_mrr"] = source_rr_sum / n_q
    result["n_chunks"] = len(texts)

    if use_llm_judge and judge_scores:
        avg_scores = [v["score"] for v in judge_scores]
        result["judge_score"] = sum(avg_scores) / len(avg_scores)
        result["judge_binary_acc"] = sum(1 for s in avg_scores if s >= 4) / len(avg_scores)

        all_model_keys: set[str] = set()
        for v in judge_scores:
            all_model_keys.update(v.get("per_model", {}).keys())
        for mk in sorted(all_model_keys):
            model_scores = [v["per_model"][mk] for v in judge_scores if mk in v.get("per_model", {})]
            if model_scores:
                result[f"judge_{mk}"] = round(sum(model_scores) / len(model_scores), 3)

    if save_details:
        result["_details"] = per_query_details

    return result


# ═══════════════════════════════════════════════════════════════════════════
# V2 runner (used by the API)
# ═══════════════════════════════════════════════════════════════════════════

def _resolve_xlsx_path(query: dict, sheet_id: str) -> str | None:
    """Find the xlsx file path for a query's sheet.

    Checks (in order): query's spreadsheet_file, temp downloads, raw dir.
    """
    if query.get("spreadsheet_file") and Path(query["spreadsheet_file"]).exists():
        return query["spreadsheet_file"]

    drive_fid = query.get("drive_file_id", "")
    if drive_fid:
        tmp_path = config.TEMP_DIR / f"{drive_fid}.xlsx"
        if tmp_path.exists():
            return str(tmp_path)
        try:
            import drive_client
            drive_client.download_file(drive_fid, str(tmp_path))
            return str(tmp_path)
        except Exception:
            pass

    raw_dir = config.DATA_DIR / "raw"
    if raw_dir.exists():
        for sub in raw_dir.iterdir():
            if sub.is_dir():
                for xlsx in sub.glob("*.xlsx"):
                    sid = drive_fid or sheet_id.rsplit("_", 1)[0]
                    if sid and sid in xlsx.stem:
                        return str(xlsx)
    return None


def _adapter_sheet_key(xlsx_path: str, sheet_name: str) -> tuple[str, str]:
    """Canonical key so the same workbook+sheet is not chunked twice."""
    try:
        p = str(Path(xlsx_path).resolve())
    except OSError:
        p = os.path.normpath(xlsx_path)
    return (p, sheet_name)


def _load_sheet_data(sheet_id: str) -> dict | None:
    """Load grid data for RAG eval (RAG structure labels from Drive/local, else raw cache)."""
    return rag_eval_storage.load_rag_eval_grid_for_sheet(sheet_id)


def _grid_has_rag_structure_labels(grid: dict | None) -> bool:
    """True if the loaded grid JSON has at least one non-trivial cell label (oracle / human structure)."""
    if not grid or not isinstance(grid, dict):
        return False
    for c in grid.get("cells") or []:
        lb = (c.get("label") or "").strip().lower()
        if lb and lb not in ("empty", "unlabeled"):
            return True
    return False


def _load_rag_existing_aggregates() -> dict[str, Any]:
    """Load ``rag_eval_results.json`` if present (for resume / merge)."""
    path = config.DATA_DIR / "rag_eval_results.json"
    if not path.exists():
        return {}
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {}


def _should_skip_rag_resume(
    key: str,
    resume: bool,
    existing_agg: dict[str, Any],
    log_fn: Callable[[str], None] | None,
    *,
    need_judge: bool = False,
) -> bool:
    """Skip re-eval if this key already has a successful checkpoint (no ``error``).

    If ``need_judge`` is True (LLM-as-a-judge requested), only skip when retrieval
    metrics *and* judge fields are present — otherwise a prior ``--no-judge`` run
    would block adding judge scores.
    """
    if not resume:
        return False
    val = existing_agg.get(key)
    if not isinstance(val, dict):
        return False
    if val.get("error"):
        return False
    has_retrieval = any(k in val for k in ("recall@1", "mrr"))
    has_judge = any(k in val for k in ("judge_score", "judge_binary_acc"))
    if need_judge:
        if has_retrieval and has_judge:
            if log_fn:
                log_fn(f"  SKIP {key} (resume — already in rag_eval_results.json)")
            return True
        return False
    if has_retrieval or has_judge:
        if log_fn:
            log_fn(f"  SKIP {key} (resume — already in rag_eval_results.json)")
        return True
    return False


def _write_rag_results_batch(
    batch: dict[str, dict],
    *,
    save_details: bool,
    log_fn: Callable[[str], None] | None,
    existing_agg: dict[str, Any],
) -> None:
    """Merge ``batch`` into ``rag_eval_results.json`` on disk; optionally ``rag_eval_details.json``.

    Call after each method/strategy (or error batch) so a crash does not lose progress.
    """
    if not batch:
        return
    path = config.DATA_DIR / "rag_eval_results.json"
    existing: dict[str, Any] = {}
    if path.exists():
        try:
            with open(path, encoding="utf-8") as f:
                existing = json.load(f)
        except Exception:
            existing = {}
    details_delta: dict[str, Any] = {}
    for key, val in batch.items():
        if isinstance(val, dict) and "_details" in val:
            details_delta[key] = val["_details"]
            val = {k: v for k, v in val.items() if k != "_details"}
        existing[key] = val
        existing_agg[key] = val
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(existing, f, indent=2, default=str)
    if log_fn:
        log_fn(f"  [checkpoint] saved {len(batch)} key(s) → {path.name}")
    if save_details and details_delta:
        dpath = config.DATA_DIR / "rag_eval_details.json"
        ex_d: dict[str, Any] = {}
        if dpath.exists():
            try:
                with open(dpath, encoding="utf-8") as f:
                    ex_d = json.load(f)
            except Exception:
                ex_d = {}
        ex_d.update(details_delta)
        with open(dpath, "w", encoding="utf-8") as f:
            json.dump(ex_d, f, indent=2, default=str)
        if log_fn:
            log_fn(f"  [checkpoint] details {len(details_delta)} key(s) → {dpath.name}")


def run_rag_evaluation_v2(
    methods: list[str] | None = None,
    use_llm_judge: bool = True,
    save_details: bool = False,
    retrieval_strategies: list[str] | None = None,
    log_fn: Callable[[str], None] | None = None,
    architectures: list[str] | None = None,
    resume: bool = True,
    *,
    resume_structure: bool = True,
    force_structure_metrics: bool = False,
    save_structure_details: bool = False,
    timing_exclude_first_sheet: bool = False,
) -> dict:
    """Run full RAG evaluation using the rag_eval_dataset.json.

    Evaluates internal chunking methods (graph, naive, single_cell) and
    external baselines from rag_baselines.py.

    Sheet grids are loaded from ``rag_eval_cache`` then ``labeled`` (see
    ``_load_sheet_data``). Training data and RAG-eval data are separate concerns:
    RAG eval provides Q&A; chunking for ``graph_*`` uses the GNN on those grids.
    Oracle ``*_gold`` methods skip the GNN and use labels in the loaded JSON only,
    which is a meaningful ceiling **only** if those files actually contain human
    structure labels for the eval sheets.

    Args:
        retrieval_strategies: subset of RETRIEVAL_STRATEGIES to run (vector, bm25,
            hybrid_rrf, hybrid_rrf_rerank). Defaults to all except rerank.
        save_details: if True, merge per-query retrieval rows into ``rag_eval_details.json``.
        resume: if True (default), skip method/strategy keys already present in
            ``rag_eval_results.json`` with metrics (no ``error``), so you can restart after a crash.
        resume_structure: skip rewriting ``rag_eval_structure_metrics.json`` when an entry exists.
        force_structure_metrics: recompute structure-vs-gold even if a checkpoint exists.
        save_structure_details: include per-sheet ``_per_sheet`` in the structure metrics JSON.
        timing_exclude_first_sheet: if True, mean / percentiles for ``inference_timing`` exclude
            the first attempted ``predict_sheet`` call (reduces one-off warmup skew); counts
            still reflect all sheets.
    """
    def log(msg: str):
        print(msg)
        if log_fn:
            log_fn(msg)

    strategies = retrieval_strategies or list(DEFAULT_RETRIEVAL_STRATEGIES)

    existing_agg = _load_rag_existing_aggregates()
    if resume and existing_agg:
        log(
            f"Resume: {len(existing_agg)} key(s) on disk — skipping rows that already "
            f"have metrics. Pass resume=False to recompute everything."
        )

    dataset = rag_eval_storage.load_rag_eval_dataset()
    if not dataset.get("queries"):
        log("No RAG eval queries found. Add Q&As at /rag-eval or upload rag_eval_dataset.json to the Drive RAG folder.")
        return {}

    queries = dataset.get("queries", [])
    held_out = set(dataset.get("held_out_sheets", []))

    log(f"Loaded {len(queries)} queries across {len(held_out)} held-out sheets")

    if methods is None:
        methods = [
            "graph", "graph_kg", "naive", "single_cell",
            "pandas", "beautifulsoup", "lxml",
        ]

    available_judges = _get_available_judges()
    if use_llm_judge and not available_judges:
        log("No LLM credentials found — set GOOGLE_APPLICATION_CREDENTIALS, ANTHROPIC_API_KEY, or OPENAI_API_KEY — skipping LLM judge")
        use_llm_judge = False
    elif use_llm_judge:
        judge_names = [f"{p}/{m}" for p, m in available_judges]
        log(f"LLM judges ({len(available_judges)}): {', '.join(judge_names)}")

    sheet_data: dict[str, dict] = {}
    for sheet_id in held_out:
        data = _load_sheet_data(sheet_id)
        if data:
            sheet_data[sheet_id] = data
    log(f"Loaded {len(sheet_data)} held-out sheet(s) with data")

    structure_metrics_doc = _load_rag_structure_metrics_doc()

    # Enrich source_cells with actual cell values from grid data so that
    # _chunk_contains_source can match by value for methods without row/col.
    for q in queries:
        sid = q.get("sheet_id", "")
        grid = sheet_data.get(sid)
        if not grid:
            continue
        cell_map = {(c["row"], c["col"]): c for c in grid.get("cells", [])}
        for sc in q.get("source_cells", []):
            if "value" not in sc:
                cell = cell_map.get((sc.get("row"), sc.get("col")))
                if cell:
                    sc["value"] = cell.get("value", "")

    log("Loading embedding model...")
    embed_model = _load_embed_model()

    # Discover trained model architectures. Build a map of arch -> list
    # of candidate checkpoint paths (sorted by answerability desc).
    # Models are loaded lazily one-at-a-time during evaluation to avoid OOM.
    from predict import load_model_from_path, ARCH_REGISTRY
    experiments_dir = config.DATA_DIR / "experiments"
    arch_candidates: dict[str, list[Path]] = {}
    if experiments_dir.exists():
        raw_candidates: dict[str, list[Path]] = {}
        for exp_dir in sorted(experiments_dir.iterdir()):
            if not exp_dir.is_dir():
                continue
            ckpt_path = exp_dir / "best_model.pt"
            if not ckpt_path.exists():
                continue
            arch = exp_dir.name.rsplit("_seed", 1)[0]
            if architectures and arch not in architectures:
                continue
            raw_candidates.setdefault(arch, []).append(exp_dir)

        for arch, dirs in raw_candidates.items():
            def _score(d):
                rp = d / "results.json"
                if rp.exists():
                    try:
                        r = json.load(open(rp))
                        return (r.get("chunk_answerability_mean", 0),
                                r.get("chunk_f1_mean", 0))
                    except Exception:
                        pass
                return (0, 0)
            dirs.sort(key=_score, reverse=True)
            arch_candidates[arch] = [d / "best_model.pt" for d in dirs]
            log(f"  Found model: {arch} ({len(dirs)} checkpoint(s), best: {dirs[0].name})")

    if not arch_candidates:
        log("  WARNING: No experiment models found — graph methods will use default model")
        default_path = config.MODELS_DIR / "best_model.pt"
        if default_path.exists():
            arch_candidates["default"] = [default_path]

    log(f"  {len(arch_candidates)} architecture(s) available for prediction")

    internal_methods = {"graph", "graph_kg", "naive", "single_cell"}
    try:
        from rag_baselines import ADAPTERS
    except ImportError:
        log("Could not import rag_baselines — external methods unavailable")
        ADAPTERS = {}

    results: dict[str, Any] = {}

    # Oracle methods use human labels only (no checkpoint). Model methods are per-architecture.
    ORACLE_GRAPH_METHODS = frozenset({"graph_gold", "graph_row_gold", "graph_kg_gold"})
    MODEL_GRAPH_METHODS = frozenset({"graph", "graph_kg", "graph_row"})
    graph_methods = ORACLE_GRAPH_METHODS | MODEL_GRAPH_METHODS
    non_graph_methods = [m for m in methods if m not in graph_methods]
    requested_graph_methods = [m for m in methods if m in graph_methods]

    if sheet_data and ORACLE_GRAPH_METHODS.intersection(set(requested_graph_methods)):
        n_oracle_ready = sum(
            1 for sid, d in sheet_data.items() if _grid_has_rag_structure_labels(d)
        )
        log(
            f"Oracle (graph_*_gold): {n_oracle_ready}/{len(sheet_data)} held-out sheets have "
            f"human cell labels in JSON — others fall back to cache-only grids (weak oracle)."
        )

    for gm in requested_graph_methods:
        # ── Oracle: human-labeled chunking (upper bound vs model predictions) ──
        if gm in ORACLE_GRAPH_METHODS:
            method_key_base = gm
            log(f"\nBuilding chunks for method: {method_key_base} (oracle / human labels)")
            all_chunks: list[dict] = []
            try:
                for si, (sheet_id, data) in enumerate(sheet_data.items()):
                    n_cells = len(data.get("cells", []))
                    if n_cells > GNN_CELL_LIMIT:
                        fallback = _naive_chunks(data)
                        all_chunks.extend(fallback)
                        log(
                            f"  [{si+1}/{len(sheet_data)}] {sheet_id[:30]}... ({n_cells} cells) "
                            f"-> naive fallback ({len(fallback)} chunks)"
                        )
                        continue
                    if gm == "graph_gold":
                        all_chunks.extend(_graph_chunks_gold(data))
                    elif gm == "graph_row_gold":
                        all_chunks.extend(_graph_row_chunks_gold(data))
                    else:
                        all_chunks.extend(_graph_kg_chunks_gold(data))
                    if (si + 1) % 5 == 0:
                        log(f"  [{si+1}/{len(sheet_data)}] {len(all_chunks)} chunks so far")
            except Exception as e:
                log(f"  ERROR building chunks for {method_key_base}: {e}")
                err_batch = {
                    f"{method_key_base}/{strat}": {"n_chunks": 0, "error": str(e)}
                    for strat in strategies
                }
                results.update(err_batch)
                _write_rag_results_batch(
                    err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                continue

            log(f"  {len(all_chunks)} chunks")
            if not all_chunks:
                err_batch = {
                    f"{method_key_base}/{strat}": {"n_chunks": 0, "error": "no chunks produced"}
                    for strat in strategies
                }
                results.update(err_batch)
                _write_rag_results_batch(
                    err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                continue

            texts = [c.get("text", "") for c in all_chunks]
            log(f"  Embedding {len(texts)} chunks for {method_key_base}...")
            chunk_embs = _embed_texts(embed_model, texts)
            faiss_index = _build_faiss_index(chunk_embs)
            bm25_index = _build_bm25_index(texts)

            for strat in strategies:
                key = f"{method_key_base}/{strat}"
                if _should_skip_rag_resume(
                    key, resume, existing_agg, log, need_judge=use_llm_judge,
                ):
                    results[key] = dict(existing_agg[key])
                    continue
                log(f"  Evaluating {key}...")
                r = evaluate_method_v2(
                    all_chunks, queries, embed_model,
                    retrieval=strat,
                    use_llm_judge=use_llm_judge,
                    save_details=save_details,
                    _prebuilt_faiss=faiss_index,
                    _prebuilt_bm25=bm25_index,
                    _prebuilt_chunk_embs=chunk_embs,
                )
                results[key] = r
                _write_rag_results_batch(
                    {key: r}, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                for k, v in sorted(r.items()):
                    if k.startswith("_"):
                        continue
                    if isinstance(v, float):
                        log(f"    {k}: {v:.4f}")
                    else:
                        log(f"    {k}: {v}")
            continue

        # ── Model-based graph methods: one run per architecture ──
        for arch_name, ckpt_paths in arch_candidates.items():
            method_key_base = f"{gm}_{arch_name}"
            log(f"\nBuilding chunks for method: {method_key_base}")

            # Try each checkpoint until one loads and produces chunks
            model = None
            for ckpt_path in ckpt_paths:
                model = load_model_from_path(ckpt_path)
                if model:
                    log(f"  Loaded {ckpt_path.parent.name}")
                    break
            if not model:
                log(f"  ERROR: no loadable checkpoint for {arch_name}")
                err_batch = {
                    f"{method_key_base}/{strat}": {"n_chunks": 0, "error": "model load failed"}
                    for strat in strategies
                }
                results.update(err_batch)
                _write_rag_results_batch(
                    err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                continue

            log("  predict_sheet cache (one inference per sheet)...")
            pred_cache, timing_bundle = _fill_prediction_cache(
                sheet_data, model, log,
                timing_exclude_first_sheet=timing_exclude_first_sheet,
            )
            inf_agg = timing_bundle.get("aggregate") or {}
            inf_per = timing_bundle.get("per_sheet") or []

            def _attach_inference_timing(entry: dict[str, Any]) -> dict[str, Any]:
                out = dict(entry)
                out["inference_timing"] = inf_agg
                if save_structure_details and inf_per:
                    out["inference_timing_per_sheet"] = list(inf_per)
                # If this run did not emit per-sheet timing, keep any existing list from ``entry``.
                return out

            if not _should_skip_structure_resume(
                method_key_base, resume_structure, structure_metrics_doc, force_structure_metrics,
            ):
                try:
                    smeta = _compute_structure_vs_gold_for_arch(
                        sheet_data, pred_cache, save_per_sheet=save_structure_details,
                    )
                    prev_entry = structure_metrics_doc.get(method_key_base)
                    prev_dict = prev_entry if isinstance(prev_entry, dict) else {}
                    entry = {
                        **prev_dict,
                        "method_base": method_key_base,
                        "chunking": gm,
                        "computed_at": datetime.now(timezone.utc).isoformat(),
                        "rag_eval_query_count": len(queries),
                        "held_out_sheet_count": len(sheet_data),
                        **smeta,
                    }
                    entry.pop("structure_metrics_error", None)
                    structure_metrics_doc[method_key_base] = _attach_inference_timing(entry)
                    _write_rag_structure_metrics_doc(structure_metrics_doc)
                    log(
                        f"  Structure vs gold: n_sheets={smeta['n_sheets_with_gold']} → "
                        f"{config.RAG_EVAL_STRUCTURE_METRICS.name}"
                    )
                except Exception as e:
                    log(f"  WARNING: structure metrics failed for {method_key_base}: {e}")
                    logger.exception("structure metrics")
                    prev = structure_metrics_doc.get(method_key_base)
                    if not isinstance(prev, dict):
                        prev = {}
                    fallback = {
                        **prev,
                        "method_base": method_key_base,
                        "chunking": gm,
                        "computed_at": datetime.now(timezone.utc).isoformat(),
                        "structure_metrics_error": str(e),
                    }
                    structure_metrics_doc[method_key_base] = _attach_inference_timing(fallback)
                    _write_rag_structure_metrics_doc(structure_metrics_doc)
            else:
                log(f"  SKIP structure metrics for {method_key_base} (resume_structure)")
                prev = structure_metrics_doc.get(method_key_base)
                if not isinstance(prev, dict):
                    prev = {}
                merged = {**prev, "method_base": method_key_base, "chunking": gm}
                merged["inference_timing_updated_at"] = datetime.now(timezone.utc).isoformat()
                structure_metrics_doc[method_key_base] = _attach_inference_timing(merged)
                _write_rag_structure_metrics_doc(structure_metrics_doc)
                log(f"  Updated inference_timing only → {config.RAG_EVAL_STRUCTURE_METRICS.name}")

            all_chunks = []
            try:
                for si, (sheet_id, data) in enumerate(sheet_data.items()):
                    n_cells = len(data.get("cells", []))
                    if n_cells > GNN_CELL_LIMIT:
                        fallback = _naive_chunks(data)
                        all_chunks.extend(fallback)
                        log(f"  [{si+1}/{len(sheet_data)}] {sheet_id[:30]}... ({n_cells} cells) -> naive fallback ({len(fallback)} chunks)")
                        continue
                    try:
                        if gm == "graph":
                            all_chunks.extend(
                                _graph_chunks(
                                    data, model=model, sheet_id=sheet_id, pred_cache=pred_cache,
                                )
                            )
                        elif gm == "graph_row":
                            all_chunks.extend(
                                _graph_row_chunks(
                                    data, model=model, sheet_id=sheet_id, pred_cache=pred_cache,
                                )
                            )
                        elif gm == "graph_kg":
                            all_chunks.extend(
                                _graph_kg_chunks(
                                    data, model=model, sheet_id=sheet_id, pred_cache=pred_cache,
                                )
                            )
                    except (RuntimeError, MemoryError) as sheet_err:
                        log(f"  [{si+1}/{len(sheet_data)}] GNN failed on {sheet_id[:30]}... ({n_cells} cells): {sheet_err}")
                        log(f"    -> falling back to naive chunking for this sheet")
                        import gc; gc.collect()
                        fallback = _naive_chunks(data)
                        all_chunks.extend(fallback)
                    if (si + 1) % 5 == 0:
                        log(f"  [{si+1}/{len(sheet_data)}] {len(all_chunks)} chunks so far")
            except Exception as e:
                log(f"  ERROR building chunks for {method_key_base}: {e}")
                err_batch = {
                    f"{method_key_base}/{strat}": {"n_chunks": 0, "error": str(e)}
                    for strat in strategies
                }
                results.update(err_batch)
                _write_rag_results_batch(
                    err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                del model
                import gc; gc.collect()
                continue

            del model
            import gc; gc.collect()

            log(f"  {len(all_chunks)} chunks")
            if not all_chunks:
                err_batch = {
                    f"{method_key_base}/{strat}": {"n_chunks": 0, "error": "no chunks produced"}
                    for strat in strategies
                }
                results.update(err_batch)
                _write_rag_results_batch(
                    err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                continue

            texts = [c.get("text", "") for c in all_chunks]
            log(f"  Embedding {len(texts)} chunks for {method_key_base}...")
            chunk_embs = _embed_texts(embed_model, texts)
            faiss_index = _build_faiss_index(chunk_embs)
            bm25_index = _build_bm25_index(texts)

            for strat in strategies:
                key = f"{method_key_base}/{strat}"
                if _should_skip_rag_resume(
                    key, resume, existing_agg, log, need_judge=use_llm_judge,
                ):
                    results[key] = dict(existing_agg[key])
                    continue
                log(f"  Evaluating {key}...")
                r = evaluate_method_v2(
                    all_chunks, queries, embed_model,
                    retrieval=strat,
                    use_llm_judge=use_llm_judge,
                    save_details=save_details,
                    _prebuilt_faiss=faiss_index,
                    _prebuilt_bm25=bm25_index,
                    _prebuilt_chunk_embs=chunk_embs,
                )
                results[key] = r
                _write_rag_results_batch(
                    {key: r}, save_details=save_details, log_fn=log, existing_agg=existing_agg,
                )
                for k, v in sorted(r.items()):
                    if k.startswith("_"):
                        continue
                    if isinstance(v, float):
                        log(f"    {k}: {v:.4f}")
                    else:
                        log(f"    {k}: {v}")

    for method in non_graph_methods:
        log(f"\nBuilding chunks for method: {method}")
        all_chunks: list[dict] = []
        method_queries = queries

        if method in internal_methods:
            for sheet_id, data in sheet_data.items():
                if method == "naive":
                    all_chunks.extend(_naive_chunks(data))
                elif method == "single_cell":
                    all_chunks.extend(_single_cell_chunks(data))
        elif method in ADAPTERS:
            adapter_fn = ADAPTERS[method]
            # One chunk pass per unique (workbook path, sheet name) per held-out sheet.
            # Re-running the adapter for every query duplicated the corpus and skewed metrics.
            rep_q_by_sheet: dict[str, dict] = {}
            for q in queries:
                sid = q.get("sheet_id", "")
                if sid and sid not in rep_q_by_sheet:
                    rep_q_by_sheet[sid] = q
            seen_sheet_keys: set[tuple[str, str]] = set()
            n_adapter_calls = 0
            for sheet_id, _grid in sheet_data.items():
                q = rep_q_by_sheet.get(sheet_id)
                if q is None:
                    log(f"  No representative query for {sheet_id!r}; skipping adapter pass")
                    continue
                xlsx_path = _resolve_xlsx_path(q, sheet_id)
                sheet_name = q.get("sheet_name", "")
                if not xlsx_path or not sheet_name:
                    log(f"  Could not resolve xlsx for {sheet_id!r}")
                    continue
                key = _adapter_sheet_key(xlsx_path, sheet_name)
                if key in seen_sheet_keys:
                    continue
                seen_sheet_keys.add(key)
                try:
                    adapter_chunks = adapter_fn(xlsx_path, sheet_name)
                    all_chunks.extend(adapter_chunks)
                    n_adapter_calls += 1
                except Exception as e:
                    log(f"  Adapter {method} failed on {sheet_id} ({sheet_name!r}): {e}")
            log(
                f"  Adapter {method}: {n_adapter_calls} unique sheet(s) "
                f"({len(sheet_data)} held-out sheet(s), {len(rep_q_by_sheet)} rep queries)"
            )
        else:
            log(f"  Unknown method: {method}, skipping")
            continue

        log(f"  {len(all_chunks)} chunks")
        if not all_chunks:
            err_batch = {
                f"{method}/{strat}": {"n_chunks": 0, "error": "no chunks produced"}
                for strat in strategies
            }
            results.update(err_batch)
            _write_rag_results_batch(
                err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg,
            )
            continue

        texts = [c.get("text", "") for c in all_chunks]
        log(f"  Embedding {len(texts)} chunks for {method}...")
        chunk_embs = _embed_texts(embed_model, texts)
        faiss_index = _build_faiss_index(chunk_embs)
        bm25_index = _build_bm25_index(texts)

        for strat in strategies:
            key = f"{method}/{strat}"
            if _should_skip_rag_resume(
                key, resume, existing_agg, log, need_judge=use_llm_judge,
            ):
                results[key] = dict(existing_agg[key])
                continue
            log(f"  Evaluating {key}...")
            r = evaluate_method_v2(
                all_chunks, method_queries, embed_model,
                retrieval=strat,
                use_llm_judge=use_llm_judge,
                save_details=save_details,
                _prebuilt_faiss=faiss_index,
                _prebuilt_bm25=bm25_index,
                _prebuilt_chunk_embs=chunk_embs,
            )
            results[key] = r
            _write_rag_results_batch(
                {key: r}, save_details=save_details, log_fn=log, existing_agg=existing_agg,
            )
            for k, v in sorted(r.items()):
                if k.startswith("_"):
                    continue
                if isinstance(v, float):
                    log(f"    {k}: {v:.4f}")
                else:
                    log(f"    {k}: {v}")

    # Separate per-query details from aggregate metrics
    details_data = {}
    agg_results = {}
    for key, val in results.items():
        if isinstance(val, dict) and "_details" in val:
            details_data[key] = val.pop("_details")
        agg_results[key] = val

    out_path = config.DATA_DIR / "rag_eval_results.json"
    if out_path.exists():
        with open(out_path) as f:
            existing = json.load(f)
        existing.update(agg_results)
        agg_results = existing
    with open(out_path, "w") as f:
        json.dump(agg_results, f, indent=2)
    log(f"\nAggregate results saved to {out_path}")

    if details_data:
        details_path = config.DATA_DIR / "rag_eval_details.json"
        existing_details = {}
        if details_path.exists():
            with open(details_path) as f:
                existing_details = json.load(f)
        existing_details.update(details_data)
        with open(details_path, "w") as f:
            json.dump(existing_details, f, indent=2)
        log(f"Per-query details saved to {details_path}")

    for key in results:
        if key in details_data:
            results[key]["_details"] = details_data[key]
    return results


# ═══════════════════════════════════════════════════════════════════════════
# Legacy v1 evaluation (backward compatible)
# ═══════════════════════════════════════════════════════════════════════════

def evaluate_method(
    chunk_texts: list[str],
    queries: list[dict],
    embed_model,
    ks: tuple[int, ...] = (1, 5, 10),
) -> dict:
    """Evaluate retrieval quality for one chunking method (legacy v1)."""
    if not chunk_texts:
        return {f"recall@{k}": 0.0 for k in ks} | {"mrr": 0.0, "n_chunks": 0}

    chunk_embs = _embed_texts(embed_model, chunk_texts)
    index = _build_index(chunk_embs)

    q_texts = [q["question"] for q in queries]
    q_embs = _embed_texts(embed_model, q_texts)

    recalls = {k: 0 for k in ks}
    rr_sum = 0.0

    for qi, q in enumerate(queries):
        expected = q.get("expected_value", "")
        _, indices = _retrieve(index, q_embs[qi], k=max(ks))
        retrieved = [chunk_texts[int(idx)] for idx in indices if idx >= 0]

        first_hit = None
        for rank, text in enumerate(retrieved, start=1):
            if expected in text:
                if first_hit is None:
                    first_hit = rank
                break

        for k in ks:
            top_k = retrieved[:k]
            if any(expected in t for t in top_k):
                recalls[k] += 1

        if first_hit is not None:
            rr_sum += 1.0 / first_hit

    n_q = max(len(queries), 1)
    return {
        f"recall@{k}": recalls[k] / n_q for k in ks
    } | {
        "mrr": rr_sum / n_q,
        "n_chunks": len(chunk_texts),
    }


def run_rag_evaluation(
    queries_path: Path | None = None,
) -> dict:
    """Run full RAG evaluation across internal methods (legacy v1)."""
    if queries_path is None:
        queries_path = config.DATA_DIR / "rag_queries.json"

    if not queries_path.exists():
        print(f"No queries file at {queries_path}. Create it first.")
        print('Format: [{"question": "...", "expected_value": "...", "sheet_id": "..."}]')
        return {}

    with open(queries_path) as f:
        queries = json.load(f)

    print(f"Loaded {len(queries)} queries from {queries_path}")

    labeled_files = sorted(config.LABELED_DIR.glob("*.json"))
    if not labeled_files:
        print("No labeled data found.")
        return {}

    graph_texts: list[str] = []
    naive_texts: list[str] = []
    single_texts: list[str] = []

    for fp in labeled_files:
        with open(fp) as f:
            data = json.load(f)
        graph_texts.extend(c["text"] for c in _graph_chunks(data))
        naive_texts.extend(c["text"] for c in _naive_chunks(data))
        single_texts.extend(c["text"] for c in _single_cell_chunks(data))

    print(f"Chunks: graph={len(graph_texts)}, naive={len(naive_texts)}, single={len(single_texts)}")

    print("Loading embedding model...")
    embed_model = _load_embed_model()

    results = {}
    for name, texts in [("graph", graph_texts), ("naive", naive_texts),
                         ("single_cell", single_texts)]:
        print(f"\nEvaluating {name}...")
        r = evaluate_method(texts, queries, embed_model)
        results[name] = r
        for k, v in r.items():
            if isinstance(v, float):
                print(f"  {k}: {v:.4f}")
            else:
                print(f"  {k}: {v}")

    # Separate per-query details from aggregate metrics
    details_data = {}
    agg_results = {}
    for key, val in results.items():
        if isinstance(val, dict) and "_details" in val:
            details_data[key] = val.pop("_details")
        agg_results[key] = val

    out_path = config.DATA_DIR / "rag_eval_results.json"
    if out_path.exists():
        with open(out_path) as f:
            existing = json.load(f)
        existing.update(agg_results)
        agg_results = existing
    with open(out_path, "w") as f:
        json.dump(agg_results, f, indent=2)
    log(f"\nAggregate results saved to {out_path}")

    if details_data:
        details_path = config.DATA_DIR / "rag_eval_details.json"
        existing_details = {}
        if details_path.exists():
            with open(details_path) as f:
                existing_details = json.load(f)
        existing_details.update(details_data)
        with open(details_path, "w") as f:
            json.dump(existing_details, f, indent=2)
        log(f"Per-query details saved to {details_path}")

    for key in results:
        if key in details_data:
            results[key]["_details"] = details_data[key]
    return results


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="RAG evaluation")
    parser.add_argument("--queries", type=str, help="Path to queries JSON file (legacy v1)")
    parser.add_argument("--methods", type=str, default=None,
                        help="Comma-separated list of methods (v2)")
    parser.add_argument("--no-judge", action="store_true",
                        help="Skip LLM-as-a-judge evaluation")
    parser.add_argument("--architectures", type=str, default=None,
                        help="Comma-separated list of model architectures to evaluate")
    parser.add_argument("--strategies", type=str, default=None,
                        help="Comma-separated retrieval strategies (vector,bm25,hybrid_rrf,hybrid_rrf_rerank)")
    parser.add_argument("--save-details", action="store_true",
                        help="Save per-query details (retrieved chunks, judge verdicts)")
    parser.add_argument(
        "--no-resume",
        action="store_true",
        help="Re-run all method/strategy keys (ignore checkpoints in rag_eval_results.json)",
    )
    parser.add_argument(
        "--pull-drive",
        action="store_true",
        help="Before eval: download rag_eval_dataset + RAG structure JSONs from Drive (OAuth or service account)",
    )
    parser.add_argument(
        "--no-resume-structure",
        action="store_true",
        help="Recompute rag_eval_structure_metrics.json entries even if present",
    )
    parser.add_argument(
        "--force-structure-metrics",
        action="store_true",
        help="Same as --no-resume-structure (recompute structure-vs-gold)",
    )
    parser.add_argument(
        "--save-structure-details",
        action="store_true",
        help="Store per-sheet _per_sheet in rag_eval_structure_metrics.json",
    )
    parser.add_argument(
        "--timing-exclude-first-sheet",
        action="store_true",
        help="For inference_timing aggregates only: drop the first predict_sheet call from mean/p95 (warmup)",
    )
    args = parser.parse_args()

    if args.pull_drive:
        info = rag_eval_storage.prefetch_rag_eval_from_drive()
        print(json.dumps(info, indent=2))

    if args.queries:
        qp = Path(args.queries)
        run_rag_evaluation(qp)
    else:
        method_list = args.methods.split(",") if args.methods else None
        arch_list = args.architectures.split(",") if args.architectures else None
        strat_list = args.strategies.split(",") if args.strategies else None
        run_rag_evaluation_v2(
            methods=method_list,
            use_llm_judge=not args.no_judge,
            architectures=arch_list,
            retrieval_strategies=strat_list,
            save_details=args.save_details,
            resume=not args.no_resume,
            resume_structure=not args.no_resume_structure,
            force_structure_metrics=args.force_structure_metrics,
            save_structure_details=args.save_structure_details,
            timing_exclude_first_sheet=args.timing_exclude_first_sheet,
        )