predict.py

"""Run trained GNN inference on unlabeled spreadsheet grids.

Loads versioned models from data/models/ and predicts cell labels
for a raw grid (as returned by xlsx_client.fetch_xlsx_sheet).

Supports batch prediction of all unlabeled sheets after training,
storing results per model version in data/predictions/v{NNN}/.
A connected-component heuristic assigns table IDs based on spatial gaps.
"""

from __future__ import annotations

import json
import logging
from collections import deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import torch
import torch.nn.functional as F

import numpy as np

import config
from features import _content_features, _formula_features, FTYPE_TO_IDX
from models import (
    GATModel, GCNModel, GraphTransformerModel, MLPModel,
    SequenceTransformerModel, AdjTransformerModel,
    SpatialEdgeTransformerModel, DualModalityGNNModel,
    LAYOUTLM_ARCHS,
    build_model,
)
from label_utils import CELL_LABELS, COARSE_LABELS
from train_gnn import (
    EMBED_DIM, EDGE_TYPE_DIM, GraphConfig,
    _build_spatial_edges_cfg, _build_row_col_edges_cfg,
    _build_virtual_col_nodes, _build_virtual_row_nodes,
    _build_content_edge_features,
    EDGE_TYPE_SPATIAL,
)
import math

ARCH_REGISTRY: dict[str, type] = {
    "gat": GATModel,
    "gcn": GCNModel,
    "graph_transformer": GraphTransformerModel,
    "mlp": MLPModel,
    "seq_transformer": SequenceTransformerModel,
    "adj_transformer": AdjTransformerModel,
    "spatial_edge_transformer": SpatialEdgeTransformerModel,
    "dual_modality_gnn": DualModalityGNNModel,
}

SpreadsheetGNN = GATModel

logger = logging.getLogger("labeler.predict")

_cached_model: Optional[SpreadsheetGNN] = None
_cached_config: Optional[dict] = None
_cached_version: Optional[str] = None
_cached_output_labels: Optional[list[str]] = None
_cached_embed_model = None


def get_latest_model_version() -> Optional[str]:
    """Scan data/models/ for the highest version number. Returns e.g. '003' or None."""
    existing = sorted(config.MODELS_DIR.glob("model_v*.pt"))
    if not existing:
        return None
    return existing[-1].stem.split("_v")[1]


def get_best_model_path() -> Optional[Path]:
    """Return best_model.pt if it exists, else the latest versioned model."""
    best = config.MODELS_DIR / "best_model.pt"
    if best.exists():
        return best
    existing = sorted(config.MODELS_DIR.glob("model_v*.pt"))
    return existing[-1] if existing else None


def _checkpoint_kwargs_from_config(cfg: dict) -> dict:
    """Extra constructor kwargs from experiment checkpoints (e.g. CNN grid)."""
    from train_gnn import CNN_GRID_CHANNELS, CNN_GRID_OUT_DIM

    extra: dict = {}
    if cfg.get("use_cnn_grid"):
        extra["use_cnn_grid"] = True
        extra["cnn_grid_channels"] = int(cfg.get("cnn_grid_channels", CNN_GRID_CHANNELS))
        extra["cnn_grid_out_dim"] = int(cfg.get("cnn_grid_out_dim", CNN_GRID_OUT_DIM))
    return extra


def _align_structural_embed_to_dim(
    x_structural: torch.Tensor,
    x_embed: torch.Tensor,
    expected_dim: int,
) -> torch.Tensor:
    """Resize structural features so ``[structural | embed]`` has width ``expected_dim``.

    Truncates or zero-pads the **structural** block (prefix) when the live feature
    pipeline drifts from the checkpoint (e.g. 89 vs 71 structural dims).
    """
    from train_gnn import EMBED_DIM

    emb_d = x_embed.size(1)
    need_struct = expected_dim - emb_d
    if need_struct < 0:
        logger.warning(
            "Checkpoint expects in_dim=%d smaller than embedding width=%d; "
            "truncating concatenated features.",
            expected_dim, emb_d,
        )
        x_full = torch.cat([x_structural, x_embed], dim=1)
        return x_full[:, :expected_dim]

    s = x_structural.size(1)
    if s == need_struct:
        return torch.cat([x_structural, x_embed], dim=1)

    if s > need_struct:
        logger.warning(
            "Aligned structural features %d -> %d to match checkpoint in_dim=%d "
            "(older / fewer structural dims in trained weights).",
            s, need_struct, expected_dim,
        )
        x_structural = x_structural[:, :need_struct].contiguous()
    else:
        pad = torch.zeros(
            x_structural.size(0), need_struct - s,
            dtype=x_structural.dtype, device=x_structural.device,
        )
        logger.warning(
            "Aligned structural features %d -> %d to match checkpoint in_dim=%d "
            "(padding missing dims with zeros).",
            s, need_struct, expected_dim,
        )
        x_structural = torch.cat([x_structural, pad], dim=1)

    return torch.cat([x_structural, x_embed], dim=1)


def load_model_from_path(model_path: Path) -> Optional[SpreadsheetGNN]:
    """Load any architecture from a checkpoint path. Not cached."""
    if not model_path.exists():
        return None
    try:
        checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
        cfg = checkpoint["config"]
        arch = checkpoint.get("arch", "gat")
        num_classes = cfg.get("num_classes", 9)
        saved_etd = cfg.get("edge_type_dim", 3)

        ck_kw = _checkpoint_kwargs_from_config(cfg)
        model = build_model(
            arch,
            in_dim=cfg["in_dim"],
            hidden_dim=cfg["hidden_dim"],
            num_classes=num_classes,
            edge_dim=saved_etd,
            **ck_kw,
        )
        # train_compare enables coarse+level heads for non-LayoutLM; legacy train_gnn GAT does not.
        state = checkpoint["model_state_dict"]
        use_two_stage = cfg.get("use_two_stage")
        if use_two_stage is None:
            use_two_stage = any(k.startswith("coarse_head") for k in state.keys())
        if use_two_stage and arch not in LAYOUTLM_ARCHS:
            model.enable_two_stage_head(cfg["hidden_dim"])

        model.load_state_dict(state, strict=False)
        model.eval()
        model._edge_type_dim = saved_etd
        model._output_labels = checkpoint.get("cell_labels", COARSE_LABELS)
        model._legacy = (num_classes <= 9)
        model._arch = arch
        model._train_in_dim = cfg.get("in_dim")

        gc_dict = cfg.get("graph_config")
        model._graph_config = GraphConfig.from_dict(gc_dict) if gc_dict else None

        return model
    except Exception:
        logger.exception("Failed to load model from %s", model_path)
        return None


def load_model(version: Optional[str] = None) -> Optional[SpreadsheetGNN]:
    """Load a trained model, caching it in memory. Returns None if no model file.

    Args:
        version: Specific version string (e.g. '001'). If None, loads the best
                 available model (best_model.pt or latest versioned).
    """
    global _cached_model, _cached_config, _cached_version, _cached_output_labels

    if version is not None:
        model_path = config.MODELS_DIR / f"model_v{version}.pt"
    else:
        model_path = get_best_model_path()
        if model_path is None:
            model_path = config.DATA_DIR / "spreadsheet_gnn.pt"

    if not model_path.exists():
        return None

    version_key = version or model_path.stem
    if _cached_model is not None and _cached_version == version_key:
        return _cached_model

    model = load_model_from_path(model_path)
    if model is None:
        return None

    _cached_model = model
    _cached_config = torch.load(model_path, map_location="cpu", weights_only=False)["config"]
    _cached_version = version_key
    _cached_output_labels = model._output_labels
    logger.info(
        "Loaded model from %s (arch=%s, version=%s)",
        model_path, getattr(model, '_arch', 'unknown'), version_key,
    )
    return model


def _cell_features_inference(cell: dict, max_row: int, max_col: int,
                             median_font_size: float = 11.0,
                             col_row_stats: dict | None = None,
                             cohesion_stats: dict | None = None) -> list[float]:
    """Delegate to the shared _cell_features in train_gnn (keeps train/inference in sync)."""
    from train_gnn import _cell_features
    return _cell_features(cell, max_row, max_col, median_font_size, col_row_stats, cohesion_stats)


def _build_flat_spatial_edges(cells: list[dict], idx_map: dict) -> list[list[int]]:
    """Build spatial edges connecting all adjacent cells (no table boundaries)."""
    coord_to_idx = {}
    for c in cells:
        key = (c["row"], c["col"])
        if key in idx_map:
            coord_to_idx[key] = idx_map[key]

    src, dst = [], []
    for (r, c), node_idx in coord_to_idx.items():
        for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            neighbor = (r + dr, c + dc)
            if neighbor in coord_to_idx:
                src.append(node_idx)
                dst.append(coord_to_idx[neighbor])

    return [src, dst]


def _build_inference_edges(
    all_cells: list[dict], idx_map: dict, graph_cfg: GraphConfig,
) -> tuple[torch.Tensor, torch.Tensor, int]:
    """Build inference-time edges matching training GraphConfig.

    Returns (edge_index, edge_attr, num_virtual_nodes).
    Inference never has table labels, so table_groups is empty.
    """
    table_groups = {}  # no table scoping at inference
    num_real = len(all_cells)

    spatial_edges, spatial_types = _build_spatial_edges_cfg(
        all_cells, idx_map, table_groups, graph_cfg)
    rc_edges, rc_types = _build_row_col_edges_cfg(
        all_cells, idx_map, table_groups, graph_cfg)

    all_src = spatial_edges[0] + rc_edges[0]
    all_dst = spatial_edges[1] + rc_edges[1]
    all_types = spatial_types + rc_types

    num_virtual = 0
    if graph_cfg.virtual_col_nodes:
        vc_edges, vc_types, n_vc = _build_virtual_col_nodes(
            all_cells, idx_map, num_real + num_virtual)
        all_src += vc_edges[0]; all_dst += vc_edges[1]; all_types += vc_types
        num_virtual += n_vc

    if graph_cfg.virtual_row_nodes:
        vr_edges, vr_types, n_vr = _build_virtual_row_nodes(
            all_cells, idx_map, num_real + num_virtual)
        all_src += vr_edges[0]; all_dst += vr_edges[1]; all_types += vr_types
        num_virtual += n_vr

    total_nodes = num_real + num_virtual

    if not all_src:
        return (torch.zeros(2, 0, dtype=torch.long),
                torch.zeros(0, graph_cfg.edge_type_dim),
                num_virtual)

    edge_index = torch.tensor([all_src, all_dst], dtype=torch.long)

    net = graph_cfg.n_edge_types
    num_edges = len(all_types)
    edge_type_onehot = torch.zeros(num_edges, net)
    et_arr = torch.tensor(all_types, dtype=torch.long).clamp(max=net - 1)
    edge_type_onehot.scatter_(1, et_arr.unsqueeze(1), 1.0)

    parts = [edge_type_onehot]

    if graph_cfg.use_rel_pe:
        node_rows_f = torch.zeros(total_nodes, dtype=torch.float)
        node_cols_f = torch.zeros(total_nodes, dtype=torch.float)
        for i, c in enumerate(all_cells):
            node_rows_f[i] = c["row"]
            node_cols_f[i] = c["col"]

        src_t = torch.tensor(all_src, dtype=torch.long)
        dst_t = torch.tensor(all_dst, dtype=torch.long)
        dr = node_rows_f[dst_t] - node_rows_f[src_t]
        dc = node_cols_f[dst_t] - node_cols_f[src_t]

        REL_PE_DIM = 16
        half = REL_PE_DIM // 4
        div_term = torch.exp(
            torch.arange(0, half, dtype=torch.float) * -(math.log(10000.0) / max(half, 1))
        )
        rel_pe = torch.cat([
            torch.sin(dr.unsqueeze(-1) * div_term),
            torch.cos(dr.unsqueeze(-1) * div_term),
            torch.sin(dc.unsqueeze(-1) * div_term),
            torch.cos(dc.unsqueeze(-1) * div_term),
        ], dim=1)
        parts.append(rel_pe)

    if graph_cfg.content_edge_features:
        content_feat = _build_content_edge_features(all_cells, all_src, all_dst)
        parts.append(content_feat)

    if graph_cfg.distance_edge_features:
        if not graph_cfg.use_rel_pe:
            node_rows_f = torch.zeros(total_nodes, dtype=torch.float)
            node_cols_f = torch.zeros(total_nodes, dtype=torch.float)
            for i, c in enumerate(all_cells):
                node_rows_f[i] = c["row"]
                node_cols_f[i] = c["col"]
            src_t = torch.tensor(all_src, dtype=torch.long)
            dst_t = torch.tensor(all_dst, dtype=torch.long)
            dr = node_rows_f[dst_t] - node_rows_f[src_t]
            dc = node_cols_f[dst_t] - node_cols_f[src_t]
        max_row = max((c["row"] for c in all_cells), default=0) + 1
        max_col = max((c["col"] for c in all_cells), default=0) + 1
        manhattan = dr.abs() + dc.abs()
        norm_dist = manhattan / max(max_row + max_col, 1)
        dist_decay = torch.exp(-0.3 * manhattan)
        dist_feat = torch.stack([norm_dist, dist_decay], dim=1)
        parts.append(dist_feat)

    edge_attr = torch.cat(parts, dim=1)
    return edge_index, edge_attr, num_virtual


def _assign_table_ids(cells: list[dict], predictions: list[dict]) -> None:
    """Heuristic: assign table IDs via connected components of non-empty cells.

    Cells predicted as metadata that are spatially isolated from data cells
    get table 0 (sheet-level). Connected components of data cells each get
    a distinct table ID starting from 1.
    """
    occupied = set()
    coord_to_pred = {}
    for p in predictions:
        key = (p["row"], p["col"])
        occupied.add(key)
        coord_to_pred[key] = p

    data_labels = {"value", "row_header_1", "col_header_1", "aggregation"}
    data_coords = set()
    for p in predictions:
        if p["label"] in data_labels or p["label"].startswith("row_header_") or p["label"].startswith("col_header_"):
            data_coords.add((p["row"], p["col"]))

    visited = set()
    components = []
    for coord in sorted(data_coords):
        if coord in visited:
            continue
        component = set()
        queue = deque([coord])
        while queue:
            cur = queue.popleft()
            if cur in visited:
                continue
            visited.add(cur)
            if cur not in occupied:
                continue
            component.add(cur)
            r, c = cur
            for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nb = (r + dr, c + dc)
                if nb in occupied and nb not in visited:
                    queue.append(nb)
        if component:
            components.append(component)

    # Merge vertically adjacent components whose column spans overlap
    # significantly. This prevents header rows from getting a different
    # table ID than their data rows when separated by an empty row.
    def _col_span(comp):
        return {c for _, c in comp}

    def _row_range(comp):
        rows = [r for r, _ in comp]
        return min(rows), max(rows)

    merged = True
    while merged:
        merged = False
        for i in range(len(components)):
            if components[i] is None:
                continue
            cols_i = _col_span(components[i])
            rmin_i, rmax_i = _row_range(components[i])
            for j in range(i + 1, len(components)):
                if components[j] is None:
                    continue
                cols_j = _col_span(components[j])
                rmin_j, rmax_j = _row_range(components[j])
                gap = min(abs(rmin_j - rmax_i), abs(rmin_i - rmax_j))
                if gap > 3:
                    continue
                overlap = len(cols_i & cols_j)
                union = len(cols_i | cols_j)
                if union > 0 and overlap / union > 0.4:
                    components[i] = components[i] | components[j]
                    components[j] = None
                    cols_i = _col_span(components[i])
                    rmin_i, rmax_i = _row_range(components[i])
                    merged = True
        components = [c for c in components if c is not None]

    for p in predictions:
        key = (p["row"], p["col"])
        if p["label"] == "metadata":
            is_near_data = False
            r, c = key
            for dr in range(-2, 3):
                for dc in range(-2, 3):
                    if (r + dr, c + dc) in data_coords:
                        is_near_data = True
                        break
                if is_near_data:
                    break
            if not is_near_data:
                p["table"] = 0
                continue

        for tid, comp in enumerate(components, start=1):
            if key in comp:
                p["table"] = tid
                break
        else:
            if p["label"] == "metadata":
                p["table"] = 0
            elif p["label"] == "empty":
                p["table"] = 0
            else:
                p["table"] = 1


def _assign_header_levels(predictions: list[dict]) -> None:
    """Assign header levels based on position within each table.

    Convention: L1 = most specific (closest to data), higher = broader.

    Column headers: bottom-to-top numbering (closest to data = L1).
    Row headers: right-to-left numbering (closest to data = L1).
    """
    from collections import defaultdict

    tables = defaultdict(list)
    for p in predictions:
        tables[p.get("table", 0)].append(p)

    for tid, preds in tables.items():
        col_header_rows = sorted({p["row"] for p in preds if p["label"] == "col_header"})
        row_to_level = {row: lvl for lvl, row in enumerate(reversed(col_header_rows), start=1)}

        row_header_cols = sorted({p["col"] for p in preds if p["label"] == "row_header"})
        col_to_level = {col: lvl for lvl, col in enumerate(reversed(row_header_cols), start=1)}

        for p in preds:
            if p["label"] == "col_header":
                level = row_to_level.get(p["row"], 1)
                p["label"] = f"col_header_{level}"
            elif p["label"] == "row_header":
                level = col_to_level.get(p["col"], 1)
                p["label"] = f"row_header_{level}"


def _get_embed_model():
    """Lazy-load the sentence-transformer model for text embeddings."""
    global _cached_embed_model
    if _cached_embed_model is None:
        try:
            from embed_text import _get_model
            _cached_embed_model = _get_model()
        except Exception:
            logger.warning("Could not load embedding model; using zero embeddings")
    return _cached_embed_model


def _is_numeric_value(val: str) -> bool:
    try:
        float(val.replace(",", "").replace(" ", "").replace("%", ""))
        return True
    except ValueError:
        return False


def _generate_embeddings(cells: list[dict]) -> torch.Tensor:
    """Generate text embeddings for cells, matching training behavior.

    Text cells get real embeddings; numeric-only cells get zeros
    (same logic as embed_text.py).
    """
    model = _get_embed_model()
    if model is None:
        return torch.zeros(len(cells), EMBED_DIM, dtype=torch.float)

    text_indices = []
    texts = []
    for i, c in enumerate(cells):
        val = c.get("value", "").strip()
        if val and not _is_numeric_value(val):
            text_indices.append(i)
            texts.append(f"passage: {val}")

    embeddings = np.zeros((len(cells), EMBED_DIM), dtype=np.float32)

    if texts:
        vecs = model.encode(texts, batch_size=64, show_progress_bar=False, normalize_embeddings=True)
        for idx, vec in zip(text_indices, vecs):
            embeddings[idx] = vec

    return torch.tensor(embeddings, dtype=torch.float)


def _cluster_by_table_edges(
    edge_index: torch.Tensor, same_table_probs: torch.Tensor,
    num_nodes: int, threshold: float = 0.5,
) -> list[int]:
    """Connected components on predicted same-table edges."""
    import networkx as nx

    G = nx.Graph()
    G.add_nodes_from(range(num_nodes))
    mask = same_table_probs > threshold
    for i in range(edge_index.size(1)):
        if mask[i]:
            G.add_edge(edge_index[0, i].item(), edge_index[1, i].item())
    components = sorted(nx.connected_components(G), key=len, reverse=True)
    table_ids = [0] * num_nodes
    for tid, comp in enumerate(components, start=1):
        for node_idx in comp:
            table_ids[node_idx] = tid
    return table_ids


import re as _re

_NUMERIC_RE = _re.compile(
    r"^[\s]*[\$€£¥]?\s*[\-\+]?\s*[\d,]+(?:\.\d+)?\s*[%]?\s*$"
)


def _postprocess_value_attribute(
    cells: list[dict], predictions: list[dict]
) -> None:
    """Reclassify attribute cells as value using three complementary heuristics.

    1. **Zero-value fallback** (original): when the model predicts zero value
       cells, flip numeric attribute cells that have text above in the same column.
    2. **Column-level voting**: if >60% of non-empty predictions in a column are
       already ``value``, flip remaining numeric ``attribute`` cells in that column.
    3. **Low-confidence flip**: if an attribute prediction has confidence < 0.7 and
       the cell is numeric with a text cell above, flip to ``value``.
    """
    from collections import defaultdict

    pred_by_coord = {(p["row"], p["col"]): p for p in predictions}
    cell_by_coord = {(c["row"], c["col"]): c for c in cells}

    has_value = any(p["label"] == "value" for p in predictions)
    reclassified_rows = set()

    def _is_numeric_cell(p: dict) -> bool:
        cell = cell_by_coord.get((p["row"], p["col"]))
        if cell is None:
            return False
        return bool(_NUMERIC_RE.match(cell.get("value", "")))

    def _has_text_above(p: dict) -> bool:
        for scan_r in range(p["row"] - 1, -1, -1):
            above_cell = cell_by_coord.get((scan_r, p["col"]))
            if above_cell is None:
                continue
            above_val = above_cell.get("value", "").strip()
            if above_val and not _NUMERIC_RE.match(above_val):
                return True
        return False

    # --- Heuristic 1: zero-value fallback (original behavior) ---
    if not has_value:
        for p in predictions:
            if p["label"] != "attribute":
                continue
            if _is_numeric_cell(p) and _has_text_above(p):
                p["label"] = "value"
                reclassified_rows.add(p["row"])

    # --- Heuristic 2: column-level voting ---
    preds_by_col: dict[int, list[dict]] = defaultdict(list)
    for p in predictions:
        if p["label"] in ("value", "attribute", "aggregation"):
            preds_by_col[p["col"]].append(p)

    for col_idx, col_preds in preds_by_col.items():
        value_count = sum(1 for p in col_preds if p["label"] == "value")
        total = len(col_preds)
        if total > 0 and value_count / total > 0.6:
            for p in col_preds:
                if p["label"] == "attribute" and _is_numeric_cell(p):
                    p["label"] = "value"
                    reclassified_rows.add(p["row"])

    # --- Heuristic 3: low-confidence flip ---
    for p in predictions:
        if p["label"] != "attribute":
            continue
        if p.get("confidence", 1.0) < 0.7:
            if _is_numeric_cell(p) and _has_text_above(p):
                p["label"] = "value"
                reclassified_rows.add(p["row"])

    # For rows that gained reclassified value cells, promote the
    # leftmost remaining text attribute cell to row_header_1 so that
    # value chunks get proper row context.
    if reclassified_rows:
        rows_attrs: dict[int, list[dict]] = defaultdict(list)
        for p in predictions:
            if p["row"] in reclassified_rows and p["label"] == "attribute":
                cell = cell_by_coord.get((p["row"], p["col"]))
                if cell:
                    val = cell.get("value", "").strip()
                    if val and not _NUMERIC_RE.match(val):
                        rows_attrs[p["row"]].append(p)
        for row, attrs in rows_attrs.items():
            attrs.sort(key=lambda a: a["col"])
            attrs[0]["label"] = "row_header_1"


@torch.no_grad()
def predict_sheet(
    grid: dict,
    use_heuristics: bool = False,
    model_override=None,
) -> Optional[list[dict]]:
    """Run GNN inference on a raw grid and return predicted labels.

    Args:
        grid: Sheet data dict with ``cells``, ``grid_size``, etc.
        use_heuristics: Force heuristic post-processing even for new models
        model_override: Use this model instead of the cached/default one.
            (useful for ablation comparison).

    Returns a list of {"row", "col", "label", "confidence", "table"} dicts,
    or None if no model is available.
    """
    model = model_override or load_model()
    if model is None:
        return None

    output_labels = getattr(model, "_output_labels", COARSE_LABELS)
    legacy = getattr(model, "_legacy", True)

    cells = grid.get("cells", [])
    all_cells = [c for c in cells if c.get("value", "") != ""]
    if len(all_cells) < 3:
        return None

    max_row = max(c["row"] for c in all_cells) + 1
    max_col = max(c["col"] for c in all_cells) + 1

    idx_map = {}
    for i, c in enumerate(all_cells):
        idx_map[(c["row"], c["col"])] = i

    font_sizes = [c.get("font_size") for c in all_cells if c.get("font_size")]
    median_fs = float(sorted(font_sizes)[len(font_sizes) // 2]) if font_sizes else 11.0

    from train_gnn import _compute_col_row_stats, _compute_cohesion_stats
    crs = _compute_col_row_stats(all_cells)
    coh = _compute_cohesion_stats(all_cells)

    x_structural = torch.tensor(
        [_cell_features_inference(c, max_row, max_col, median_fs, crs, coh) for c in all_cells],
        dtype=torch.float,
    )

    x_embed = _generate_embeddings(all_cells)

    expected_dim = getattr(model, "_train_in_dim", None)
    if expected_dim is None and _cached_config is not None:
        expected_dim = _cached_config.get("in_dim")
    if expected_dim is None:
        expected_dim = x_structural.size(1) + x_embed.size(1)

    x = _align_structural_embed_to_dim(x_structural, x_embed, int(expected_dim))

    graph_cfg = getattr(model, "_graph_config", None)
    num_virtual = 0

    if graph_cfg is not None:
        edge_index, edge_attr, num_virtual = _build_inference_edges(
            all_cells, idx_map, graph_cfg)
        if edge_index.size(1) == 0:
            return None
        if num_virtual > 0:
            x = torch.cat([x, torch.zeros(num_virtual, x.size(1))], dim=0)
    else:
        edges = _build_flat_spatial_edges(all_cells, idx_map)
        if not edges[0]:
            return None
        edge_index = torch.tensor(edges, dtype=torch.long)
        model_etd = getattr(model, "_edge_type_dim", EDGE_TYPE_DIM)
        edge_attr = torch.zeros(edge_index.size(1), model_etd)
        edge_attr[:, 0] = 1.0

    grid_kwargs = {}
    if getattr(model, "use_cnn_grid", False):
        from train_gnn import _build_cnn_grid, CNN_GRID_MAX_ROWS, CNN_GRID_MAX_COLS
        grid_features = _build_cnn_grid(all_cells, max_row, max_col)
        grid_H = min(max_row, CNN_GRID_MAX_ROWS)
        grid_W = min(max_col, CNN_GRID_MAX_COLS)
        num_real = len(all_cells)
        grid_row_col = torch.zeros(num_real + num_virtual, 2, dtype=torch.long)
        grid_valid = torch.zeros(num_real + num_virtual, dtype=torch.bool)
        for i, c in enumerate(all_cells):
            r, col = c["row"], c["col"]
            if r < grid_H and col < grid_W:
                grid_row_col[i, 0] = r
                grid_row_col[i, 1] = col
                grid_valid[i] = True
        grid_kwargs = {
            "grid_features": grid_features,
            "grid_row_col": grid_row_col,
            "grid_valid": grid_valid,
        }

    is_graph_learner = getattr(model, "_is_graph_learner", False)
    if is_graph_learner:
        node_rows = torch.tensor([c["row"] for c in all_cells], dtype=torch.long)
        node_cols = torch.tensor([c["col"] for c in all_cells], dtype=torch.long)
        if num_virtual > 0:
            node_rows = torch.cat([node_rows, torch.zeros(num_virtual, dtype=torch.long)])
            node_cols = torch.cat([node_cols, torch.zeros(num_virtual, dtype=torch.long)])
        if hasattr(model, "_struct_refine_active"):
            model._struct_refine_active = True
        h = model(x, edge_index, edge_attr, node_rows=node_rows, node_cols=node_cols)
    else:
        h = model(x, edge_index, edge_attr, **grid_kwargs)

    logits = model.classify_nodes(h)
    probs = F.softmax(logits, dim=-1)
    pred_classes = probs.argmax(dim=-1).cpu().numpy()
    confidences = probs.max(dim=-1).values.cpu().numpy()

    # Build label predictions using the correct label list for this model
    predictions = []
    for i, c in enumerate(all_cells):
        cls_idx = int(pred_classes[i])
        label_name = output_labels[cls_idx] if cls_idx < len(output_labels) else "value"

        predictions.append({
            "row": c["row"],
            "col": c["col"],
            "label": label_name,
            "confidence": round(float(confidences[i]), 3),
            "table": 1,
        })

    pred_adjs = {}
    if legacy or use_heuristics:
        _assign_table_ids(all_cells, predictions)
        _assign_header_levels(predictions)
    elif is_graph_learner and hasattr(model, "predict_adjacency"):
        import networkx as _nx
        from models import SpatialEdgeTransformerModel, SpatialBilinearTransformerModel

        if isinstance(model, SpatialEdgeTransformerModel):
            adj_pred = model.predict_adjacency(h, node_rows, node_cols)
        elif isinstance(model, SpatialBilinearTransformerModel):
            adj_pred = model.predict_adjacency(h, node_rows, node_cols)
        else:
            adj_pred = model.predict_adjacency(h)

        for adj_type, logit_mat in adj_pred.items():
            prob_mat = torch.sigmoid(logit_mat)
            sym = (prob_mat + prob_mat.T) / 2 if adj_type != "hier" else prob_mat
            edge_mask = sym > 0.5
            pairs = edge_mask.nonzero(as_tuple=False)
            pred_adjs[adj_type] = {(r[0].item(), r[1].item()) for r in pairs}

        if "table" in pred_adjs and pred_adjs["table"]:
            G = _nx.Graph()
            G.add_nodes_from(range(len(all_cells)))
            G.add_edges_from(pred_adjs["table"])
            for comp_idx, comp in enumerate(
                sorted(_nx.connected_components(G), key=len, reverse=True), start=1
            ):
                for node in comp:
                    if node < len(predictions):
                        predictions[node]["table"] = comp_idx
        else:
            _assign_table_ids(all_cells, predictions)
    else:
        if hasattr(model, "predict_table_edges"):
            table_logits = model.predict_table_edges(h, edge_index)
            table_probs = torch.sigmoid(table_logits)
            table_ids = _cluster_by_table_edges(
                edge_index, table_probs, len(all_cells),
            )
            for i, p in enumerate(predictions):
                p["table"] = table_ids[i]
        else:
            _assign_table_ids(all_cells, predictions)

    _postprocess_value_attribute(all_cells, predictions)

    return predictions


# ═══════════════════════════════════════════════════════════════════════════
# Batch prediction & pre-computed prediction storage
# ═══════════════════════════════════════════════════════════════════════════

def _safe_filename(drive_file_id: str, sheet_name: str) -> str:
    """Build a safe filename from drive_file_id and sheet_name."""
    return f"{drive_file_id}_{sheet_name}".replace("/", "_").replace(" ", "_") + ".json"


def batch_predict_all(version: str) -> None:
    """Predict all unlabeled sheets with the given model version and save to disk.

    Stores results in data/predictions/v{version}/{drive_file_id}_{sheet_name}.json.
    Skips sheets that already have a prediction file for this version (idempotent).
    """
    import drive_client
    import metadata_client
    import xlsx_client

    model = load_model(version)
    if model is None:
        print(f"No model found for version v{version}. Skipping batch prediction.")
        return

    version_dir = config.PREDICTIONS_DIR / f"v{version}"
    version_dir.mkdir(parents=True, exist_ok=True)

    rows = metadata_client.get_all_rows()
    pending = [r for r in rows if r["status"] in {"unlabelled", "in_progress"}]

    if not pending:
        print("No unlabeled sheets to predict.")
        return

    print(f"Predicting {len(pending)} sheets with model v{version}...")
    predicted = 0
    skipped = 0
    errors = 0

    for i, row in enumerate(pending, 1):
        fname = _safe_filename(row["drive_file_id"], row["sheet_name"])
        pred_path = version_dir / fname

        if pred_path.exists():
            skipped += 1
            continue

        try:
            local_path = str(config.TEMP_DIR / f"{row['drive_file_id']}.xlsx")
            if not Path(local_path).exists():
                drive_client.download_file(row["drive_file_id"], local_path)

            grid = xlsx_client.fetch_xlsx_sheet(local_path, row["sheet_name"])
            preds = predict_sheet(grid)

            if preds:
                pred_data = {
                    "model_version": f"v{version}",
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                    "drive_file_id": row["drive_file_id"],
                    "sheet_name": row["sheet_name"],
                    "predictions": preds,
                }
                pred_path.write_text(
                    json.dumps(pred_data, ensure_ascii=False, indent=1)
                )
                predicted += 1
                print(f"  [{i}/{len(pending)}] Predicted {row['sheet_name']} from {row['file_name']}")
            else:
                skipped += 1
                print(f"  [{i}/{len(pending)}] Skipped {row['sheet_name']} (too few cells)")

        except Exception as e:
            errors += 1
            logger.warning("Error predicting %s/%s: %s", row["file_name"], row["sheet_name"], e)
            print(f"  [{i}/{len(pending)}] Error: {row['sheet_name']} - {e}")

    print(f"\nBatch prediction complete: {predicted} predicted, {skipped} skipped, {errors} errors")


def load_predictions(
    drive_file_id: str, sheet_name: str
) -> tuple[Optional[list[dict]], Optional[str]]:
    """Load pre-computed predictions from the latest model version directory.

    Returns (predictions_list, model_version) or (None, None) if not found.
    """
    latest = get_latest_model_version()
    if latest is None:
        return None, None

    version_dir = config.PREDICTIONS_DIR / f"v{latest}"
    if not version_dir.exists():
        return None, None

    fname = _safe_filename(drive_file_id, sheet_name)
    pred_path = version_dir / fname

    if not pred_path.exists():
        return None, None

    try:
        data = json.loads(pred_path.read_text())
        return data.get("predictions"), data.get("model_version", f"v{latest}")
    except (json.JSONDecodeError, OSError) as e:
        logger.warning("Error reading prediction file %s: %s", pred_path, e)
        return None, None