script.py

"""S23DR 2026 — submission script (trained WireframeDiffusion).

HF Competitions entry point: the test runner executes ``python script.py``
in the submission repo's working directory. We therefore:

  * Load ``params.json`` (mirrors ``usm3d/empty_submission_2026``).
  * Resolve the dataset (test server: ``/tmp/data``; local: snapshot_download).
  * Stream samples from the webdataset, **load the diffusion model exactly
    once** on CUDA, run inference, and write ``submission.json``.

Online preprocessing
--------------------
Uses V10 online preprocessing by default — ``build_scene_input`` produces
the V10 feature point cloud (COLMAP+camera tokens + depth-unprojected
tokens, per-point RGB, soft top-2 gestalt). Matches the cache layout the
model was trained on.

Efficiency notes
----------------
* Model + checkpoint are loaded once at module-main, not per worker — a
  process-pool would otherwise pay full GPU-init cost per sample.
* CPU preprocessing (COLMAP parse, depth back-projection, gestalt rasterise)
  runs in a small ``ThreadPoolExecutor`` so the next sample is ready by the
  time the GPU finishes the current one. Threads — not processes — keep the
  model on a single device while NumPy / Pillow / pycolmap drop the GIL.
* On any per-sample failure we fall back to the empty solution; the run
  never aborts mid-dataset.

Specifying the checkpoint
-------------------------
Resolution order (first hit wins):
  1. ``--ckpt PATH``               (CLI, for local testing)
  2. ``S23DR_CKPT`` env var        (set by SLURM / docker entrypoint)
  3. ``params['ckpt']``            (extra field in params.json)
  4. ``./test_checkpoint.pth``     (default — what you'd commit to the
                                    submission repo)

The script auto-resolves model architecture (small vs. large) from the
``args`` dict embedded in the checkpoint, so the same script handles both.

Local sanity check
------------------
Before pushing a submission repo, run::

    python script.py --sanity

This loads the public ``usm3d/hoho22k_2026_trainval`` dataset (which has
GT wireframes), runs the trained model on a few validation samples, and
prints per-sample + mean HSS / F1 / IoU. It writes no submission.json —
its only job is to verify that the checkpoint loads, the V10 online
preprocessing path works end-to-end, and the model produces sane output
before you submit.
"""

from __future__ import annotations

import argparse
import io
import json
import os
import queue
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch
try:
    from datasets import load_dataset
except ModuleNotFoundError:
    load_dataset = None
from tqdm import tqdm

# --- our project imports ----------------------------------------------------
from hoho2025.example_solutions import read_colmap_rec, _cam_matrix_from_image
from data.preprocess import (
    build_scene_input,
    _image_to_rgb_array,
    gestalt_img_to_ids,
    _camera_for_image,
    _resolve_colmap_image,
    GESTALT_CLASSES,
)
from data.dataset import _seed_from_order_id
from data.stage2_build import build_stage2_scene
from data.stage2_dataset import stage2_row_to_sample, _pad_init_verts
from data.fps import fps_subsample_stage2_batch
from models.scene_encoder import SceneEncoder
from models.denoiser import VertexDenoiser
from models.diffusion import WireframeDiffusion
from models.stage2 import Stage2Diffusion


# ---------------------------------------------------------------------------
# Empty / safe fallback
# ---------------------------------------------------------------------------

def empty_solution() -> Tuple[np.ndarray, List[Tuple[int, int]]]:
    """Minimal valid solution: 2 coincident vertices + 1 edge."""
    return np.zeros((2, 3)), [(0, 1)]


# ---------------------------------------------------------------------------
# Model loading (called once)
# ---------------------------------------------------------------------------

def resolve_ckpt_path(cli_arg: Optional[str], params: Dict[str, Any]) -> Path:
    for candidate, src in [
        (cli_arg,                     "--ckpt"),
        (os.environ.get("S23DR_CKPT"), "$S23DR_CKPT"),
        (params.get("ckpt"),          "params.json[ckpt]"),
        ("./test_checkpoint.pth",     "default ./test_checkpoint.pth"),
    ]:
        if candidate:
            print(f"[ckpt] using {candidate}  (from {src})", flush=True)
            return Path(candidate)
    raise SystemExit("[ckpt] no checkpoint specified")


def resolve_stage2_ckpt_path(cli_arg: Optional[str], params: Dict[str, Any]) -> Path:
    """Stage-2 checkpoint resolution. The two-stage architecture is mandatory;
    if no stage-2 ckpt is configured the script aborts. Resolution mirrors
    ``resolve_ckpt_path``: CLI → env → params.json → committed default."""
    for candidate, src in [
        (cli_arg,                                "--stage2_ckpt"),
        (os.environ.get("S23DR_STAGE2_CKPT"),    "$S23DR_STAGE2_CKPT"),
        (params.get("stage2_ckpt"),              "params.json[stage2_ckpt]"),
        ("./test_checkpoint_stage2.pth",         "default ./test_checkpoint_stage2.pth"),
    ]:
        if candidate:
            print(f"[stage2/ckpt] using {candidate}  (from {src})", flush=True)
            return Path(candidate)
    raise SystemExit("[stage2/ckpt] no stage-2 checkpoint specified — "
                     "two-stage pipeline is mandatory.")


def load_model(ckpt_path: Path, device: torch.device):
    """Return (model, ckpt_args) — model is .eval() on `device`.

    Architecture hyperparameters (d_model, n_heads, n_layers, d_ff,
    n_pool, encoder layer counts, ...) are read from the ``args`` dict
    embedded in the checkpoint. Same loader handles small and large
    configs transparently.
    """
    print(f"[model] loading {ckpt_path} on {device}", flush=True)
    t0 = time.perf_counter()
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    ckpt_args = ckpt.get("args", {}) or {}

    # RGB branch is present in v10+ caches and was used when training models
    # on those shards. Older sem_v7 checkpoints disable the RGB MLP.
    cache_version = str(ckpt_args.get("cache_version", "sem_v10")).lower()
    use_rgb = cache_version not in {"v7", "sem_v7"}

    enc = SceneEncoder(
        d_model=ckpt_args.get("d_model", 256),
        n_heads=ckpt_args.get("n_heads", 8),
        d_ff=ckpt_args.get("d_ff", 1024),
        n_full_layers=ckpt_args.get("n_encoder_full_layers", 2),
        n_pool=ckpt_args.get("n_pool", 1024),
        n_pool_cross_layers=ckpt_args.get("n_encoder_pool_cross_layers", 1),
        n_pool_layers=ckpt_args.get("n_encoder_pool_layers", 2),
        use_rgb=use_rgb,
    )
    den = VertexDenoiser(
        d_model=ckpt_args.get("d_model", 256),
        n_heads=ckpt_args.get("n_heads", 8),
        n_layers=ckpt_args.get("n_layers", 8),
        d_ff=ckpt_args.get("d_ff", 1024),
        k_verts=ckpt_args.get("k_verts", 64),
    )
    model = WireframeDiffusion(
        enc, den,
        noise_sigma_xyz=ckpt_args.get("noise_sigma_xyz", 1.0),
        init_from_scene=ckpt_args.get("init_from_scene", False),
        scene_init_jitter=ckpt_args.get("scene_init_jitter", 0.05),
    ).to(device)

    state = ckpt.get("model") or ckpt.get("state_dict")
    if state is None:
        raise RuntimeError(f"checkpoint {ckpt_path} has no 'model' or 'state_dict' key")
    # Strip DDP prefix if present.
    if any(k.startswith("module.") for k in state):
        state = {k[len("module."):]: v for k, v in state.items()}
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing or unexpected:
        print(f"[model] state_dict mismatch: missing={missing}  unexpected={unexpected}",
              flush=True)
    model.eval()

    print(f"[model] loaded in {time.perf_counter() - t0:.1f}s  "
          f"d_model={ckpt_args.get('d_model')} n_layers={ckpt_args.get('n_layers')} "
          f"use_rgb={use_rgb}  step={ckpt.get('step')}  best_hss={ckpt.get('best_hss')}",
          flush=True)
    return model, ckpt_args


def load_stage2_model(ckpt_path: Path, device: torch.device):
    """Build a Stage2Diffusion model from its checkpoint. Architecture is read
    from the embedded ``args`` dict so the same loader handles small / large
    variants. Returns (model, ckpt_args).
    """
    print(f"[stage2/model] loading {ckpt_path} on {device}", flush=True)
    t0 = time.perf_counter()
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    a = ckpt.get("args", {}) or {}
    if not isinstance(a, dict):
        a = vars(a)

    enc = SceneEncoder(
        d_model=a.get("d_model", 256),
        n_heads=a.get("n_heads", 8),
        d_ff=a.get("d_ff", 1024),
        n_full_layers=a.get("n_encoder_full_layers", 2),
        n_pool=a.get("n_pool", 1024),
        n_pool_cross_layers=a.get("n_encoder_pool_cross_layers", 1),
        n_pool_layers=a.get("n_encoder_pool_layers", 2),
        use_rgb=True,  # stage-2 cache always carries RGB
    )
    den = VertexDenoiser(
        d_model=a.get("d_model", 256),
        n_heads=a.get("n_heads", 8),
        n_layers=a.get("n_layers", 8),
        d_ff=a.get("d_ff", 1024),
        k_verts=a.get("k_verts", 64),
    )
    model = Stage2Diffusion(
        enc, den,
        noise_sigma_xyz=a.get("noise_sigma_xyz", 0.5),
        init_from_scene=a.get("init_from_scene", True),
        scene_init_jitter=a.get("scene_init_jitter", 0.05),
    ).to(device)

    state = ckpt.get("model") or ckpt.get("state_dict")
    if state is None:
        raise RuntimeError(f"checkpoint {ckpt_path} has no 'model' or 'state_dict' key")
    if any(k.startswith("module.") for k in state):
        state = {k[len("module."):]: v for k, v in state.items()}
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing or unexpected:
        print(f"[stage2/model] state_dict mismatch: missing={missing}  unexpected={unexpected}",
              flush=True)
    model.eval()

    print(f"[stage2/model] loaded in {time.perf_counter() - t0:.1f}s  "
          f"d_model={a.get('d_model')} n_layers={a.get('n_layers')}  "
          f"step={ckpt.get('step')}  best_hss={ckpt.get('best_hss')}",
          flush=True)
    return model, a


# ---------------------------------------------------------------------------
# Per-sample preprocessing (CPU, runs in worker threads)
# ---------------------------------------------------------------------------

def preprocess_sample(
    sample: Dict[str, Any],
    n_pts: int,
    use_depth: bool,
    voxel_size_m: float = 0.1,
    keep_raw: bool = False,
) -> Dict[str, Any]:
    """Heavy CPU work: COLMAP parse + V10 scene input build. No torch / GPU here.

    Mirrors `data.dataset.process_sample` exactly: same `voxel_size_m`
    (configs set this to 0.1; build_scene_input's bare default is 0.0
    which would skip voxel downsampling and produce a different point
    set) and same per-order_id deterministic RNG seed (so the same
    points are sampled as were cached at training time).

    With ``keep_raw=True`` the return value is ``{"scene": ..., "raw_sample":
    ..., "colmap_rec": ...}`` so the stage-2 builder can reuse the already-
    decoded COLMAP record instead of parsing the zip twice.
    """
    colmap_rec = read_colmap_rec(sample["colmap"])
    order_id = sample.get("order_id", "")
    rng = np.random.default_rng(_seed_from_order_id(order_id))
    scene = build_scene_input(
        {**sample, "colmap": colmap_rec},
        n_pts=n_pts, use_depth=use_depth,
        voxel_size_m=voxel_size_m,
        rng=rng,
        return_cache=keep_raw,
    )
    if not keep_raw:
        return scene
    cache = scene.pop("_cache", None)
    return {
        "scene":      scene,
        "raw_sample": sample,
        "colmap_rec": colmap_rec,
        "cache":      cache,
    }


def scene_to_batch(scene: Dict[str, np.ndarray], device: torch.device) -> Dict[str, torch.Tensor]:
    batch = {
        "scene_xyz":         torch.from_numpy(scene["scene_xyz"]).unsqueeze(0).to(device),
        "scene_type_ids":    torch.from_numpy(scene["scene_type_ids"]).unsqueeze(0).to(device),
        "scene_gestalt_ids": torch.from_numpy(scene["scene_gestalt_ids"]).unsqueeze(0).to(device),
        "scene_ade_ids":     torch.from_numpy(scene["scene_ade_ids"]).unsqueeze(0).to(device),
        "bbox_center":       torch.from_numpy(scene["bbox_center"]).unsqueeze(0).to(device),
        "bbox_scale":        torch.tensor([scene["bbox_scale"]], device=device),
    }
    # Optional V10 fields — included when produced by build_scene_input.
    for key in ("scene_gestalt_id2", "scene_gestalt_w1",
                "scene_geom_conf", "scene_sem_conf", "scene_rgb", "bbox_R"):
        if key in scene:
            batch[key] = torch.from_numpy(scene[key]).unsqueeze(0).to(device)
    return batch


# ---------------------------------------------------------------------------
# Stage-2 inference helpers
# ---------------------------------------------------------------------------

class Stage2Config:
    """Frozen knobs for the stage-2 refinement path. Built once in `main`."""

    __slots__ = ("model", "n_pts", "k_verts", "n_sample_steps", "validity_thresh",
                 "above_m", "below_m", "side_m",
                 "n_col_oversample", "n_dep_oversample",
                 "n_depth_per_image_cap", "sampling",
                 "density_voxel_size_m", "density_kernel_radius",
                 "density_kernel_axis", "density_response_power",
                 "density_planarity_suppression", "density_planarity_radius",
                 "density_planarity_min_points", "density_min_per_voxel",
                 "ensemble_n", "ensemble_merge_m", "ensemble_mode",
                 "ensemble_strategy", "ensemble_refine_positions",
                 "ensemble_refine_agg", "ensemble_top_k",
                 "ensemble_edge_vote", "ensemble_edge_vote_frac",
                 "ensemble_vertex_vote_frac",
                 "ensemble_stage2_last_k", "ensemble_selector",
                 "ensemble_ranker_path", "ensemble_ranker",
                 "ensemble_fixed_candidate",
                 "ensemble_topk_fuse",
                 "fps_max_exact_iters")

    def __init__(self, model, n_pts: int, k_verts: int, n_sample_steps: int,
                 validity_thresh: float, above_m: float, below_m: float, side_m: float,
                 n_col_oversample: int, n_dep_oversample: int,
                 n_depth_per_image_cap: int, sampling: str = "random",
                 density_voxel_size_m: float = 0.25,
                 density_kernel_radius: int = 3,
                 density_kernel_axis: str = "cube",
                 density_response_power: float = 1.0,
                 density_planarity_suppression: float = 1.0,
                 density_planarity_radius: int = 1,
                 density_planarity_min_points: int = 12,
                 density_min_per_voxel: int = 0,
                 ensemble_n: int = 1,
                 ensemble_merge_m: float = 0.5,
                 ensemble_mode: str = "medoid",
                 ensemble_strategy: str = "union_hull",
                 ensemble_refine_positions: bool = True,
                 ensemble_refine_agg: str = "mean",
                 ensemble_top_k: int = 0,
                 ensemble_edge_vote: bool = True,
                 ensemble_edge_vote_frac: float = 0.5,
                 ensemble_vertex_vote_frac: float = 0.0,
                 ensemble_stage2_last_k: int = 1,
                 ensemble_selector: str = "legacy",
                 ensemble_ranker_path: Optional[str] = None,
                 ensemble_ranker: Optional[Dict[str, np.ndarray]] = None,
                 ensemble_fixed_candidate: str = "confidence_mean_add",
                 ensemble_topk_fuse: int = 1,
                 fps_max_exact_iters: Optional[int] = None):
        self.model = model
        self.n_pts = int(n_pts)
        self.k_verts = int(k_verts)
        self.n_sample_steps = int(n_sample_steps)
        self.validity_thresh = float(validity_thresh)
        self.above_m = float(above_m); self.below_m = float(below_m); self.side_m = float(side_m)
        self.n_col_oversample = int(n_col_oversample)
        self.n_dep_oversample = int(n_dep_oversample)
        self.n_depth_per_image_cap = int(n_depth_per_image_cap)
        if sampling not in ("random", "fps", "voxel"):
            raise ValueError(f"sampling must be 'random', 'fps', or 'voxel', got {sampling!r}")
        self.sampling = sampling
        self.density_voxel_size_m = float(density_voxel_size_m)
        self.density_kernel_radius = int(density_kernel_radius)
        self.density_kernel_axis = str(density_kernel_axis)
        self.density_response_power = float(density_response_power)
        self.density_planarity_suppression = float(density_planarity_suppression)
        self.density_planarity_radius = int(density_planarity_radius)
        self.density_planarity_min_points = int(density_planarity_min_points)
        self.density_min_per_voxel = int(density_min_per_voxel)
        self.ensemble_n = max(1, int(ensemble_n))
        self.ensemble_merge_m = float(ensemble_merge_m)
        if ensemble_mode not in ("medoid", "consensus", "confidence"):
            raise ValueError(f"ensemble_mode must be 'medoid', 'consensus', or 'confidence', got {ensemble_mode!r}")
        self.ensemble_mode = str(ensemble_mode)
        if ensemble_strategy not in ("union_hull", "per_seed"):
            raise ValueError(f"ensemble_strategy must be 'union_hull' or 'per_seed', got {ensemble_strategy!r}")
        self.ensemble_strategy = str(ensemble_strategy)
        self.ensemble_refine_positions = bool(ensemble_refine_positions)
        if ensemble_refine_agg not in ("median", "mean", "wmean"):
            raise ValueError(
                f"ensemble_refine_agg must be 'median', 'mean', or 'wmean', "
                f"got {ensemble_refine_agg!r}")
        self.ensemble_refine_agg = str(ensemble_refine_agg)
        self.ensemble_top_k = max(0, int(ensemble_top_k))
        self.ensemble_edge_vote = bool(ensemble_edge_vote)
        if not (0.0 < float(ensemble_edge_vote_frac) <= 1.0):
            raise ValueError(
                f"ensemble_edge_vote_frac must be in (0, 1], got {ensemble_edge_vote_frac!r}")
        self.ensemble_edge_vote_frac = float(ensemble_edge_vote_frac)
        if not (0.0 <= float(ensemble_vertex_vote_frac) <= 1.0):
            raise ValueError(
                f"ensemble_vertex_vote_frac must be in [0, 1], "
                f"got {ensemble_vertex_vote_frac!r}")
        self.ensemble_vertex_vote_frac = float(ensemble_vertex_vote_frac)
        self.ensemble_stage2_last_k = max(1, int(ensemble_stage2_last_k))
        if ensemble_selector not in ("legacy", "ranker", "fixed"):
            raise ValueError(
                f"ensemble_selector must be 'legacy', 'ranker', or 'fixed', got {ensemble_selector!r}")
        self.ensemble_selector = str(ensemble_selector)
        self.ensemble_ranker_path = ensemble_ranker_path
        self.ensemble_ranker = ensemble_ranker
        self.ensemble_fixed_candidate = str(ensemble_fixed_candidate)
        self.ensemble_topk_fuse = max(1, int(ensemble_topk_fuse))
        # None => full exact FPS. Set to the value the model trained with so
        # the per-provenance anchor count matches; otherwise the encoder sees
        # a markedly different point distribution at inference (all exact
        # anchors vs. exact prefix + deterministic random fill).
        self.fps_max_exact_iters = (None if fps_max_exact_iters is None
                                    else int(fps_max_exact_iters))


def _stage2_sample_to_batch(sample: Dict[str, Any], device: torch.device) -> Dict[str, torch.Tensor]:
    """Same shape as ``scene_to_batch`` but for the stage-2 sample dict produced
    by ``stage2_row_to_sample`` (includes ``init_verts`` + valid mask)."""
    keys = (
        "scene_xyz", "scene_type_ids",
        "scene_gestalt_ids", "scene_gestalt_id2", "scene_gestalt_w1",
        "scene_ade_ids", "scene_geom_conf", "scene_sem_conf",
        "scene_rgb", "init_verts", "init_verts_valid",
        "bbox_center", "bbox_scale", "bbox_R",
        # FPS-mode extras (only present when pre_subsample=False):
        "scene_valid_mask", "init_verts_world", "verts_gt_world",
    )
    return {k: sample[k].unsqueeze(0).to(device) for k in keys if k in sample}


def _replicate_batch(batch: Dict[str, torch.Tensor], n: int) -> Dict[str, torch.Tensor]:
    """Return a shallow-copied batch with every leading-dim-1 tensor repeated
    to leading dim `n`. Non-tensor / non-(1, ...) entries pass through. Used to
    convert a B=1 single-scene batch into a B=N replicated batch so the
    diffusion sampler produces N independent trajectories sharing one scene."""
    if n <= 1:
        return batch
    out: Dict[str, torch.Tensor] = {}
    for k, v in batch.items():
        if isinstance(v, torch.Tensor) and v.dim() >= 1 and v.shape[0] == 1:
            out[k] = v.repeat(n, *([1] * (v.dim() - 1)))
        else:
            out[k] = v
    return out


def _stage1_seed_ensemble(
    model: WireframeDiffusion,
    batch1: Dict[str, torch.Tensor],
    n_steps: int,
    validity_thresh: float,
    n_seeds: int,
) -> List[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
    """Run stage-1 sampling `n_seeds` times with independent random init x0.

    Mirrors the body of ``WireframeDiffusion.sample`` so we can extract the
    final-step validity logits per slot (needed for confidence-mode ensemble
    selection); ``model.sample`` only returns a boolean mask and drops the
    raw logits.

    Returns a list of ``(verts_world (M, 3) float32, edges (E, 2) int64,
    valid_logits (M,) float32)`` tuples, one per seed. ``valid_logits``
    contains the raw logit value for each kept vertex (post-threshold) so
    callers can compute mean/max self-confidence."""
    n_seeds = max(1, int(n_seeds))
    batch_n = _replicate_batch(batch1, n_seeds) if n_seeds > 1 else batch1
    device = batch_n["scene_xyz"].device

    # Inline equivalent of model.sample(), but keeping last_logit available.
    scene_feats, scene_xyz, query_xyz = model._encode_scene_with_query(batch_n)
    K = model.denoiser.k_verts
    B = scene_feats.shape[0]
    dt = 1.0 / max(1, n_steps)
    x = model._init_x0(batch_n, K, device, B, query_xyz=query_xyz)
    last_logit = torch.zeros(B, K, device=device)
    last_edge_logit = torch.zeros(B, K, K, device=device)
    for i in range(n_steps):
        t = torch.full((B,), i * dt, device=device)
        v, last_logit, last_edge_logit = model.denoiser(
            x, t, scene_feats, scene_xyz,
        )
        v = v.float()
        last_logit = last_logit.float()
        last_edge_logit = last_edge_logit.float()
        x = x + dt * v

    out: List[Tuple[np.ndarray, np.ndarray, np.ndarray]] = []
    bbox_R = batch_n.get("bbox_R")
    for i in range(B):
        x_i = torch.nan_to_num(x[i].float(),
                               nan=0.0, posinf=model.xyz_clip,
                               neginf=-model.xyz_clip
                               ).clamp(-model.xyz_clip, model.xyz_clip)
        logit_i = last_logit[i]
        idx = torch.nonzero(logit_i > validity_thresh, as_tuple=False).flatten()
        x_valid = x_i.index_select(0, idx)
        center = batch_n["bbox_center"][i].to(device=x_valid.device, dtype=torch.float32)
        scale = batch_n["bbox_scale"][i].to(device=x_valid.device, dtype=torch.float32)
        if isinstance(bbox_R, torch.Tensor):
            R = bbox_R[i].to(device=x_valid.device, dtype=torch.float32)
            verts_world = (x_valid * scale) @ R + center
        else:
            verts_world = x_valid * scale + center
        verts_np = verts_world.float().cpu().numpy().reshape(-1, 3).astype(np.float32)
        logits_kept = logit_i.index_select(0, idx).detach().cpu().numpy().astype(np.float32)

        n_valid = int(idx.numel())
        edges_arr = np.zeros((0, 2), dtype=np.int64)
        if n_valid >= 2:
            sub = last_edge_logit[i].float().index_select(0, idx).index_select(1, idx)
            tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device,
                                        dtype=torch.bool), diagonal=1)
            pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False)
            if pairs.numel() > 0:
                edges_arr = pairs.detach().cpu().numpy().astype(np.int64)
        out.append((verts_np, edges_arr, logits_kept))
    return out


def _fuse_wireframes(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    n_total: int,
    merge_tau_m: float,
    vertex_vote_frac: float = 0.5,
    edge_vote_frac: float = 0.5,
    aggregator: str = "mean",
) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
    """Vertex-merge consensus fusion.

    1. Greedy single-link cluster all run vertices in world coords (threshold
       ``merge_tau_m`` metres). Each cluster centroid is the running mean of
       its members (or a logit-weighted mean when ``aggregator="wmean"``).
    2. Keep clusters supported by at least
       ``ceil(n_total * vertex_vote_frac)`` distinct runs.
    3. For each pair of kept clusters, vote an edge if it appears in at least
       ``ceil(n_total * edge_vote_frac)`` runs.

    Falls back to the single run with the most predicted vertices when no
    cluster meets the support bar — we never collapse to empty when individual
    runs were non-empty. ``runs`` entries are ``(verts, edges, logits)``;
    ``logits`` is only consulted under ``aggregator="wmean"``."""
    if aggregator not in ("mean", "median", "wmean"):
        raise ValueError(
            f"aggregator must be 'mean', 'median', or 'wmean', got {aggregator!r}")
    valid: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = []
    for r in runs:
        v_raw = r[0]
        if v_raw is None or len(v_raw) == 0:
            continue
        v = np.asarray(v_raw, dtype=np.float64).reshape(-1, 3)
        e = list(r[1]) if len(r) > 1 and r[1] is not None else []
        lg_raw = r[2] if len(r) > 2 else None
        lg = (np.asarray(lg_raw, dtype=np.float64).reshape(-1)
              if lg_raw is not None else np.zeros(v.shape[0], dtype=np.float64))
        if lg.shape[0] != v.shape[0]:
            lg = np.zeros(v.shape[0], dtype=np.float64)
        valid.append((v, e, lg))
    if not valid:
        return np.zeros((0, 3), dtype=np.float64), []

    v_support = max(1, int(np.ceil(n_total * float(vertex_vote_frac))))
    e_support = max(1, int(np.ceil(n_total * float(edge_vote_frac))))

    # Per-cluster bookkeeping: member positions, member logits, and the set
    # of run indices contributing. Centroid is computed once at the end so
    # 'median' / 'wmean' can use all members (running-mean only worked for
    # the plain 'mean' aggregator).
    cluster_members: List[List[np.ndarray]] = []
    cluster_member_lg: List[List[float]] = []
    cluster_centroids: List[np.ndarray] = []  # running mean — used for matching
    cluster_count: List[int] = []
    cluster_runs: List[set] = []
    vertex_to_cluster: List[List[int]] = []

    for r, (verts, _edges, lg) in enumerate(valid):
        local_map: List[int] = []
        for vi in range(verts.shape[0]):
            v = verts[vi]
            best_c, best_d = -1, float("inf")
            for ci, cc in enumerate(cluster_centroids):
                d = float(np.linalg.norm(v - cc))
                if d < best_d:
                    best_d, best_c = d, ci
            if best_c >= 0 and best_d <= merge_tau_m:
                n = cluster_count[best_c]
                cluster_centroids[best_c] = (cluster_centroids[best_c] * n + v) / (n + 1)
                cluster_count[best_c] = n + 1
                cluster_runs[best_c].add(r)
                cluster_members[best_c].append(v)
                cluster_member_lg[best_c].append(float(lg[vi]))
                local_map.append(best_c)
            else:
                cluster_centroids.append(v.copy())
                cluster_count.append(1)
                cluster_runs.append({r})
                cluster_members.append([v])
                cluster_member_lg.append([float(lg[vi])])
                local_map.append(len(cluster_centroids) - 1)
        vertex_to_cluster.append(local_map)

    kept = [ci for ci, rs in enumerate(cluster_runs) if len(rs) >= v_support]
    if not kept:
        best = max(range(len(valid)), key=lambda r: valid[r][0].shape[0])
        return valid[best][0].astype(np.float64), list(valid[best][1])

    # Final centroids per kept cluster, per aggregator.
    fused_verts = np.empty((len(kept), 3), dtype=np.float64)
    for new_i, ci in enumerate(kept):
        stack = np.stack(cluster_members[ci], axis=0)
        if aggregator == "median":
            fused_verts[new_i] = np.median(stack, axis=0)
        elif aggregator == "wmean":
            lg = np.asarray(cluster_member_lg[ci], dtype=np.float64)
            lg = lg - lg.max()
            w = np.exp(lg)
            ws = float(w.sum())
            if ws <= 1e-12 or not np.isfinite(ws):
                fused_verts[new_i] = stack.mean(axis=0)
            else:
                fused_verts[new_i] = (stack * (w / ws)[:, None]).sum(axis=0)
        else:  # "mean"
            fused_verts[new_i] = stack.mean(axis=0)

    old_to_new = {old: new for new, old in enumerate(kept)}
    n_kept = len(kept)
    edge_votes = np.zeros((n_kept, n_kept), dtype=np.int32)
    for r, (_verts, edges, _lg) in enumerate(valid):
        seen_pairs: set = set()
        for (a, b) in edges:
            if a < 0 or b < 0 or a == b:
                continue
            if a >= len(vertex_to_cluster[r]) or b >= len(vertex_to_cluster[r]):
                continue
            ca = vertex_to_cluster[r][a]
            cb = vertex_to_cluster[r][b]
            if ca == cb or ca not in old_to_new or cb not in old_to_new:
                continue
            i, j = old_to_new[ca], old_to_new[cb]
            if i > j:
                i, j = j, i
            if (i, j) in seen_pairs:
                continue
            seen_pairs.add((i, j))
            edge_votes[i, j] += 1
    ii, jj = np.where(edge_votes >= e_support)
    fused_edges = [(int(a), int(b)) for a, b in zip(ii.tolist(), jj.tolist())]
    return fused_verts, fused_edges


def _wireframe_pair_distance(
    va: np.ndarray, ea: List[Tuple[int, int]],
    vb: np.ndarray, eb: List[Tuple[int, int]],
    match_tau_m: float = 0.4,
) -> float:
    """Symmetric distance between two wireframes for medoid selection.

    Uses Hungarian (optimal 1-1) vertex assignment, which is naturally
    symmetric and aligned with how the HSS metric matches predictions to GT.

    Components:
      * **Vertex term** (metres, capped at ``match_tau_m``): solve the
        rectangular linear-sum-assignment on the pairwise distance matrix
        capped at ``match_tau_m``. Matched pair cost = capped distance;
        unmatched leftover vertices (size mismatch) cost ``match_tau_m`` each.
        Mean over ``max(|Va|, |Vb|)`` slots so the term lives in
        ``[0, match_tau_m]``.
      * **Edge term** (Jaccard, in ``[0, 1]``): the Hungarian assignment
        gives a 1-1 vertex map; use it to relabel B's edges into A's index
        space (only for matches actually within ``match_tau_m``) and compute
        ``1 - IoU`` over the two edge sets.

    Returns a single scalar; lower = more similar. Degenerate empty inputs
    return ``+inf`` so they cannot be picked as the medoid.
    """
    from scipy.optimize import linear_sum_assignment

    va = np.asarray(va, dtype=np.float64).reshape(-1, 3)
    vb = np.asarray(vb, dtype=np.float64).reshape(-1, 3)
    Na, Nb = va.shape[0], vb.shape[0]
    if Na == 0 or Nb == 0:
        return float("inf")

    D = np.linalg.norm(va[:, None, :] - vb[None, :, :], axis=-1)
    D_capped = np.minimum(D, match_tau_m)

    row_ind, col_ind = linear_sum_assignment(D_capped)
    matched_cost = float(D_capped[row_ind, col_ind].sum())
    n_unmatched = abs(Na - Nb)  # leftover on the larger side
    vertex_dist = (matched_cost + n_unmatched * match_tau_m) / max(Na, Nb)

    # Map B → A via the Hungarian assignment, but only when the match is
    # actually within tau (capped matches at exactly tau are unreliable).
    b_to_a = -np.ones(Nb, dtype=np.int64)
    for r, c in zip(row_ind, col_ind):
        if D[r, c] <= match_tau_m:
            b_to_a[c] = r

    ea_norm = set((min(int(i), int(j)), max(int(i), int(j)))
                  for i, j in ea if i != j)
    eb_mapped = set()
    for (i, j) in eb:
        if i == j or i < 0 or j < 0 or i >= Nb or j >= Nb:
            continue
        ai, aj = int(b_to_a[i]), int(b_to_a[j])
        if ai < 0 or aj < 0 or ai == aj:
            continue
        eb_mapped.add((min(ai, aj), max(ai, aj)))

    if not ea_norm and not eb_mapped:
        edge_dist = 0.0
    elif not ea_norm or not eb_mapped:
        edge_dist = 1.0
    else:
        inter = len(ea_norm & eb_mapped)
        union = len(ea_norm | eb_mapped)
        edge_dist = 1.0 - (inter / union)

    return vertex_dist + edge_dist


def _medoid_rank(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    match_tau_m: float = 0.4,
) -> List[int]:
    """Return indices into ``runs`` sorted by centrality (lowest pairwise
    Hungarian distance first → most-central run is index ``[0]``). Empty /
    edgeless runs are excluded from the pairwise computation; if none qualify
    the singleton list ``[longest_run_index]`` is returned. Always returns
    at least one valid index when ``runs`` is non-empty."""
    if not runs:
        return []
    candidates_idx = [r for r, (v, e, *_) in enumerate(runs)
                      if len(v) > 0 and len(e) > 0]
    if not candidates_idx:
        return [max(range(len(runs)),
                    key=lambda r: (len(runs[r][0]) if runs[r][0] is not None else -1,
                                   len(runs[r][1]) if runs[r][1] is not None else -1))]
    if len(candidates_idx) == 1:
        return [candidates_idx[0]]

    n = len(candidates_idx)
    verts = [np.asarray(runs[idx][0], dtype=np.float64).reshape(-1, 3)
             for idx in candidates_idx]
    edges = [list(runs[idx][1]) for idx in candidates_idx]
    dist = np.zeros((n, n), dtype=np.float64)
    for i in range(n):
        for j in range(i + 1, n):
            d = _wireframe_pair_distance(
                verts[i], edges[i], verts[j], edges[j],
                match_tau_m=match_tau_m,
            )
            dist[i, j] = d
            dist[j, i] = d
    centrality = dist.sum(axis=1)
    order = np.argsort(centrality, kind="stable")
    return [candidates_idx[int(k)] for k in order]


def _medoid_select_idx(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    match_tau_m: float = 0.4,
) -> int:
    """Return the index into ``runs`` of the medoid — the run whose wireframe
    is most similar to the others by Hungarian-matched distance. Thin wrapper
    around ``_medoid_rank`` that returns just the most-central index."""
    ranking = _medoid_rank(runs, match_tau_m=match_tau_m)
    return ranking[0] if ranking else 0


def _confidence_select_idx(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
) -> int:
    """Return the index into ``runs`` of the highest mean-validity-logit run.
    Skips empty / all-NaN runs in the scoring; falls back to the longest run
    when no run has meaningful confidence."""
    if not runs:
        return 0
    scored: List[Tuple[float, int]] = []
    for i, (v, _e, logits) in enumerate(runs):
        if v is None or len(v) == 0 or logits is None or len(logits) == 0:
            continue
        arr = np.asarray(logits, dtype=np.float64)
        if not np.all(np.isfinite(arr)):
            arr = arr[np.isfinite(arr)]
            if arr.size == 0:
                continue
        scored.append((float(arr.mean()), i))
    if not scored:
        return max(range(len(runs)),
                   key=lambda r: (len(runs[r][0]) if runs[r][0] is not None else -1,
                                  len(runs[r][1]) if runs[r][1] is not None else -1))
    return max(scored, key=lambda t: t[0])[1]


def _refine_positions(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    picked_idx: int,
    match_tau_m: float,
    aggregator: str = "mean",
    member_indices: Optional[Sequence[int]] = None,
    vertex_vote_frac: float = 0.0,
) -> Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]:
    """Position refinement on top of a single picked run.

    Keep the picked run's vertices and edges (topology unchanged). For each
    vertex ``v`` of the picked run, find its Hungarian-matched correspondent
    in every other run (only matches within ``match_tau_m`` count), then
    replace ``v`` with the per-axis ``aggregator`` of ``{v}`` ∪
    correspondents.

    ``aggregator`` ∈ {"mean", "median", "wmean"}:
      * ``"mean"`` (default) — simple average; lower variance under iid noise
        but sensitive to outliers.
      * ``"median"`` — per-axis median. Robust to a single bad correspondent
        (e.g., an off-by-one Hungarian match at the τ boundary).
      * ``"wmean"`` — softmax-of-validity-logit weighted mean.

    Rationale: vertex prediction is the dominant noise source (F1 hits HSS
    harder than IoU). Selectors (medoid / confidence) pick a *vertex set*
    that's already good; this step denoises each vertex's *position* by
    aggregating with consensus detections from the other N−1 runs. Topology
    (edges) is untouched so we never lose recall on edges.

    Returns ``(verts_refined, edges_picked)``. Falls back to the picked
    run's verts/edges verbatim when no refinement signal is available."""
    from scipy.optimize import linear_sum_assignment

    if aggregator not in ("median", "mean", "wmean"):
        raise ValueError(
            f"aggregator must be 'median', 'mean', or 'wmean', got {aggregator!r}")
    if not runs:
        return np.zeros((0, 3), dtype=np.float64), [], np.zeros((0,), dtype=np.int64)
    picked_v_raw, picked_e_raw, picked_logits = runs[picked_idx]
    picked_v = np.asarray(picked_v_raw, dtype=np.float64).reshape(-1, 3)
    picked_e = list(picked_e_raw)
    if picked_v.shape[0] == 0:
        return picked_v, picked_e, np.zeros((0,), dtype=np.int64)

    M = picked_v.shape[0]
    picked_lg = (np.asarray(picked_logits, dtype=np.float64).reshape(-1)
                 if picked_logits is not None else np.zeros(M, dtype=np.float64))
    if picked_lg.shape[0] != M:
        picked_lg = np.zeros(M, dtype=np.float64)

    # Per-picked-vertex correspondent positions and per-correspondent logits
    # (used for 'wmean' aggregation). Picked run is always the first entry.
    correspondents: List[List[np.ndarray]] = [[picked_v[i]] for i in range(M)]
    corr_logits: List[List[float]] = [[float(picked_lg[i])] for i in range(M)]

    if member_indices is None:
        contributor_iter: List[int] = list(range(len(runs)))
    else:
        contributor_iter = list(member_indices)
        if picked_idx not in contributor_iter:
            contributor_iter.append(picked_idx)
    # Voter set includes picked (self-vote on each vertex).
    M_voters = max(1, len(set(contributor_iter)))

    for j in contributor_iter:
        if j == picked_idx or j < 0 or j >= len(runs):
            continue
        vj_raw, _ej, lj_raw = runs[j]
        vj = np.asarray(vj_raw, dtype=np.float64).reshape(-1, 3)
        if vj.shape[0] == 0:
            continue
        lj = (np.asarray(lj_raw, dtype=np.float64).reshape(-1)
              if lj_raw is not None else np.zeros(vj.shape[0], dtype=np.float64))
        if lj.shape[0] != vj.shape[0]:
            lj = np.zeros(vj.shape[0], dtype=np.float64)
        D = np.linalg.norm(picked_v[:, None, :] - vj[None, :, :], axis=-1)
        D_capped = np.minimum(D, match_tau_m)
        row_ind, col_ind = linear_sum_assignment(D_capped)
        for r, c in zip(row_ind, col_ind):
            if D[r, c] <= match_tau_m:
                correspondents[r].append(vj[c])
                corr_logits[r].append(float(lj[c]))

    # Vertex consensus drop. A picked vertex survives iff at least
    # ceil(M_voters * vertex_vote_frac) voters support it (including the
    # picked run's self-vote).
    if vertex_vote_frac > 0.0:
        threshold = int(np.ceil(M_voters * float(vertex_vote_frac)))
        keep_mask = np.array(
            [len(correspondents[i]) >= threshold for i in range(M)],
            dtype=bool,
        )
        if not keep_mask.any():
            # Don't collapse to empty — fall back to keeping all picked verts.
            keep_mask[:] = True
    else:
        keep_mask = np.ones(M, dtype=bool)

    keep_idx = np.flatnonzero(keep_mask).astype(np.int64)

    refined = np.empty((keep_idx.size, 3), dtype=np.float64)
    for new_i, old_i in enumerate(keep_idx):
        stack = np.stack(correspondents[int(old_i)], axis=0)
        if aggregator == "median":
            refined[new_i] = np.median(stack, axis=0)
        elif aggregator == "mean":
            refined[new_i] = stack.mean(axis=0)
        else:  # "wmean"
            lg = np.asarray(corr_logits[int(old_i)], dtype=np.float64)
            # Softmax over logits for stable positive weights; falls back to
            # uniform when logits are all-equal / missing.
            lg = lg - lg.max()
            w = np.exp(lg)
            w_sum = float(w.sum())
            if w_sum <= 1e-12 or not np.isfinite(w_sum):
                refined[new_i] = stack.mean(axis=0)
            else:
                w = w / w_sum
                refined[new_i] = (stack * w[:, None]).sum(axis=0)

    # Remap picked-run edges into the kept-vertex index space; drop edges
    # whose endpoints didn't survive the vertex vote.
    old_to_new = -np.ones(M, dtype=np.int64)
    for new_i, old_i in enumerate(keep_idx):
        old_to_new[int(old_i)] = new_i
    edges_remapped: List[Tuple[int, int]] = []
    for (a, b) in picked_e:
        ia, ib = int(a), int(b)
        if ia < 0 or ib < 0 or ia >= M or ib >= M or ia == ib:
            continue
        na, nb = int(old_to_new[ia]), int(old_to_new[ib])
        if na < 0 or nb < 0 or na == nb:
            continue
        edges_remapped.append((na, nb))

    return refined, edges_remapped, keep_idx


def _vote_edges(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    picked_idx: int,
    match_tau_m: float,
    vote_frac: float = 0.5,
    member_indices: Optional[Sequence[int]] = None,
    keep_idx: Optional[np.ndarray] = None,
) -> List[Tuple[int, int]]:
    """Majority-vote edge refinement.

    For each voter run, Hungarian-match its vertices to the picked run's
    vertices (matches outside ``match_tau_m`` ignored), relabel its edges
    into picked-vertex index space, and tally votes per ``(i, j)`` pair.
    An edge survives if at least ``ceil(M * vote_frac)`` voters (out of
    ``M`` = ``len(member_indices)`` or ``len(runs)``) vote for it. The
    picked run itself is always a voter — its own edges therefore enter
    with one self-vote.

    Returns the filtered edge list in picked-vertex index space.
    """
    from scipy.optimize import linear_sum_assignment

    if not runs:
        return []
    picked_v = np.asarray(runs[picked_idx][0], dtype=np.float64).reshape(-1, 3)
    if picked_v.shape[0] == 0:
        return list(runs[picked_idx][1])

    if member_indices is None:
        voters: List[int] = list(range(len(runs)))
    else:
        voters = [int(j) for j in member_indices if 0 <= int(j) < len(runs)]
        if picked_idx not in voters:
            voters.append(int(picked_idx))
    M = max(1, len(voters))
    threshold = int(np.ceil(M * float(vote_frac)))

    votes: Dict[Tuple[int, int], int] = {}
    Mp = picked_v.shape[0]
    for j in voters:
        vj_raw, ej_raw, _lj = runs[j]
        vj = np.asarray(vj_raw, dtype=np.float64).reshape(-1, 3)
        if j == picked_idx:
            j_to_p = np.arange(Mp, dtype=np.int64)
        else:
            if vj.shape[0] == 0:
                continue
            D = np.linalg.norm(picked_v[:, None, :] - vj[None, :, :], axis=-1)
            D_capped = np.minimum(D, match_tau_m)
            row_ind, col_ind = linear_sum_assignment(D_capped)
            j_to_p = -np.ones(vj.shape[0], dtype=np.int64)
            for r, c in zip(row_ind, col_ind):
                if D[r, c] <= match_tau_m:
                    j_to_p[c] = r

        for (a, b) in ej_raw:
            ia, ib = int(a), int(b)
            if ia == ib or ia < 0 or ib < 0 or ia >= len(j_to_p) or ib >= len(j_to_p):
                continue
            pa, pb = int(j_to_p[ia]), int(j_to_p[ib])
            if pa < 0 or pb < 0 or pa == pb:
                continue
            key = (min(pa, pb), max(pa, pb))
            votes[key] = votes.get(key, 0) + 1

    kept_old = [(int(i), int(j)) for (i, j), c in votes.items() if c >= threshold]
    if not kept_old:
        # Degenerate vote tally (e.g. only the picked run voted on edges with
        # vote_frac > 1/M). Fall back to picked edges so we never lose all
        # edges — F1 would zero out HSS.
        kept_old = [(int(a), int(b)) for (a, b) in runs[picked_idx][1]
                    if int(a) != int(b)]

    if keep_idx is None:
        return kept_old

    # Remap from picked-vertex index space → kept-vertex index space (the
    # output of _refine_positions when vertex_vote_frac > 0 drops some
    # picked vertices). Edges referencing dropped vertices are filtered out.
    old_to_new = -np.ones(Mp, dtype=np.int64)
    for new_i, old_i in enumerate(np.asarray(keep_idx).reshape(-1)):
        oi = int(old_i)
        if 0 <= oi < Mp:
            old_to_new[oi] = new_i
    out: List[Tuple[int, int]] = []
    for (a, b) in kept_old:
        na, nb = int(old_to_new[a]), int(old_to_new[b])
        if na < 0 or nb < 0 or na == nb:
            continue
        out.append((min(na, nb), max(na, nb)))
    return out


# ------------------------------------------------------------------
# 2D gestalt-mask reprojection features
# ------------------------------------------------------------------
# Pixels of these gestalt classes mark wireframe edges (lines).
_GESTALT_EDGE_NAMES = (
    "ridge", "eave", "rake", "hip", "valley",
    "fascia", "transition_line", "flashing", "step_flashing",
    "ground_line", "soffit",
)
# Pixels of these classes mark wireframe vertices (line endpoints).
_GESTALT_VERT_NAMES = (
    "apex", "eave_end_point", "flashing_end_point",
)


def _ids_from_class_names(names: Tuple[str, ...]) -> np.ndarray:
    return np.array(
        sorted(GESTALT_CLASSES.index(n) for n in names if n in GESTALT_CLASSES),
        dtype=np.int64,
    )


_EDGE_GESTALT_IDS = _ids_from_class_names(_GESTALT_EDGE_NAMES)
_VERT_GESTALT_IDS = _ids_from_class_names(_GESTALT_VERT_NAMES)


def _build_reprojection_views(
    raw_sample: Dict[str, Any],
    colmap_rec,
    edge_dilate: int = 2,
    vert_dilate: int = 4,
) -> List[Dict[str, Any]]:
    """Pre-compute per-view camera + dilated gestalt masks for reprojection.

    Returns a list of dicts with keys
      ``R`` (3,3), ``t`` (3,), ``fx``, ``fy``, ``cx``, ``cy``, ``W``, ``H``,
      ``edge_mask`` and ``vert_mask`` (both uint8 HxW).

    The masks are 1 where pixels belong to wireframe edge / vertex gestalt
    classes (dilated by ``edge_dilate`` / ``vert_dilate`` 3x3 iterations to
    tolerate small projection error and label boundary noise).
    """
    try:
        from scipy.ndimage import binary_dilation
    except Exception:
        binary_dilation = None  # type: ignore

    image_ids = raw_sample.get("image_ids", []) or []
    gestalt_imgs = raw_sample.get("gestalt", []) or []
    if not image_ids or not gestalt_imgs:
        return []

    colmap_by_name = {
        col_img.name: col_img for col_img in colmap_rec.images.values()
    }
    structure = np.ones((3, 3), dtype=bool)

    out: List[Dict[str, Any]] = []
    for i, img_id in enumerate(image_ids):
        if i >= len(gestalt_imgs) or gestalt_imgs[i] is None:
            continue
        col_img = _resolve_colmap_image(colmap_by_name, img_id)
        if col_img is None:
            continue
        cam = _camera_for_image(colmap_rec, col_img)
        cam_w = int(getattr(cam, "width", 0) or 0)
        cam_h = int(getattr(cam, "height", 0) or 0)
        if cam_w <= 0 or cam_h <= 0:
            continue
        try:
            R, t = _cam_matrix_from_image(col_img)
            R = np.asarray(R, dtype=np.float64).reshape(3, 3)
            t = np.asarray(t, dtype=np.float64).reshape(3)
            K = cam.calibration_matrix()
        except Exception:
            continue
        try:
            gest_np = _image_to_rgb_array(gestalt_imgs[i])
        except Exception:
            continue
        ids = gestalt_img_to_ids(gest_np)
        edge_tight = np.isin(ids, _EDGE_GESTALT_IDS)
        vert_tight = np.isin(ids, _VERT_GESTALT_IDS)
        edge_mask = edge_tight.copy()
        vert_mask = vert_tight.copy()
        if binary_dilation is not None:
            if edge_dilate > 0 and edge_mask.any():
                edge_mask = binary_dilation(
                    edge_mask, structure=structure, iterations=int(edge_dilate))
            if vert_dilate > 0 and vert_mask.any():
                vert_mask = binary_dilation(
                    vert_mask, structure=structure, iterations=int(vert_dilate))
        H, W = ids.shape[:2]
        out.append({
            "R": R, "t": t,
            "fx": float(K[0, 0]), "fy": float(K[1, 1]),
            "cx": float(K[0, 2]), "cy": float(K[1, 2]),
            "cam_w": float(cam_w), "cam_h": float(cam_h),
            "W": int(W), "H": int(H),
            "edge_mask": np.ascontiguousarray(edge_mask, dtype=bool),
            "vert_mask": np.ascontiguousarray(vert_mask, dtype=bool),
            "edge_mask_tight": np.ascontiguousarray(edge_tight, dtype=bool),
            "vert_mask_tight": np.ascontiguousarray(vert_tight, dtype=bool),
            "ids": np.ascontiguousarray(ids, dtype=np.int32),
        })
    return out


def _project_to_mask(
    pts: np.ndarray,
    view: Dict[str, Any],
    which: str,
) -> Tuple[np.ndarray, np.ndarray]:
    """Project (N,3) world points into one view's mask.

    Returns (in_view, hit) bool arrays of length N. ``in_view`` is True when
    the point is in front of the camera and lands inside the image; ``hit``
    is True when in_view AND the mask is set at the corresponding pixel.
    """
    R = view["R"]; t = view["t"]
    p_cam = pts @ R.T + t
    z = p_cam[:, 2]
    valid = z > 1e-6
    u = np.empty_like(z); v = np.empty_like(z)
    zsafe = np.where(valid, z, 1.0)
    u_proj = p_cam[:, 0] / zsafe * view["fx"] + view["cx"]
    v_proj = p_cam[:, 1] / zsafe * view["fy"] + view["cy"]
    sx = view["W"] / view["cam_w"]; sy = view["H"] / view["cam_h"]
    ui = np.rint(u_proj * sx).astype(np.int64)
    vi = np.rint(v_proj * sy).astype(np.int64)
    in_view = valid & (ui >= 0) & (ui < view["W"]) & (vi >= 0) & (vi < view["H"])
    mask = view[which]
    hit = np.zeros(pts.shape[0], dtype=bool)
    if in_view.any():
        idx = np.where(in_view)[0]
        hit[idx] = mask[vi[idx], ui[idx]]
    return in_view, hit


def _reprojection_features(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    views: List[Dict[str, Any]],
    n_edge_samples: int = 8,
) -> Tuple[float, float, float, float]:
    """Return ``(vert_mean, vert_min, edge_mean, edge_min)`` — legacy stub.

    Mostly kept for backwards-compatibility; new code should call
    :func:`_reprojection_features_extended` for the richer per-vertex /
    per-edge breakdown.
    """
    res = _reprojection_features_extended(verts, edges, views, n_edge_samples)
    return res["vert_mean"], res["vert_min"], res["edge_mean"], res["edge_min"]


def _reprojection_features_extended(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    views: List[Dict[str, Any]],
    n_edge_samples: int = 8,
) -> Dict[str, Any]:
    """Full reprojection scoring with per-vertex / per-edge arrays.

    Returns a dict with scalar fields:
      vert_mean, vert_min, vert_worst3_mean, vert_bad_frac,
      vert_soft_mean,
      edge_mean, edge_min, edge_worst3_mean, edge_bad_frac,
      edge_soft_mean,
      endpoint_corner_mean, endpoint_corner_min,
    plus arrays:
      per_v (Nv,), per_e (Ne,)
    All scalars in [0, 1]. "soft" uses the non-dilated (tight) masks;
    "bad" thresholds at 0.1; "worst3" averages the 3 lowest scores.
    """
    out = {
        "vert_mean": 0.0, "vert_min": 0.0,
        "vert_worst3_mean": 0.0, "vert_bad_frac": 0.0,
        "vert_soft_mean": 0.0,
        "edge_mean": 0.0, "edge_min": 0.0,
        "edge_worst3_mean": 0.0, "edge_bad_frac": 0.0,
        "edge_soft_mean": 0.0,
        "endpoint_corner_mean": 0.0, "endpoint_corner_min": 0.0,
        "per_v": np.zeros((0,), dtype=np.float64),
        "per_e": np.zeros((0,), dtype=np.float64),
    }
    if not views or verts.size == 0:
        return out
    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    Nv = v.shape[0]

    # --- vertices: dilated + tight masks ---
    vert_hits = np.zeros(Nv, dtype=np.float64)
    vert_soft = np.zeros(Nv, dtype=np.float64)
    vert_views = np.zeros(Nv, dtype=np.float64)
    for view in views:
        in_view, hit = _project_to_mask(v, view, "vert_mask")
        _, hit_soft = _project_to_mask(v, view, "vert_mask_tight")
        vert_views += in_view.astype(np.float64)
        vert_hits += hit.astype(np.float64)
        vert_soft += hit_soft.astype(np.float64)
    denom = np.maximum(vert_views, 1.0)
    per_v = np.where(vert_views > 0, vert_hits / denom, 0.0)
    per_v_soft = np.where(vert_views > 0, vert_soft / denom, 0.0)
    out["per_v"] = per_v
    out["vert_mean"] = float(per_v.mean())
    out["vert_min"] = float(per_v.min())
    k3 = min(3, Nv)
    out["vert_worst3_mean"] = float(np.sort(per_v)[:k3].mean()) if k3 else 0.0
    out["vert_bad_frac"] = float(np.mean(per_v < 0.1))
    out["vert_soft_mean"] = float(per_v_soft.mean())

    if not edges:
        return out
    e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2)
    Ne = e_arr.shape[0]

    # --- edges: sample interior, dilated + tight masks ---
    ts = np.linspace(0.0, 1.0, n_edge_samples + 2)[1:-1]
    a = v[e_arr[:, 0]]; b = v[e_arr[:, 1]]
    samples = a[:, None, :] + (b - a)[:, None, :] * ts[None, :, None]
    sflat = samples.reshape(-1, 3)
    eh = np.zeros(Ne * n_edge_samples, dtype=np.float64)
    eh_soft = np.zeros_like(eh)
    ev = np.zeros_like(eh)
    for view in views:
        in_view, hit = _project_to_mask(sflat, view, "edge_mask")
        _, hit_soft = _project_to_mask(sflat, view, "edge_mask_tight")
        ev += in_view.astype(np.float64)
        eh += hit.astype(np.float64)
        eh_soft += hit_soft.astype(np.float64)
    edenom = np.maximum(ev, 1.0)
    per_sample = np.where(ev > 0, eh / edenom, 0.0)
    per_sample_soft = np.where(ev > 0, eh_soft / edenom, 0.0)
    per_e = per_sample.reshape(Ne, n_edge_samples).mean(axis=1)
    per_e_soft = per_sample_soft.reshape(Ne, n_edge_samples).mean(axis=1)
    out["per_e"] = per_e
    out["edge_mean"] = float(per_e.mean())
    out["edge_min"] = float(per_e.min())
    k3e = min(3, Ne)
    out["edge_worst3_mean"] = float(np.sort(per_e)[:k3e].mean()) if k3e else 0.0
    out["edge_bad_frac"] = float(np.mean(per_e < 0.1))
    out["edge_soft_mean"] = float(per_e_soft.mean())

    # --- endpoint-on-corner: per-edge min(endpoint vertex scores) ---
    if Nv > 0:
        e_endpoint_min = np.minimum(per_v[e_arr[:, 0]], per_v[e_arr[:, 1]])
        out["endpoint_corner_mean"] = float(e_endpoint_min.mean())
        out["endpoint_corner_min"] = float(e_endpoint_min.min())
    return out


def _stage1_displacement_features(
    verts: np.ndarray,
    s1_runs: Optional[List[Tuple[np.ndarray, Any, Any]]],
) -> Tuple[float, float, float]:
    """Distance from candidate vertices to nearest stage-1 prediction.

    Returns ``(disp_mean, disp_max, disp_far_frac)``, where ``disp_far_frac``
    is the fraction of candidate vertices > 0.5 m from any stage-1 vertex.
    """
    if not s1_runs or verts.size == 0:
        return 0.0, 0.0, 0.0
    cv = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    if cv.shape[0] == 0:
        return 0.0, 0.0, 0.0
    s1_v: List[np.ndarray] = []
    for entry in s1_runs:
        rv = entry[0] if isinstance(entry, (tuple, list)) else entry
        rv = np.asarray(rv, dtype=np.float64).reshape(-1, 3)
        if rv.shape[0] > 0:
            s1_v.append(rv)
    if not s1_v:
        return 0.0, 0.0, 0.0
    pool = np.concatenate(s1_v, axis=0)
    D = np.linalg.norm(cv[:, None, :] - pool[None, :, :], axis=-1)
    d_min = D.min(axis=1)
    return (float(d_min.mean()), float(d_min.max()),
            float(np.mean(d_min > 0.5)))


def _connectivity_features(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
) -> Tuple[float, float, float]:
    """``(n_components_norm, cycle_rank_norm, degree_entropy)``.

    * ``n_components_norm`` = (n_components - 1) / max(1, Nv) — 0 when fully
      connected, grows with disconnected fragments.
    * ``cycle_rank_norm``  = (Ne - Nv + n_components) / max(1, Nv) — Betti-1
      density; very low (tree-like) or very high (over-connected) are both
      suspicious.
    * ``degree_entropy``  = Shannon entropy of the degree distribution in nats.
    """
    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    Nv = v.shape[0]
    if Nv == 0 or not edges:
        return 0.0, 0.0, 0.0
    e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2)
    Ne = e_arr.shape[0]
    adj: List[List[int]] = [[] for _ in range(Nv)]
    for a, b in e_arr:
        if 0 <= a < Nv and 0 <= b < Nv and a != b:
            adj[a].append(b); adj[b].append(a)
    visited = [False] * Nv
    n_components = 0
    for s in range(Nv):
        if visited[s]:
            continue
        n_components += 1
        stack = [s]
        while stack:
            u = stack.pop()
            if visited[u]: continue
            visited[u] = True
            for w in adj[u]:
                if not visited[w]:
                    stack.append(w)
    cycle_rank = Ne - Nv + n_components
    deg = np.array([len(a) for a in adj], dtype=np.float64)
    p = np.bincount(deg.astype(np.int64))
    p = p / p.sum() if p.sum() > 0 else p
    p = p[p > 0]
    deg_entropy = float(-(p * np.log(p)).sum()) if p.size > 0 else 0.0
    return (
        float((n_components - 1) / max(1, Nv)),
        float(cycle_rank / max(1, Nv)),
        deg_entropy,
    )


def _symmetry_residual(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
) -> float:
    """Best reflection-plane residual normalised by mean edge length.

    Computes PCA on the candidate vertices, tries reflecting across each of
    the three principal planes (normal = each eigenvector through the
    centroid), Hungarian-matches reflected vs. original verts, and returns
    the minimum mean matched distance / mean edge length. Lower = more
    symmetric. 0 when the wireframe has fewer than 4 vertices or any edge.
    """
    from scipy.optimize import linear_sum_assignment
    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    Nv = v.shape[0]
    if Nv < 4 or not edges:
        return 0.0
    e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2)
    lengths = np.linalg.norm(v[e_arr[:, 0]] - v[e_arr[:, 1]], axis=1)
    mean_len = float(lengths.mean()) if lengths.size > 0 else 0.0
    if mean_len < 1e-6:
        return 0.0
    centroid = v.mean(axis=0)
    Q = v - centroid
    try:
        _, S, Vt = np.linalg.svd(Q, full_matrices=False)
    except np.linalg.LinAlgError:
        return 0.0
    axes = Vt  # rows are unit principal directions
    best = float("inf")
    for k in range(min(3, axes.shape[0])):
        n = axes[k]
        # Reflect Q across plane through centroid with normal n.
        ref = Q - 2.0 * (Q @ n)[:, None] * n[None, :]
        D = np.linalg.norm(Q[:, None, :] - ref[None, :, :], axis=-1)
        row, col = linear_sum_assignment(D)
        matched_d = D[row, col]
        cost = float(matched_d.mean())
        if cost < best:
            best = cost
    if not np.isfinite(best):
        return 0.0
    return float(best / mean_len)


def _per_edge_class_consistency(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    views: List[Dict[str, Any]],
    n_edge_samples: int = 8,
) -> Tuple[float, float]:
    """Gestalt-class purity per edge, aggregated to candidate-level stats.

    For each edge we sample interior points, project to each view, and look
    up the gestalt-class id at the pixel. Across all (sample × view) hits we
    take the modal class and define edge-purity as the fraction of hits
    belonging to that mode. Returns ``(mean_purity, high_purity_frac)``
    where ``high_purity_frac = mean(purity > 0.5)``. Both are 0 with no
    views or edges.
    """
    if not views or not edges or verts.size == 0:
        return 0.0, 0.0
    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2)
    Ne = e_arr.shape[0]
    ts = np.linspace(0.0, 1.0, n_edge_samples + 2)[1:-1]
    a = v[e_arr[:, 0]]; b = v[e_arr[:, 1]]
    samples = a[:, None, :] + (b - a)[:, None, :] * ts[None, :, None]
    sflat = samples.reshape(-1, 3)
    # (Ne*S, V) classes; -1 = no hit
    cls_per_view: List[np.ndarray] = []
    for view in views:
        R = view["R"]; t = view["t"]
        p_cam = sflat @ R.T + t
        z = p_cam[:, 2]
        valid = z > 1e-6
        zsafe = np.where(valid, z, 1.0)
        u_proj = p_cam[:, 0] / zsafe * view["fx"] + view["cx"]
        v_proj = p_cam[:, 1] / zsafe * view["fy"] + view["cy"]
        sx = view["W"] / view["cam_w"]; sy = view["H"] / view["cam_h"]
        ui = np.rint(u_proj * sx).astype(np.int64)
        vi = np.rint(v_proj * sy).astype(np.int64)
        in_view = (valid & (ui >= 0) & (ui < view["W"])
                   & (vi >= 0) & (vi < view["H"]))
        cls = -np.ones(sflat.shape[0], dtype=np.int64)
        if in_view.any():
            idx = np.where(in_view)[0]
            cls[idx] = view["ids"][vi[idx], ui[idx]]
        # Only keep edge-class hits (other classes are background-ish noise)
        is_edge_cls = np.isin(cls, _EDGE_GESTALT_IDS)
        cls = np.where(is_edge_cls, cls, -1)
        cls_per_view.append(cls)
    if not cls_per_view:
        return 0.0, 0.0
    cls_stack = np.stack(cls_per_view, axis=0).reshape(
        len(views), Ne, n_edge_samples)
    # For each edge: flatten (V*S) hits, drop -1, compute mode purity.
    purities = np.zeros(Ne, dtype=np.float64)
    for e_i in range(Ne):
        hits = cls_stack[:, e_i, :].reshape(-1)
        hits = hits[hits >= 0]
        if hits.size == 0:
            purities[e_i] = 0.0
            continue
        counts = np.bincount(hits.astype(np.int64))
        purities[e_i] = float(counts.max()) / float(hits.size)
    return float(purities.mean()), float(np.mean(purities > 0.5))


def _structural_features(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
) -> Tuple[float, float]:
    """Cheap structural priors: degree-1 vertex fraction + edge-length CV."""
    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    Nv = v.shape[0]
    if Nv == 0 or not edges:
        return 0.0, 0.0
    e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2)
    deg = np.zeros(Nv, dtype=np.int64)
    for a, b in e_arr:
        if 0 <= a < Nv: deg[a] += 1
        if 0 <= b < Nv: deg[b] += 1
    degree1_frac = float(np.mean(deg == 1))
    lengths = np.linalg.norm(v[e_arr[:, 0]] - v[e_arr[:, 1]], axis=1)
    mean_len = float(lengths.mean()) if lengths.size > 0 else 0.0
    cv = float(lengths.std() / mean_len) if mean_len > 1e-6 else 0.0
    return degree1_frac, cv


def _coplanarity_score(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    max_cycles: int = 64,
) -> float:
    """Median plane-fit residual of 4-cycles, mapped to [0, 1] (higher = flatter).

    Buildings are mostly composed of planar quads (roof faces, wall panels).
    For every length-4 cycle in the wireframe we compute the third singular
    value of the centred (4,3) vertex matrix — the out-of-plane spread —
    normalised by the cycle's average edge length. The median residual is
    then mapped to ``exp(-5*r)`` so the feature is in [0, 1] with 1 = perfectly
    planar. Returns 0 when no cycles exist.
    """
    Nv = verts.shape[0]
    if Nv < 4 or not edges:
        return 0.0
    adj = [set() for _ in range(Nv)]
    for a, b in edges:
        a, b = int(a), int(b)
        if 0 <= a < Nv and 0 <= b < Nv and a != b:
            adj[a].add(b); adj[b].add(a)
    cycles: List[Tuple[int, int, int, int]] = []
    for u in range(Nv):
        for v_idx in range(u + 1, Nv):
            common = adj[u] & adj[v_idx]
            if len(common) < 2:
                continue
            comm = sorted(c for c in common if c != u and c != v_idx)
            for i in range(len(comm)):
                for j in range(i + 1, len(comm)):
                    cycles.append((u, comm[i], v_idx, comm[j]))
                    if len(cycles) >= max_cycles:
                        break
                if len(cycles) >= max_cycles: break
            if len(cycles) >= max_cycles: break
        if len(cycles) >= max_cycles: break
    if not cycles:
        return 0.0
    vals = []
    for cyc in cycles:
        pts = verts[list(cyc)]
        centroid = pts.mean(axis=0)
        Q = pts - centroid
        try:
            S = np.linalg.svd(Q, compute_uv=False)
        except np.linalg.LinAlgError:
            continue
        if S.size < 3:
            continue
        dev = float(S[-1])
        perim = 0.0
        for k in range(4):
            a, b = cyc[k], cyc[(k + 1) % 4]
            perim += float(np.linalg.norm(verts[a] - verts[b]))
        mean_len = perim / 4.0
        if mean_len < 1e-6:
            continue
        vals.append(dev / mean_len)
    if not vals:
        return 0.0
    return float(np.exp(-5.0 * float(np.median(vals))))


def _junction_angle_dev(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
) -> float:
    """Mean angular deviation of edge-pair junctions from {90°, 180°} in radians.

    At each vertex with degree >= 2, we compute the angle between every pair
    of incident edges and take the smaller deviation to 90° or 180° (the two
    angles that dominate roof/wall geometry). Lower is better; 0 means every
    junction sits exactly at a right or straight angle.
    """
    Nv = verts.shape[0]
    if Nv == 0 or not edges:
        return 0.0
    nbrs: List[List[int]] = [[] for _ in range(Nv)]
    for a, b in edges:
        a, b = int(a), int(b)
        if 0 <= a < Nv and 0 <= b < Nv and a != b:
            nbrs[a].append(b); nbrs[b].append(a)
    devs: List[float] = []
    for u in range(Nv):
        if len(nbrs[u]) < 2:
            continue
        dirs: List[np.ndarray] = []
        for v_idx in nbrs[u]:
            d = verts[v_idx] - verts[u]
            n = float(np.linalg.norm(d))
            if n > 1e-6:
                dirs.append(d / n)
        for i in range(len(dirs)):
            for j in range(i + 1, len(dirs)):
                c = float(np.clip(np.dot(dirs[i], dirs[j]), -1.0, 1.0))
                ang = float(np.arccos(c))
                devs.append(min(abs(ang - np.pi / 2), abs(ang - np.pi)))
    if not devs:
        return 0.0
    return float(np.mean(devs))


def _principal_axis_features(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    cos_thresh: float = 0.9659,  # cos(15°)
) -> Tuple[float, float]:
    """Manhattan structure features from per-edge direction PCA.

    Returns ``(manhattan_frac, top1_frac)``:
      * ``manhattan_frac`` — fraction of edges within 15° of any of the top-3
        eigenvectors of the symmetric second-moment matrix ``sum_i d_i d_i^T``
        (antipodal-symmetric: signs don't matter).
      * ``top1_frac`` — same, but only against the *dominant* eigenvector;
        a proxy for "fraction of edges aligned with the building's principal
        axis" (often vertical in COLMAP world).
    """
    if not edges or verts.shape[0] == 0:
        return 0.0, 0.0
    e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2)
    if e_arr.shape[0] == 0:
        return 0.0, 0.0
    d = verts[e_arr[:, 1]] - verts[e_arr[:, 0]]
    nrm = np.linalg.norm(d, axis=1)
    valid = nrm > 1e-6
    if not valid.any():
        return 0.0, 0.0
    u = d[valid] / nrm[valid, None]
    M = u.T @ u
    try:
        eigvals, eigvecs = np.linalg.eigh(M)
    except np.linalg.LinAlgError:
        return 0.0, 0.0
    order = np.argsort(-eigvals)
    eigvecs = eigvecs[:, order]
    dots = np.abs(u @ eigvecs)
    max_any = dots.max(axis=1)
    max_top1 = dots[:, 0]
    return (float(np.mean(max_any >= cos_thresh)),
            float(np.mean(max_top1 >= cos_thresh)))


_RANKER_FEATURE_NAMES = [
    "n_v",
    "n_e",
    "edge_density",
    "logit_mean",
    "logit_std",
    "logit_min",
    "logit_max",
    "vertex_support_mean",
    "vertex_support_std",
    "vertex_support_min",
    "edge_support_mean",
    "edge_support_max",
    "centrality_mean",
    "centrality_min",
    "centrality_rank",
    "is_raw_member",
    "is_medoid_source",
    "is_conf_source",
    "is_refined",
    "uses_edge_vote",
    "uses_add_edges",
    "uses_topk",
    "source_idx",
    # --- 2D reprojection + structural priors (added) ---
    "vert_reproj_mean",
    "vert_reproj_min",
    "edge_reproj_mean",
    "edge_reproj_min",
    "degree1_frac",
    "edge_length_cv",
    # --- geometric priors (added) ---
    "coplanarity_score",
    "junction_angle_dev",
    "manhattan_frac",
    "principal_axis_align",
    # --- v5 features (selected from v4 ablation: only those that survived
    # 5-fold CV with positive lift in addition to v3) ---
    "s1_disp_far_frac",
    "degree_entropy",
    "symmetry_residual",
]


def _edge_key_set(edges: List[Tuple[int, int]]) -> set:
    out = set()
    for a, b in edges:
        ia, ib = int(a), int(b)
        if ia == ib:
            continue
        out.add((min(ia, ib), max(ia, ib)))
    return out


def _candidate_logits_from_source(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    source_idx: int,
    keep_idx: Optional[np.ndarray],
    n_verts: int,
) -> np.ndarray:
    if source_idx < 0 or source_idx >= len(runs):
        return np.zeros((n_verts,), dtype=np.float64)
    lg = np.asarray(runs[source_idx][2], dtype=np.float64).reshape(-1)
    if keep_idx is not None and lg.size > 0:
        ki = np.asarray(keep_idx, dtype=np.int64).reshape(-1)
        valid = (ki >= 0) & (ki < lg.size)
        if valid.all() and ki.size == n_verts:
            return lg[ki]
    if lg.size == n_verts:
        return lg
    if lg.size == 0:
        return np.zeros((n_verts,), dtype=np.float64)
    return np.full((n_verts,), float(np.nanmean(lg)), dtype=np.float64)


def _candidate_support_features(
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    match_tau_m: float,
) -> Tuple[float, float, float, float, float]:
    """Agreement features for one candidate against the ensemble members."""
    from scipy.optimize import linear_sum_assignment

    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    Nv = v.shape[0]
    if Nv == 0 or not runs:
        return 0.0, 0.0, 0.0, 0.0, 0.0

    v_support = np.zeros((Nv,), dtype=np.float64)
    e_keys = _edge_key_set(edges)
    e_votes: Dict[Tuple[int, int], int] = {k: 0 for k in e_keys}

    for rv_raw, re_raw, _rlg in runs:
        rv = np.asarray(rv_raw, dtype=np.float64).reshape(-1, 3)
        if rv.shape[0] == 0:
            continue
        D = np.linalg.norm(v[:, None, :] - rv[None, :, :], axis=-1)
        D_capped = np.minimum(D, match_tau_m)
        row_ind, col_ind = linear_sum_assignment(D_capped)
        r_to_c = -np.ones(rv.shape[0], dtype=np.int64)
        for r, c in zip(row_ind, col_ind):
            if D[r, c] <= match_tau_m:
                v_support[r] += 1.0
                r_to_c[c] = r
        if e_keys:
            seen = set()
            for a, b in re_raw:
                ia, ib = int(a), int(b)
                if (ia == ib or ia < 0 or ib < 0
                        or ia >= r_to_c.size or ib >= r_to_c.size):
                    continue
                ca, cb = int(r_to_c[ia]), int(r_to_c[ib])
                if ca < 0 or cb < 0 or ca == cb:
                    continue
                key = (min(ca, cb), max(ca, cb))
                if key in e_votes:
                    seen.add(key)
            for key in seen:
                e_votes[key] += 1

    denom = float(max(1, len(runs)))
    v_frac = v_support / denom
    if e_votes:
        e_frac = np.asarray(list(e_votes.values()), dtype=np.float64) / denom
        e_mean = float(e_frac.mean())
        e_max = float(e_frac.max())
    else:
        e_mean = 0.0
        e_max = 0.0
    return (
        float(v_frac.mean()),
        float(v_frac.std()),
        float(v_frac.min()),
        e_mean,
        e_max,
    )


def _candidate_feature_vector(
    name: str,
    verts: np.ndarray,
    edges: List[Tuple[int, int]],
    logits: np.ndarray,
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    match_tau_m: float,
    source_idx: int,
    medoid_idx: int,
    confidence_idx: int,
    ranking: List[int],
    is_raw_member: bool,
    is_refined: bool,
    uses_edge_vote: bool,
    uses_add_edges: bool,
    uses_topk: bool,
    reproj_views: Optional[List[Dict[str, Any]]] = None,
    s1_runs: Optional[List[Tuple[np.ndarray, Any, Any]]] = None,
) -> np.ndarray:
    v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
    e = list(edges)
    lg = np.asarray(logits, dtype=np.float64).reshape(-1)
    lg = lg[np.isfinite(lg)]
    if lg.size == 0:
        lg = np.zeros((1,), dtype=np.float64)

    n_v = float(v.shape[0])
    n_e = float(len(e))
    max_edges = max(1.0, n_v * max(0.0, n_v - 1.0) / 2.0)
    v_sup_mean, v_sup_std, v_sup_min, e_sup_mean, e_sup_max = (
        _candidate_support_features(v, e, runs, match_tau_m)
    )

    dists = []
    for rv, re, _rl in runs:
        try:
            dists.append(_wireframe_pair_distance(
                v, e, np.asarray(rv, dtype=np.float64), list(re),
                match_tau_m=match_tau_m,
            ))
        except Exception:
            continue
    d_arr = np.asarray([d for d in dists if np.isfinite(d)], dtype=np.float64)
    if d_arr.size == 0:
        centrality_mean = 10.0
        centrality_min = 10.0
    else:
        centrality_mean = float(d_arr.mean())
        centrality_min = float(d_arr.min())

    rank_pos = float(len(ranking))
    if source_idx in ranking:
        rank_pos = float(ranking.index(source_idx))

    if reproj_views:
        vr_mean, vr_min, er_mean, er_min = _reprojection_features(
            v, e, reproj_views,
        )
    else:
        vr_mean = vr_min = er_mean = er_min = 0.0
    deg1_frac, edge_len_cv = _structural_features(v, e)
    coplanarity = _coplanarity_score(v, e)
    junction_dev = _junction_angle_dev(v, e)
    manhattan_frac, princ_axis_align = _principal_axis_features(v, e)

    _, _, s1_disp_far = _stage1_displacement_features(v, s1_runs)
    _, _, deg_ent = _connectivity_features(v, e)
    sym_res = _symmetry_residual(v, e)

    feats = np.asarray([
        n_v / 64.0,
        n_e / 128.0,
        n_e / max_edges,
        float(lg.mean()),
        float(lg.std()),
        float(lg.min()),
        float(lg.max()),
        v_sup_mean,
        v_sup_std,
        v_sup_min,
        e_sup_mean,
        e_sup_max,
        centrality_mean,
        centrality_min,
        rank_pos / float(max(1, len(ranking))),
        1.0 if is_raw_member else 0.0,
        1.0 if source_idx == medoid_idx else 0.0,
        1.0 if source_idx == confidence_idx else 0.0,
        1.0 if is_refined else 0.0,
        1.0 if uses_edge_vote else 0.0,
        1.0 if uses_add_edges else 0.0,
        1.0 if uses_topk else 0.0,
        float(source_idx) / float(max(1, len(runs) - 1)),
        vr_mean,
        vr_min,
        er_mean,
        er_min,
        deg1_frac,
        edge_len_cv,
        coplanarity,
        junction_dev,
        manhattan_frac,
        princ_axis_align,
        # --- v5 features ---
        s1_disp_far,
        deg_ent,
        sym_res,
    ], dtype=np.float64)
    if feats.shape[0] != len(_RANKER_FEATURE_NAMES):
        raise RuntimeError(f"ranker feature mismatch for {name}")
    return feats


def _ensemble_heuristic_score(features: np.ndarray) -> float:
    f = {k: float(v) for k, v in zip(_RANKER_FEATURE_NAMES, features)}
    return (
        1.20 * f["vertex_support_mean"]
        + 0.55 * f["edge_support_mean"]
        + 0.08 * f["logit_mean"]
        - 0.30 * f["centrality_mean"]
        - 0.08 * abs(f["n_v"] - 0.33)
        - 0.04 * abs(f["n_e"] - 0.18)
        + 0.04 * f["is_refined"]
        + 0.02 * f["is_medoid_source"]
        - 0.03 * f["uses_edge_vote"]
    )


def _score_ranker_candidate(
    features: np.ndarray,
    ranker: Optional[Dict[str, np.ndarray]],
) -> float:
    if ranker is None:
        return _ensemble_heuristic_score(features)
    mean = np.asarray(ranker["mean"], dtype=np.float64)
    scale = np.asarray(ranker["scale"], dtype=np.float64)
    weights = np.asarray(ranker["weights"], dtype=np.float64)
    x = (features - mean) / np.maximum(scale, 1e-8)
    return float(np.dot(x, weights[1:]) + weights[0])


def _fuse_topk_candidates(
    candidates: List[Dict[str, Any]],
    scores: np.ndarray,
    k: int,
    match_tau_m: float,
    edge_vote_frac: float,
) -> Tuple[np.ndarray, List[Tuple[int, int]], int]:
    """Fuse the top-K ranker-scored candidates into one wireframe.

    Strategy mirrors ``medoid_mean_add`` but applies it to the top-K
    *ranker-quality-filtered* candidates rather than all ensemble members:

      1. Sort candidates by ``scores`` descending; keep ``top_k = min(k, len)``.
      2. Pivot = top-1 candidate. For each pivot vertex, Hungarian-match to
         each of the other top-K candidates' vertices (cap τ = ``match_tau_m``)
         and average the positions of matched correspondents.
      3. Union pivot's edges with edges voted by at least
         ``edge_vote_frac`` of the top-K (counted on matched vertex pairs).

    Returns ``(verts, edges, best_idx)`` where ``best_idx`` is the position
    of the top-1 candidate in the input list (for diagnostics).
    """
    scores_arr = np.asarray(scores, dtype=np.float64).reshape(-1)
    n = scores_arr.size
    if n == 0:
        return np.zeros((0, 3), dtype=np.float64), [], 0
    k_eff = max(1, min(int(k), n))
    order = np.argsort(-scores_arr)[:k_eff].astype(np.int64)
    best_idx = int(order[0])
    if k_eff == 1:
        c = candidates[best_idx]
        return (np.asarray(c["verts"], dtype=np.float64),
                list(c["edges"]), best_idx)

    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = []
    for j in order:
        c = candidates[int(j)]
        v = np.asarray(c["verts"], dtype=np.float64).reshape(-1, 3)
        e = list(c["edges"])
        lg_src = c.get("logits", None)
        lg = (np.asarray(lg_src, dtype=np.float64).reshape(-1)
              if lg_src is not None and len(np.asarray(lg_src).ravel()) == v.shape[0]
              else np.zeros((v.shape[0],), dtype=np.float64))
        runs.append((v, e, lg))

    refined_v, refined_e, keep_idx = _refine_positions(
        runs, picked_idx=0, match_tau_m=match_tau_m,
        aggregator="mean", member_indices=None, vertex_vote_frac=0.0,
    )
    voted = _vote_edges(
        runs, picked_idx=0, match_tau_m=match_tau_m,
        vote_frac=edge_vote_frac, member_indices=None,
        keep_idx=keep_idx,
    )
    edge_set = set()
    for (a, b) in refined_e:
        ia, ib = int(a), int(b)
        if ia == ib: continue
        edge_set.add((min(ia, ib), max(ia, ib)))
    for (a, b) in voted:
        ia, ib = int(a), int(b)
        if ia == ib: continue
        edge_set.add((min(ia, ib), max(ia, ib)))
    return refined_v, sorted(edge_set), best_idx


def _load_ensemble_ranker(path: Optional[str]) -> Optional[Dict[str, np.ndarray]]:
    if not path:
        return None
    p = Path(path)
    if not p.exists():
        print(f"[ranker] {p} not found; using built-in heuristic scorer", flush=True)
        return None
    data = np.load(p, allow_pickle=False)
    names = [str(x) for x in data["feature_names"].tolist()]
    if names != _RANKER_FEATURE_NAMES:
        print(
            f"[ranker] feature-name mismatch in {p} "
            f"(have {len(names)}-d, expected {len(_RANKER_FEATURE_NAMES)}-d); "
            f"falling back to built-in heuristic scorer", flush=True)
        return None
    ranker = {
        "mean": np.asarray(data["mean"], dtype=np.float64),
        "scale": np.asarray(data["scale"], dtype=np.float64),
        "weights": np.asarray(data["weights"], dtype=np.float64),
    }
    print(f"[ranker] loaded {p}", flush=True)
    return ranker


def _build_ranker_candidates(
    runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]],
    match_tau_m: float,
    edge_vote_frac: float,
    ranking: Optional[List[int]] = None,
    medoid_idx: Optional[int] = None,
    confidence_idx: Optional[int] = None,
    reproj_views: Optional[List[Dict[str, Any]]] = None,
    s1_runs: Optional[List[Tuple[np.ndarray, Any, Any]]] = None,
) -> List[Dict[str, Any]]:
    """Produce candidate wireframes for learned/heuristic ensemble selection."""
    if not runs:
        return []
    if ranking is None:
        ranking = _medoid_rank(runs, match_tau_m=match_tau_m)
    medoid_idx = int(medoid_idx if medoid_idx is not None
                     else (ranking[0] if ranking else 0))
    confidence_idx = int(confidence_idx if confidence_idx is not None
                         else _confidence_select_idx(runs))

    candidates: List[Dict[str, Any]] = []

    def add_candidate(
        name: str,
        verts: np.ndarray,
        edges: List[Tuple[int, int]],
        source_idx: int,
        keep_idx: Optional[np.ndarray],
        is_raw_member: bool,
        is_refined: bool,
        uses_edge_vote: bool,
        uses_add_edges: bool,
        uses_topk: bool,
    ) -> None:
        v = np.asarray(verts, dtype=np.float64).reshape(-1, 3)
        e = [(int(a), int(b)) for a, b in edges if int(a) != int(b)]
        if v.shape[0] == 0 or len(e) == 0:
            return
        logits = _candidate_logits_from_source(runs, source_idx, keep_idx, v.shape[0])
        features = _candidate_feature_vector(
            name=name, verts=v, edges=e, logits=logits, runs=runs,
            match_tau_m=match_tau_m, source_idx=source_idx,
            medoid_idx=medoid_idx, confidence_idx=confidence_idx,
            ranking=ranking, is_raw_member=is_raw_member,
            is_refined=is_refined, uses_edge_vote=uses_edge_vote,
            uses_add_edges=uses_add_edges, uses_topk=uses_topk,
            reproj_views=reproj_views,
            s1_runs=s1_runs,
        )
        candidates.append({
            "name": name,
            "verts": v,
            "edges": e,
            "logits": logits,
            "features": features,
            "source_idx": int(source_idx),
        })

    for i, (v, e, _lg) in enumerate(runs):
        add_candidate(
            name=f"member_{i:02d}", verts=v, edges=list(e), source_idx=i,
            keep_idx=None, is_raw_member=True, is_refined=False,
            uses_edge_vote=False, uses_add_edges=False, uses_topk=False,
        )

    def refined_candidates(source_idx: int, label: str, member_indices: Optional[List[int]]) -> None:
        uses_topk = member_indices is not None
        rv, re, keep_idx = _refine_positions(
            runs, source_idx, match_tau_m=match_tau_m,
            aggregator="mean", member_indices=member_indices,
            vertex_vote_frac=0.0,
        )
        add_candidate(
            name=f"{label}_mean_picked" + ("_topk" if uses_topk else ""),
            verts=rv, edges=re, source_idx=source_idx, keep_idx=keep_idx,
            is_raw_member=False, is_refined=True,
            uses_edge_vote=False, uses_add_edges=False, uses_topk=uses_topk,
        )
        voted = _vote_edges(
            runs, source_idx, match_tau_m=match_tau_m,
            vote_frac=edge_vote_frac, member_indices=member_indices,
            keep_idx=keep_idx,
        )
        add_candidate(
            name=f"{label}_mean_voted" + ("_topk" if uses_topk else ""),
            verts=rv, edges=voted, source_idx=source_idx, keep_idx=keep_idx,
            is_raw_member=False, is_refined=True,
            uses_edge_vote=True, uses_add_edges=False, uses_topk=uses_topk,
        )
        add_edges = sorted(_edge_key_set(re) | _edge_key_set(voted))
        add_candidate(
            name=f"{label}_mean_add" + ("_topk" if uses_topk else ""),
            verts=rv, edges=add_edges, source_idx=source_idx, keep_idx=keep_idx,
            is_raw_member=False, is_refined=True,
            uses_edge_vote=True, uses_add_edges=True, uses_topk=uses_topk,
        )

    refined_candidates(medoid_idx, "medoid", None)
    if confidence_idx != medoid_idx:
        refined_candidates(confidence_idx, "confidence", None)

    for k in (5, 8):
        if ranking and len(ranking) >= k:
            members = list(ranking[:k])
            if medoid_idx not in members:
                members.append(medoid_idx)
            refined_candidates(medoid_idx, f"medoid_top{k}", members)

    return candidates


def train_ensemble_ranker_jsonl(jsonl_path: str, out_path: str, l2: float = 1e-2) -> None:
    """Fit a tiny ridge ranker from within-sample candidate comparisons.

    Absolute HSS is dominated by sample difficulty ("easy building" vs.
    "hard building"), but inference only needs to choose among candidates for
    the same sample. Training on pairwise feature differences cancels that
    nuisance and directly optimizes the ordering we care about.
    """
    rows_by_sample: Dict[str, List[Tuple[np.ndarray, float]]] = {}
    with open(jsonl_path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            feats = np.asarray(row.get("features", []), dtype=np.float64)
            if feats.shape[0] != len(_RANKER_FEATURE_NAMES):
                continue
            y = float(row.get("hss", 0.0))
            if not np.isfinite(y) or not np.all(np.isfinite(feats)):
                continue
            sid = str(row.get("order_id", row.get("sample_index", "")))
            rows_by_sample.setdefault(sid, []).append((feats, y))

    xs = [feat for rows in rows_by_sample.values() for feat, _y in rows]
    if len(xs) < 10 or len(rows_by_sample) < 2:
        raise SystemExit(
            f"[ranker/train] need >=10 rows and >=2 samples, got "
            f"{len(xs)} rows / {len(rows_by_sample)} samples")

    X = np.stack(xs, axis=0)
    mean = X.mean(axis=0)
    scale = X.std(axis=0)
    scale[scale < 1e-8] = 1.0

    pair_x: List[np.ndarray] = []
    pair_y: List[float] = []
    for rows in rows_by_sample.values():
        if len(rows) < 2:
            continue
        feats = [(feat - mean) / scale for feat, _y in rows]
        ys = [float(y) for _feat, y in rows]
        for i in range(len(rows)):
            for j in range(i + 1, len(rows)):
                dy = ys[i] - ys[j]
                if abs(dy) < 1e-9:
                    continue
                pair_x.append(feats[i] - feats[j])
                pair_y.append(dy)
    if len(pair_x) < 10:
        raise SystemExit(f"[ranker/train] need >=10 informative pairs, got {len(pair_x)}")

    DX = np.stack(pair_x, axis=0)
    dy = np.asarray(pair_y, dtype=np.float64)
    reg = np.eye(DX.shape[1], dtype=np.float64) * float(l2)
    w = np.linalg.solve(DX.T @ DX + reg, DX.T @ dy)
    weights = np.concatenate([[0.0], w], axis=0)
    pred = DX @ w
    mse = float(np.mean((pred - dy) ** 2))
    corr = float(np.corrcoef(pred, dy)[0, 1]) if pred.std() > 0 and dy.std() > 0 else 0.0

    # Grouped oracle estimate: how good would this ranker be on the same dump?
    # This is optimistic when train==eval, but useful as a smoke test.
    picked = []
    oracle = []
    for rows in rows_by_sample.values():
        scored = []
        for feat, y in rows:
            x = np.concatenate([[1.0], (feat - mean) / scale])
            scored.append((float(np.dot(x, weights)), float(y)))
        picked.append(max(scored, key=lambda t: t[0])[1])
        oracle.append(max(scored, key=lambda t: t[1])[1])

    np.savez(
        out_path,
        feature_names=np.asarray(_RANKER_FEATURE_NAMES, dtype="<U64"),
        mean=mean,
        scale=scale,
        weights=weights,
    )
    print(f"[ranker/train] rows={len(xs)} samples={len(rows_by_sample)} "
          f"pairs={len(pair_x)} pair_mse={mse:.6f} pair_corr={corr:.4f}",
          flush=True)
    if picked:
        print(f"[ranker/train] train-picked mean={float(np.mean(picked)):.4f}  "
              f"candidate-oracle mean={float(np.mean(oracle)):.4f}", flush=True)
    print(f"[ranker/train] wrote {out_path}", flush=True)


def _run_stage2_for_seed(
    raw_sample: Dict[str, Any],
    colmap_rec,
    pred_verts: np.ndarray,
    pred_edges: np.ndarray,
    stage2_cfg: Stage2Config,
    device: torch.device,
    preprocess_cache: Optional[Dict[str, Any]],
    timings_accum: Optional[Dict[str, float]],
    order_id: str,
) -> Tuple[Optional[np.ndarray], Optional[List[Tuple[int, int]]],
           Optional[np.ndarray], bool]:
    """Build stage-2 scene from a single stage-1 prediction and run stage-2
    refinement. Returns ``(verts_world, edges, valid_logits, hull_valid)``.
    ``valid_logits`` is a (M,) float32 array of raw validity logits for the
    kept vertices (used by confidence-mode ensemble selection). When the hull
    is empty or stage-2 produces nothing, returns
    ``(None, None, None, False)`` so the caller can fall back to stage-1."""
    def _sync():
        if device.type == "cuda":
            torch.cuda.synchronize()

    rng = np.random.default_rng(_seed_from_order_id(order_id))
    t0 = time.perf_counter()
    s2_row = build_stage2_scene(
        raw_sample=raw_sample,
        colmap_rec=colmap_rec,
        pred_verts=pred_verts,
        pred_edges=pred_edges,
        above_m=stage2_cfg.above_m,
        below_m=stage2_cfg.below_m,
        side_m=stage2_cfg.side_m,
        n_col_oversample=stage2_cfg.n_col_oversample,
        n_dep_oversample=stage2_cfg.n_dep_oversample,
        n_depth_per_image_cap=stage2_cfg.n_depth_per_image_cap,
        sampling_mode=("voxel" if stage2_cfg.sampling == "voxel" else "random"),
        density_voxel_size_m=stage2_cfg.density_voxel_size_m,
        density_kernel_radius=stage2_cfg.density_kernel_radius,
        density_kernel_axis=stage2_cfg.density_kernel_axis,
        density_response_power=stage2_cfg.density_response_power,
        density_planarity_suppression=stage2_cfg.density_planarity_suppression,
        density_planarity_radius=stage2_cfg.density_planarity_radius,
        density_planarity_min_points=stage2_cfg.density_planarity_min_points,
        density_min_per_voxel=stage2_cfg.density_min_per_voxel,
        rng=rng,
        cache=preprocess_cache,
    )
    if timings_accum is not None:
        timings_accum["s2_build"] = timings_accum.get("s2_build", 0.0) + (time.perf_counter() - t0)

    t0 = time.perf_counter()
    s2_sample = stage2_row_to_sample(
        s2_row,
        n_pts=stage2_cfg.n_pts,
        k_verts=stage2_cfg.k_verts,
        augment=False,
        flip=False, yaw=False, jitter_sigma=0.0,
        pre_subsample=(stage2_cfg.sampling != "fps"),
    )
    if timings_accum is not None:
        timings_accum["s2_transform"] = timings_accum.get("s2_transform", 0.0) + (time.perf_counter() - t0)

    batch2 = _stage2_sample_to_batch(s2_sample, device)
    if stage2_cfg.sampling == "fps":
        seed = torch.zeros(1, device=device, dtype=torch.long)
        batch2 = fps_subsample_stage2_batch(
            batch2,
            n_col=stage2_cfg.n_pts // 2,
            n_dep=stage2_cfg.n_pts - stage2_cfg.n_pts // 2,
            seed_per_sample=seed,
            renormalize_bbox=True,
            max_exact_iters_per_provenance=stage2_cfg.fps_max_exact_iters,
        )

    s2_model = stage2_cfg.model
    n_steps2 = int(stage2_cfg.n_sample_steps)
    vth2 = float(stage2_cfg.validity_thresh)
    K2 = s2_model.denoiser.k_verts
    B2 = 1
    step_dt2 = 1.0 / max(1, n_steps2)
    with torch.inference_mode():
        t0 = time.perf_counter()
        scene_feats2, scene_xyz_norm2 = s2_model._encode_scene(batch2)
        _sync()
        if timings_accum is not None:
            timings_accum["s2_enc"] = timings_accum.get("s2_enc", 0.0) + (time.perf_counter() - t0)

        x = s2_model._init_x0(batch2, K2, device, B2)
        last_logit = torch.zeros(B2, K2, device=device)
        last_edge_logit = torch.zeros(B2, K2, K2, device=device)
        last_k = max(1, int(getattr(stage2_cfg, "ensemble_stage2_last_k", 1)))
        start_avg = n_steps2 - last_k
        x_sum: Optional[torch.Tensor] = None
        logit_sum: Optional[torch.Tensor] = None
        edge_sum: Optional[torch.Tensor] = None
        n_accum = 0
        t0 = time.perf_counter()
        for s_idx in range(n_steps2):
            t = torch.full((B2,), s_idx * step_dt2, device=device)
            v, last_logit, last_edge_logit = s2_model.denoiser(
                x, t, scene_feats2, scene_xyz_norm2,
            )
            x = x + step_dt2 * v
            if last_k > 1 and s_idx >= start_avg:
                if x_sum is None:
                    x_sum = x.clone()
                    logit_sum = last_logit.clone()
                    edge_sum = last_edge_logit.clone()
                else:
                    x_sum.add_(x)
                    logit_sum.add_(last_logit)
                    edge_sum.add_(last_edge_logit)
                n_accum += 1
        if x_sum is not None and n_accum > 0:
            inv = 1.0 / float(n_accum)
            x = x_sum * inv
            last_logit = logit_sum * inv
            last_edge_logit = edge_sum * inv
        _sync()
        if timings_accum is not None:
            timings_accum["s2_den"] = timings_accum.get("s2_den", 0.0) + (time.perf_counter() - t0)

        xyz_norm = x[0]
        valid = last_logit[0] > vth2
        edge_logit_ = last_edge_logit[0].float()
        valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
        xyz_valid = xyz_norm.index_select(0, valid_idx)
        center = batch2["bbox_center"][0].to(xyz_valid)
        scale = batch2["bbox_scale"][0].to(xyz_valid)
        bbox_R = batch2.get("bbox_R")
        if isinstance(bbox_R, torch.Tensor):
            R = bbox_R[0].to(xyz_valid)
            verts_world = (xyz_valid * scale) @ R + center
        else:
            verts_world = xyz_valid * scale + center
        s2_verts = verts_world.cpu().numpy()
        s2_logits = (last_logit[0].index_select(0, valid_idx)
                     .detach().cpu().numpy().astype(np.float32))
        s2_edges: List[Tuple[int, int]] = []
        n_valid = int(valid_idx.numel())
        if n_valid >= 2:
            sub = edge_logit_.index_select(0, valid_idx).index_select(1, valid_idx)
            tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device,
                                        dtype=torch.bool), diagonal=1)
            pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False)
            s2_edges = [(int(a), int(b)) for a, b in pairs.detach().cpu().tolist()]

    hull_valid = bool(s2_row.get("hull_valid", False))
    if not hull_valid or len(s2_verts) == 0 or len(s2_edges) == 0:
        return None, None, None, hull_valid
    return np.asarray(s2_verts, dtype=float), s2_edges, s2_logits, hull_valid


def _predict_union_hull_ensemble(
    raw_sample: Dict[str, Any],
    colmap_rec,
    s1_runs: List[Tuple[np.ndarray, np.ndarray, np.ndarray]],
    stage2_cfg: Stage2Config,
    device: torch.device,
    preprocess_cache: Optional[Dict[str, Any]],
    timings_accum: Optional[Dict[str, float]],
    order_id: str,
) -> Tuple[List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], bool]:
    """Union-hull batched ensemble.

    Strategy:
      1. Concatenate the N stage-1 vertex predictions into one cloud and build
         **one** stage-2 scene (single hull, single point cloud, single bbox).
      2. Transform once → one shared sample. Encode once (B=1).
      3. Build B=N denoiser batch sharing the scene but with per-slot
         ``init_verts`` derived from each stage-1 prediction (each slot
         starts its trajectory from a different stage-1 hypothesis).
      4. Run the stage-2 Euler loop at B=N — broadcasts the shared encoded
         features via ``.expand`` (no extra GPU memory) for ``N`` parallel
         trajectories with diverse inits.

    This shares the heaviest stage-2 cost (encoder + build) across all N
    seeds, so total per-sample cost is ~1.0-1.2x single-run wall vs 1.5-1.8x
    for the per-seed strategy at N=5. Returns ``(seed_results, hull_valid)``
    where ``seed_results`` is ``[(verts_world, edges, valid_logits)]`` per
    slot, in the same shape as ``_stage1_seed_ensemble``. Empty list on
    catastrophic failure; caller handles the fallback.
    """
    if not s1_runs:
        return [], False

    # 1) Union of stage-1 vertex predictions for the hull.
    union_verts_parts = []
    union_edges_parts = []
    offset = 0
    for v, e, _l in s1_runs:
        if len(v) == 0:
            continue
        union_verts_parts.append(np.asarray(v, dtype=np.float32).reshape(-1, 3))
        if len(e) > 0:
            ea = np.asarray(e, dtype=np.int64).reshape(-1, 2)
            union_edges_parts.append(ea + offset)
        offset += len(v)
    if not union_verts_parts:
        return [], False
    union_verts = np.concatenate(union_verts_parts, axis=0)
    union_edges = (np.concatenate(union_edges_parts, axis=0)
                   if union_edges_parts else np.zeros((0, 2), dtype=np.int64))

    # 2) Build the single shared stage-2 scene.
    rng = np.random.default_rng(_seed_from_order_id(order_id))
    t0 = time.perf_counter()
    s2_row = build_stage2_scene(
        raw_sample=raw_sample,
        colmap_rec=colmap_rec,
        pred_verts=union_verts,
        pred_edges=union_edges,
        above_m=stage2_cfg.above_m,
        below_m=stage2_cfg.below_m,
        side_m=stage2_cfg.side_m,
        n_col_oversample=stage2_cfg.n_col_oversample,
        n_dep_oversample=stage2_cfg.n_dep_oversample,
        n_depth_per_image_cap=stage2_cfg.n_depth_per_image_cap,
        sampling_mode=("voxel" if stage2_cfg.sampling == "voxel" else "random"),
        density_voxel_size_m=stage2_cfg.density_voxel_size_m,
        density_kernel_radius=stage2_cfg.density_kernel_radius,
        density_kernel_axis=stage2_cfg.density_kernel_axis,
        density_response_power=stage2_cfg.density_response_power,
        density_planarity_suppression=stage2_cfg.density_planarity_suppression,
        density_planarity_radius=stage2_cfg.density_planarity_radius,
        density_planarity_min_points=stage2_cfg.density_planarity_min_points,
        density_min_per_voxel=stage2_cfg.density_min_per_voxel,
        rng=rng,
        cache=preprocess_cache,
    )
    if timings_accum is not None:
        timings_accum["s2_build"] = timings_accum.get("s2_build", 0.0) + (time.perf_counter() - t0)

    hull_valid = bool(s2_row.get("hull_valid", False))
    if not hull_valid:
        return [], False

    # 3) Transform once → one shared sample (random-sampling path: scene gets
    #    bbox-normalised here; init_verts is the field we'll override per slot).
    t0 = time.perf_counter()
    s2_sample = stage2_row_to_sample(
        s2_row,
        n_pts=stage2_cfg.n_pts,
        k_verts=stage2_cfg.k_verts,
        augment=False,
        flip=False, yaw=False, jitter_sigma=0.0,
        pre_subsample=(stage2_cfg.sampling != "fps"),
    )
    if timings_accum is not None:
        timings_accum["s2_transform"] = timings_accum.get("s2_transform", 0.0) + (time.perf_counter() - t0)

    batch2_single = _stage2_sample_to_batch(s2_sample, device)
    if stage2_cfg.sampling == "fps":
        seed = torch.zeros(1, device=device, dtype=torch.long)
        batch2_single = fps_subsample_stage2_batch(
            batch2_single,
            n_col=stage2_cfg.n_pts // 2,
            n_dep=stage2_cfg.n_pts - stage2_cfg.n_pts // 2,
            seed_per_sample=seed,
            renormalize_bbox=True,
            max_exact_iters_per_provenance=stage2_cfg.fps_max_exact_iters,
        )

    # 4) Per-slot init_verts: each slot starts from its own stage-1 prediction,
    #    normalised into the shared bbox frame (same bbox the encoder sees).
    K = stage2_cfg.k_verts
    N = len(s1_runs)
    bbox_center_np = batch2_single["bbox_center"][0].detach().cpu().numpy().astype(np.float32)
    bbox_scale_val = float(batch2_single["bbox_scale"][0].detach().cpu().item())
    init_padded_n = np.zeros((N, K, 3), dtype=np.float32)
    init_valid_n = np.zeros((N, K), dtype=bool)
    seed_root = _seed_from_order_id(order_id)
    for i, (verts_w, _e, _l) in enumerate(s1_runs):
        if len(verts_w) == 0:
            continue
        init_norm = ((np.asarray(verts_w, dtype=np.float32).reshape(-1, 3)
                      - bbox_center_np) / max(bbox_scale_val, 1e-6)).astype(np.float32)
        rng_i = np.random.default_rng(seed_root + i + 1)
        ip, iv = _pad_init_verts(init_norm, K, rng=rng_i)
        init_padded_n[i] = ip
        init_valid_n[i] = iv

    # 5) Build B=N denoiser batch sharing every per-point field, overriding
    #    only init_verts / init_verts_valid per slot.
    batch2_N = _replicate_batch(batch2_single, N)
    batch2_N["init_verts"] = torch.from_numpy(init_padded_n).to(device)
    batch2_N["init_verts_valid"] = torch.from_numpy(init_valid_n).to(device)

    # 6) Encoder once (B=1), broadcast features to B=N for the denoiser.
    s2_model = stage2_cfg.model
    n_steps2 = int(stage2_cfg.n_sample_steps)
    vth2 = float(stage2_cfg.validity_thresh)
    step_dt2 = 1.0 / max(1, n_steps2)
    def _sync():
        if device.type == "cuda":
            torch.cuda.synchronize()

    with torch.inference_mode():
        t0 = time.perf_counter()
        scene_feats_1, scene_xyz_1 = s2_model._encode_scene(batch2_single)
        scene_feats_N = scene_feats_1.expand(N, -1, -1).contiguous()
        scene_xyz_N = scene_xyz_1.expand(N, -1, -1).contiguous()
        _sync()
        if timings_accum is not None:
            timings_accum["s2_enc"] = timings_accum.get("s2_enc", 0.0) + (time.perf_counter() - t0)

        x = s2_model._init_x0(batch2_N, K, device, N)
        last_logit = torch.zeros(N, K, device=device)
        last_edge_logit = torch.zeros(N, K, K, device=device)
        last_k = max(1, int(getattr(stage2_cfg, "ensemble_stage2_last_k", 1)))
        start_avg = n_steps2 - last_k
        x_sum: Optional[torch.Tensor] = None
        logit_sum: Optional[torch.Tensor] = None
        edge_sum: Optional[torch.Tensor] = None
        n_accum = 0
        t0 = time.perf_counter()
        for s_idx in range(n_steps2):
            t = torch.full((N,), s_idx * step_dt2, device=device)
            v, last_logit, last_edge_logit = s2_model.denoiser(
                x, t, scene_feats_N, scene_xyz_N,
            )
            x = x + step_dt2 * v
            if last_k > 1 and s_idx >= start_avg:
                if x_sum is None:
                    x_sum = x.clone()
                    logit_sum = last_logit.clone()
                    edge_sum = last_edge_logit.clone()
                else:
                    x_sum.add_(x)
                    logit_sum.add_(last_logit)
                    edge_sum.add_(last_edge_logit)
                n_accum += 1
        if x_sum is not None and n_accum > 0:
            inv = 1.0 / float(n_accum)
            x = x_sum * inv
            last_logit = logit_sum * inv
            last_edge_logit = edge_sum * inv
        _sync()
        if timings_accum is not None:
            timings_accum["s2_den"] = timings_accum.get("s2_den", 0.0) + (time.perf_counter() - t0)

        # 7) Extract per-slot predictions in world coords.
        results: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = []
        bbox_R = batch2_N.get("bbox_R")
        for i in range(N):
            valid = last_logit[i] > vth2
            idx = torch.nonzero(valid, as_tuple=False).flatten()
            x_valid = x[i].index_select(0, idx)
            center = batch2_N["bbox_center"][i].to(device=x_valid.device, dtype=torch.float32)
            scale = batch2_N["bbox_scale"][i].to(device=x_valid.device, dtype=torch.float32)
            if isinstance(bbox_R, torch.Tensor):
                R = bbox_R[i].to(device=x_valid.device, dtype=torch.float32)
                verts_world = (x_valid * scale) @ R + center
            else:
                verts_world = x_valid * scale + center
            verts_np = verts_world.detach().cpu().numpy().astype(np.float32).reshape(-1, 3)
            logits_kept = (last_logit[i].index_select(0, idx)
                           .detach().cpu().numpy().astype(np.float32))

            n_valid = int(idx.numel())
            edges_i: List[Tuple[int, int]] = []
            if n_valid >= 2:
                sub = last_edge_logit[i].float().index_select(0, idx).index_select(1, idx)
                tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device,
                                            dtype=torch.bool), diagonal=1)
                pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False)
                edges_i = [(int(a), int(b)) for a, b in pairs.detach().cpu().tolist()]
            results.append((verts_np, edges_i, logits_kept))
    return results, True


def predict_two_stage(
    raw_sample: Dict[str, Any],
    colmap_rec,
    stage1_scene: Dict[str, np.ndarray],
    stage1_model: WireframeDiffusion,
    stage1_n_steps: int,
    stage1_validity_thresh: float,
    stage2_cfg: Stage2Config,
    device: torch.device,
    preprocess_cache: Optional[Dict[str, Any]] = None,
    stats: Optional[Dict[str, int]] = None,
    timings: Optional[Dict[str, float]] = None,
    diagnostics: Optional[Dict[str, Any]] = None,
) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
    """Run stage 1 (coarse) → build stage-2 scene from its prediction → run
    stage 2 (refinement). Returns the stage-2 verts + edges in world coords.

    When ``stage2_cfg.ensemble_n > 1``, runs N independent (stage-1, stage-2)
    trajectories with different random init x0 (stage-1 denoiser batched over
    seeds, stage-2 sequential since each seed produces its own hull) and fuses
    the N wireframes by vertex-merge consensus.

    If ``timings`` is provided, populates per-phase wallclock seconds summed
    across seeds: ``s2_build``, ``s2_transform``, ``s2_enc``, ``s2_den``.
    Stage-1 inference is NOT timed here because the sanity check inlines
    stage 1 with its own per-phase timing.
    """
    order_id = str(raw_sample.get("order_id", ""))
    n_ensemble = max(1, int(getattr(stage2_cfg, "ensemble_n", 1)))
    strategy = str(getattr(stage2_cfg, "ensemble_strategy", "union_hull"))

    # 1) Stage-1 inference (GPU) — N seeds batched on one shared scene.
    batch1 = scene_to_batch(stage1_scene, device)
    with torch.inference_mode():
        s1_runs = _stage1_seed_ensemble(
            stage1_model, batch1,
            n_steps=stage1_n_steps,
            validity_thresh=stage1_validity_thresh,
            n_seeds=n_ensemble,
        )

    # 2) Stage-2 ensembling.
    #
    #   - "union_hull": build ONE stage-2 scene from the union of all N
    #     stage-1 vertex predictions, encode ONCE, then run a B=N denoiser
    #     batch with per-slot init_verts from each stage-1 hypothesis.
    #     Cheap (~1.0-1.2x single-run wall) and gives N stage-2 predictions
    #     that share a common reference frame.
    #
    #   - "per_seed": legacy path — N independent (build → transform → enc →
    #     den) passes, sequential because hulls differ. Used as fallback when
    #     the union-hull build fails (no valid hull).
    seed_results: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = []
    n_hull_valid = 0

    if strategy == "union_hull" and n_ensemble > 1:
        union_results, hull_valid = _predict_union_hull_ensemble(
            raw_sample=raw_sample,
            colmap_rec=colmap_rec,
            s1_runs=s1_runs,
            stage2_cfg=stage2_cfg,
            device=device,
            preprocess_cache=preprocess_cache,
            timings_accum=timings,
            order_id=order_id,
        )
        if hull_valid and union_results:
            n_hull_valid = n_ensemble
            for i, (v, e, l) in enumerate(union_results):
                if len(v) > 0 and len(e) > 0:
                    seed_results.append((
                        np.asarray(v, dtype=float), list(e),
                        np.asarray(l, dtype=np.float32),
                    ))
                else:
                    # This slot's stage-2 returned nothing — fall back to its
                    # stage-1 prediction so it still contributes to selection.
                    s1_v, s1_e_arr, s1_l = s1_runs[i]
                    s1_e_list = [(int(a), int(b)) for a, b in s1_e_arr.tolist()]
                    if len(s1_v) > 0 and len(s1_e_list) > 0:
                        seed_results.append((
                            np.asarray(s1_v, dtype=float), s1_e_list,
                            np.asarray(s1_l, dtype=np.float32),
                        ))
        # If hull was invalid OR union path produced nothing, fall through
        # to per-seed below so we don't silently lose this sample.

    if not seed_results:
        for s1_verts, s1_edges_arr, s1_logits in s1_runs:
            s2_verts, s2_edges, s2_logits, hull_valid = _run_stage2_for_seed(
                raw_sample=raw_sample,
                colmap_rec=colmap_rec,
                pred_verts=s1_verts,
                pred_edges=s1_edges_arr,
                stage2_cfg=stage2_cfg,
                device=device,
                preprocess_cache=preprocess_cache,
                timings_accum=timings,
                order_id=order_id,
            )
            if hull_valid:
                n_hull_valid += 1
            if (s2_verts is not None and s2_edges is not None
                    and len(s2_verts) > 0 and len(s2_edges) > 0):
                seed_results.append((
                    np.asarray(s2_verts, dtype=float),
                    list(s2_edges),
                    np.asarray(s2_logits, dtype=np.float32),
                ))
            else:
                s1_e_list = [(int(a), int(b)) for a, b in s1_edges_arr.tolist()]
                if len(s1_verts) > 0 and len(s1_e_list) > 0:
                    seed_results.append((
                        np.asarray(s1_verts, dtype=float),
                        s1_e_list,
                        np.asarray(s1_logits, dtype=np.float32),
                    ))

    # 3) Fuse / select. When n_ensemble=1 the single run is returned verbatim.
    if not seed_results:
        if stats is not None:
            stats["empty_fallback"] = stats.get("empty_fallback", 0) + 1
        if diagnostics is not None:
            diagnostics["members"] = []
            diagnostics["final"] = empty_solution()
        ev, ee = empty_solution()
        return np.asarray(ev, dtype=float), [(int(a), int(b)) for a, b in ee]

    if diagnostics is not None:
        diagnostics["members"] = [
            (np.asarray(v, dtype=float), list(e), np.asarray(l, dtype=np.float32))
            for v, e, l in seed_results
        ]
        diagnostics["s1_runs"] = [
            (np.asarray(v, dtype=float), list(e), np.asarray(l, dtype=np.float32))
            for v, e, l in s1_runs
        ]
        diagnostics["n_hull_valid"] = int(n_hull_valid)

    if n_ensemble == 1:
        fused_v, fused_e = seed_results[0][0], seed_results[0][1]
        if diagnostics is not None:
            diagnostics["single"] = (np.asarray(fused_v, dtype=float), list(fused_e))
    else:
        selector = str(getattr(stage2_cfg, "ensemble_selector", "legacy"))
        mode = str(getattr(stage2_cfg, "ensemble_mode", "medoid"))
        tau = float(getattr(stage2_cfg, "ensemble_merge_m", 0.4))
        refine = bool(getattr(stage2_cfg, "ensemble_refine_positions", True))
        top_k = int(getattr(stage2_cfg, "ensemble_top_k", 0))
        edge_vote = bool(getattr(stage2_cfg, "ensemble_edge_vote", False))
        edge_vote_frac = float(getattr(stage2_cfg, "ensemble_edge_vote_frac", 0.5))
        vertex_vote_frac = float(getattr(stage2_cfg, "ensemble_vertex_vote_frac", 0.0))
        agg = str(getattr(stage2_cfg, "ensemble_refine_agg", "mean"))
        ranking: Optional[List[int]] = None
        if mode == "medoid" or top_k > 0 or selector == "ranker":
            ranking = _medoid_rank(seed_results, match_tau_m=tau)
        medoid_pick = (ranking[0] if ranking else
                       _medoid_select_idx(seed_results, match_tau_m=tau))
        conf_pick = _confidence_select_idx(seed_results)

        if diagnostics is not None:
            diagnostics["single"] = (
                np.asarray(seed_results[0][0], dtype=float),
                list(seed_results[0][1]),
            )
            diagnostics["medoid_pick"] = (
                np.asarray(seed_results[medoid_pick][0], dtype=float),
                list(seed_results[medoid_pick][1]),
            )
            diagnostics["confidence_pick"] = (
                np.asarray(seed_results[conf_pick][0], dtype=float),
                list(seed_results[conf_pick][1]),
            )
            diagnostics["medoid_idx"] = int(medoid_pick)
            diagnostics["confidence_idx"] = int(conf_pick)

        if selector in ("ranker", "fixed"):
            reproj_views_cached = _build_reprojection_views(
                raw_sample, colmap_rec,
            ) if selector == "ranker" else None
            candidates = _build_ranker_candidates(
                seed_results, match_tau_m=tau, edge_vote_frac=edge_vote_frac,
                ranking=ranking, medoid_idx=medoid_pick,
                confidence_idx=conf_pick,
                reproj_views=reproj_views_cached,
                s1_runs=s1_runs,
            )
            if candidates:
                if selector == "fixed":
                    wanted = str(getattr(
                        stage2_cfg, "ensemble_fixed_candidate",
                        "confidence_mean_add"))
                    fallback = "medoid_mean_add"
                    scores = [
                        (2.0 if c["name"] == wanted else
                         1.0 if c["name"] == fallback else 0.0)
                        for c in candidates
                    ]
                else:
                    ranker = getattr(stage2_cfg, "ensemble_ranker", None)
                    scores = [
                        _score_ranker_candidate(c["features"], ranker)
                        for c in candidates
                    ]
                scores_arr = np.asarray(scores, dtype=np.float64)
                best_c = int(np.argmax(scores_arr))
                chosen = candidates[best_c]
                topk_k = int(getattr(stage2_cfg, "ensemble_topk_fuse", 1))
                if (selector == "ranker" and topk_k > 1
                        and len(candidates) > 1):
                    fused_v, fused_e, _ = _fuse_topk_candidates(
                        candidates, scores_arr, topk_k,
                        match_tau_m=tau, edge_vote_frac=edge_vote_frac,
                    )
                else:
                    fused_v = np.asarray(chosen["verts"], dtype=np.float64)
                    fused_e = list(chosen["edges"])
                if diagnostics is not None:
                    diagnostics["candidates"] = [
                        {
                            "name": str(c["name"]),
                            "features": np.asarray(c["features"], dtype=np.float64),
                            "score": float(scores[j]),
                            "verts": np.asarray(c["verts"], dtype=float),
                            "edges": list(c["edges"]),
                        }
                        for j, c in enumerate(candidates)
                    ]
                    diagnostics["selected_candidate"] = str(chosen["name"])
                    diagnostics["selected_candidate_score"] = float(scores[best_c])
                    diagnostics["pick_idx"] = int(chosen.get("source_idx", -1))
                    diagnostics["topk_fuse_k"] = int(topk_k)
            else:
                fused_v = np.asarray(seed_results[medoid_pick][0], dtype=np.float64)
                fused_e = list(seed_results[medoid_pick][1])
                if diagnostics is not None:
                    diagnostics["pick_idx"] = int(medoid_pick)
        elif mode == "consensus":
            # Consensus already aggregates across runs — refinement is
            # logically subsumed. Vertex / edge vote fractions come straight
            # from the same flags (vertex frac defaults to 0.5 here since
            # consensus needs *some* support bar, not 0).
            v_frac_consensus = vertex_vote_frac if vertex_vote_frac > 0.0 else 0.5
            e_frac_consensus = edge_vote_frac if edge_vote else 0.5
            fused_v, fused_e = _fuse_wireframes(
                list(seed_results),
                n_total=n_ensemble, merge_tau_m=tau,
                vertex_vote_frac=v_frac_consensus,
                edge_vote_frac=e_frac_consensus,
                aggregator=agg,
            )
        else:
            if mode == "confidence":
                pick = conf_pick
            else:  # "medoid"
                pick = medoid_pick

            if diagnostics is not None:
                diagnostics["pick_idx"] = int(pick)

            if top_k > 0 and ranking:
                members: Optional[List[int]] = list(ranking[:max(1, top_k)])
                if pick not in members:
                    members.append(pick)
            else:
                members = None  # use all runs

            keep_idx: Optional[np.ndarray] = None
            if refine:
                fused_v, fused_e, keep_idx = _refine_positions(
                    seed_results, pick, match_tau_m=tau,
                    aggregator=agg, member_indices=members,
                    vertex_vote_frac=vertex_vote_frac,
                )
            else:
                fused_v = np.asarray(seed_results[pick][0], dtype=np.float64).reshape(-1, 3)
                fused_e = list(seed_results[pick][1])
                keep_idx = np.arange(fused_v.shape[0], dtype=np.int64)

            if edge_vote:
                fused_e = _vote_edges(
                    seed_results, pick, match_tau_m=tau,
                    vote_frac=edge_vote_frac,
                    member_indices=members,
                    keep_idx=keep_idx,
                )

    if stats is not None:
        if n_hull_valid == 0:
            stats["stage1_fallback"] = stats.get("stage1_fallback", 0) + 1
        else:
            stats["stage2_used"] = stats.get("stage2_used", 0) + 1

    if len(fused_v) == 0 or len(fused_e) == 0:
        if stats is not None:
            stats["empty_fallback"] = stats.get("empty_fallback", 0) + 1
        if diagnostics is not None:
            diagnostics["final"] = empty_solution()
        ev, ee = empty_solution()
        return np.asarray(ev, dtype=float), [(int(a), int(b)) for a, b in ee]
    if diagnostics is not None:
        diagnostics["final"] = (np.asarray(fused_v, dtype=float), list(fused_e))
    return np.asarray(fused_v, dtype=float), fused_e


# ---------------------------------------------------------------------------
# Dataset loading (mirrors empty_submission_2026)
# ---------------------------------------------------------------------------

def load_test_dataset(params: Dict[str, Any]):
    if load_dataset is None:
        raise SystemExit("[data] missing Python package 'datasets'")
    data_path_test_server = Path("/tmp/data")
    if data_path_test_server.exists():
        print("[data] test environment detected (/tmp/data exists)", flush=True)
    else:
        print(f"[data] local run — snapshot_download '{params['dataset']}' → /tmp/data",
              flush=True)
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id=params["dataset"],
            local_dir="/tmp/data",
            repo_type="dataset",
        )
    data_path = data_path_test_server

    data_files = {
        "validation": [str(p) for p in data_path.rglob("*public*/**/*.tar")],
        "test":       [str(p) for p in data_path.rglob("*private*/**/*.tar")],
    }
    print(f"[data] resolved data_files: "
          f"{ {k: len(v) for k, v in data_files.items()} } shards", flush=True)

    return load_dataset(
        str(data_path / "hoho22k_2026_test_x_anon.py"),
        data_files=data_files,
        trust_remote_code=True,
        writer_batch_size=100,
    )


# ---------------------------------------------------------------------------
# Inference loop with prefetched preprocessing
# ---------------------------------------------------------------------------

def predict_split(
    split_ds,
    model: WireframeDiffusion,
    device: torch.device,
    n_pts: int,
    use_depth: bool,
    n_sample_steps: int,
    validity_thresh: float,
    n_prefetch: int,
    stage2_cfg: Stage2Config,
) -> List[Dict[str, Any]]:
    """Sequential GPU inference; CPU preprocessing prefetched n_prefetch ahead.

    Stage 1 runs on the V10 scene, its prediction defines the hull, the
    stage-2 scene is built from the hull-cropped raw points, and stage 2
    produces the final wireframe. The CPU prefetch holds onto the raw sample
    + decoded COLMAP record so the stage-2 builder doesn't re-parse the
    colmap zip.
    """
    results: List[Dict[str, Any]] = []
    use_stage2 = stage2_cfg is not None

    pool = ThreadPoolExecutor(max_workers=max(1, n_prefetch))
    pending: "queue.Queue[Tuple[Any, Any]]" = queue.Queue()

    sample_iter = iter(split_ds)

    def submit_next() -> bool:
        try:
            sample = next(sample_iter)
        except StopIteration:
            return False
        order_id = sample.get("order_id")
        fut = pool.submit(preprocess_sample, sample, n_pts, use_depth,
                          0.1, use_stage2)
        pending.put((order_id, fut))
        return True

    for _ in range(max(1, n_prefetch)):
        if not submit_next():
            break

    pbar = tqdm(desc="predict", unit="sample")
    while not pending.empty():
        order_id, fut = pending.get()
        submit_next()

        t0 = time.perf_counter()
        try:
            pre = fut.result()
            if use_stage2:
                scene = pre["scene"]
                raw_sample = pre["raw_sample"]
                colmap_rec = pre["colmap_rec"]
                cache = pre.get("cache")
                verts, edges = predict_two_stage(
                    raw_sample=raw_sample,
                    colmap_rec=colmap_rec,
                    stage1_scene=scene,
                    stage1_model=model,
                    stage1_n_steps=n_sample_steps,
                    stage1_validity_thresh=validity_thresh,
                    stage2_cfg=stage2_cfg,
                    device=device,
                    preprocess_cache=cache,
                )
                verts = np.asarray(verts, dtype=float)
            else:
                scene = pre
                batch = scene_to_batch(scene, device)
                with torch.inference_mode():
                    verts, edges = model.predict_wireframe(
                        batch,
                        n_steps=n_sample_steps,
                        validity_thresh=validity_thresh,
                    )
                verts = np.asarray(verts, dtype=float)
                edges = [(int(a), int(b)) for a, b in edges]
            dt = time.perf_counter() - t0
            tag = "stage2" if use_stage2 else "stage1"
            print(f"[predict/{tag}] {order_id}: {len(verts)}v {len(edges)}e  ({dt:.2f}s)",
                  flush=True)
        except Exception as exc:
            print(f"[predict] {order_id}: FAILED ({type(exc).__name__}: {exc}) "
                  f"— empty_solution", flush=True)
            verts, edges = empty_solution()
            edges = [(int(a), int(b)) for a, b in edges]

        results.append({
            "order_id":    order_id,
            "wf_vertices": verts.tolist(),
            "wf_edges":    edges,
        })
        pbar.update(1)

    pbar.close()
    pool.shutdown(wait=True)
    return results


# ---------------------------------------------------------------------------
# Local sanity check (--sanity): val split, scored against GT
# ---------------------------------------------------------------------------

def run_sanity_check(
    model: WireframeDiffusion,
    device: torch.device,
    n_pts: int,
    use_depth: bool,
    n_sample_steps: int,
    validity_thresh: float,
    n_samples: int,
    split: str,
    dataset_id: str,
    cache_dir: str,
    stage2_cfg: Stage2Config,
    ensemble_diagnostics: bool = False,
    candidate_dump_path: Optional[str] = None,
) -> float:
    """Run the model on a handful of val samples with GT and report HSS.

    Streams from `dataset_id` (default the public hoho22k trainval split)
    so we never need /tmp/data or the anonymised test shards. Returns
    mean HSS across the evaluated samples.
    """
    import numpy as np
    from hoho2025.metric_helper import hss as hss_fn, HSSReturnType

    use_all = n_samples <= 0
    n_label = "all" if use_all else str(n_samples)
    print(f"[sanity] dataset={dataset_id}  split={split}  n={n_label}  "
          f"n_pts={n_pts}  use_depth={use_depth}  "
          f"n_sample_steps={n_sample_steps}  validity_thresh={validity_thresh}",
          flush=True)

    t_load = time.perf_counter()
    if load_dataset is None:
        raise SystemExit("[sanity] missing Python package 'datasets'")
    ds = load_dataset(
        dataset_id,
        streaming=True,
        trust_remote_code=True,
        cache_dir=str(Path(cache_dir) / "hf"),
    )
    print(f"[sanity] dataset ready in {time.perf_counter() - t_load:.1f}s  "
          f"(available splits: {list(ds.keys())})", flush=True)
    if split not in ds:
        raise SystemExit(f"[sanity] split '{split}' not in dataset; "
                         f"available: {list(ds.keys())}")
    split_ds = ds[split]

    hss_scores: List[float] = []
    f1_scores:  List[float] = []
    iou_scores: List[float] = []
    diag_hss: Dict[str, List[float]] = {}
    diag_counts: Dict[str, List[int]] = {}
    candidate_dump_f = open(candidate_dump_path, "w") if candidate_dump_path else None
    if candidate_dump_f is not None:
        print(f"[ranker/dump] writing candidates to {candidate_dump_path}", flush=True)
    t_pre_list: List[float] = []
    t_enc_list: List[float] = []
    t_den_list: List[float] = []

    def _sync():
        if device.type == "cuda":
            torch.cuda.synchronize()

    use_stage2 = stage2_cfg is not None
    if use_stage2:
        print(f"[sanity] STAGE 2 enabled  n_pts={stage2_cfg.n_pts}  "
              f"n_sample_steps={stage2_cfg.n_sample_steps}  "
              f"hull=(+{stage2_cfg.above_m:.2f}/-{stage2_cfg.below_m:.2f}/"
              f"±{stage2_cfg.side_m:.2f})", flush=True)
        if ensemble_diagnostics and stage2_cfg.ensemble_n > 1:
            print("[diag] enabled: scoring single/member0, medoid pick, "
                  "confidence pick, final fused, and oracle-best member",
                  flush=True)

    t_s2_list: List[float] = []
    t_s2_build_list: List[float] = []
    t_s2_xform_list: List[float] = []
    t_s2_enc_list:   List[float] = []
    t_s2_den_list:   List[float] = []
    stage2_stats: Dict[str, int] = {}

    for i, sample in enumerate(split_ds):
        if not use_all and i >= n_samples:
            break
        order_id = sample.get("order_id", i)
        t_pre = t_enc = t_den = t_s2 = float("nan")
        s2_timings: Dict[str, float] = {}
        t_total0 = time.perf_counter()
        try:
            diag: Dict[str, Any] = {}
            # 1) Preprocessing (CPU) + H2D transfer
            t0 = time.perf_counter()
            pre = preprocess_sample(sample, n_pts=n_pts, use_depth=use_depth,
                                    keep_raw=use_stage2)
            scene = pre["scene"] if use_stage2 else pre
            batch = scene_to_batch(scene, device)
            _sync()
            t_pre = time.perf_counter() - t0

            with torch.inference_mode():
                # 2) Scene encoder
                t0 = time.perf_counter()
                scene_feats, scene_xyz_norm = model._encode_scene(batch)
                _sync()
                t_enc = time.perf_counter() - t0

                # 3) Vertex denoiser — n_sample_steps Euler steps
                K = model.denoiser.k_verts
                B = scene_feats.shape[0]
                step_dt = 1.0 / n_sample_steps
                x = model._init_x0(batch, K, device, B)
                last_logit = torch.zeros(B, K, device=device)
                last_edge_logit = torch.zeros(B, K, K, device=device)
                t0 = time.perf_counter()
                for s in range(n_sample_steps):
                    t = torch.full((B,), s * step_dt, device=device)
                    v, last_logit, last_edge_logit = model.denoiser(
                        x, t, scene_feats, scene_xyz_norm,
                    )
                    x = x + step_dt * v
                _sync()
                t_den = time.perf_counter() - t0

                # 4) Postprocess → world-space verts/edges (mirrors
                # WireframeDiffusion.predict_wireframe; not timed).
                xyz_norm = x[0]
                valid = (last_logit[0] > validity_thresh)
                edge_logit = last_edge_logit[0].float()
                valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
                xyz_valid = xyz_norm.index_select(0, valid_idx)
                center = batch["bbox_center"][0].to(xyz_valid)
                scale = batch["bbox_scale"][0].to(xyz_valid)
                bbox_R = batch.get("bbox_R")
                if isinstance(bbox_R, torch.Tensor):
                    R = bbox_R[0].to(xyz_valid)
                    verts_world = (xyz_valid * scale) @ R + center
                else:
                    verts_world = xyz_valid * scale + center
                pred_v = verts_world.cpu().numpy()
                pred_e: List[Tuple[int, int]] = []
                n_valid = int(valid_idx.numel())
                if n_valid >= 2:
                    sub = edge_logit.index_select(0, valid_idx).index_select(1, valid_idx)
                    tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device,
                                                dtype=torch.bool), diagonal=1)
                    pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False)
                    pred_e = [(int(a), int(b)) for a, b in pairs.detach().cpu().tolist()]

                # 5) Stage 2 (optional) — build cropped scene + refine. The
                # `s2_timings` dict gets populated with per-phase wallclock
                # so we can print stage-2 with the same breakdown as stage 1.
                if use_stage2:
                    t0 = time.perf_counter()
                    pred_v_s2, pred_e_s2 = predict_two_stage(
                        raw_sample=pre["raw_sample"],
                        colmap_rec=pre["colmap_rec"],
                        stage1_scene=scene,
                        stage1_model=model,
                        stage1_n_steps=n_sample_steps,
                        stage1_validity_thresh=validity_thresh,
                        stage2_cfg=stage2_cfg,
                        device=device,
                        preprocess_cache=pre.get("cache"),
                        stats=stage2_stats,
                        timings=s2_timings,
                        diagnostics=(diag if ensemble_diagnostics else None),
                    )
                    _sync()
                    t_s2 = time.perf_counter() - t0
                    pred_v = np.asarray(pred_v_s2, dtype=float)
                    pred_e = list(pred_e_s2)
        except Exception as exc:
            print(f"[sanity] {order_id}: PREDICT FAILED "
                  f"({type(exc).__name__}: {exc})", flush=True)
            pred_v, pred_e = empty_solution()

        gt_v = np.asarray(sample["wf_vertices"])
        gt_e = np.asarray(sample["wf_edges"])

        def _score_diag_variant(name: str,
                                verts: np.ndarray,
                                edges: List[Tuple[int, int]]) -> float:
            try:
                r = hss_fn(verts, edges, gt_v, gt_e)
                val = float(r.hss)
            except Exception:
                val = 0.0
            diag_hss.setdefault(name, []).append(val)
            diag_counts.setdefault(f"{name}_v", []).append(int(len(verts)))
            diag_counts.setdefault(f"{name}_e", []).append(int(len(edges)))
            return val

        diag_line = ""
        if use_stage2 and ensemble_diagnostics and "diag" in locals() and diag:
            scored: Dict[str, float] = {}
            for name in ("single", "medoid_pick", "confidence_pick", "final"):
                if name in diag:
                    vv, ee = diag[name]
                    scored[name] = _score_diag_variant(
                        name, np.asarray(vv, dtype=float), list(ee))
            member_scores: List[float] = []
            for vv, ee, _ll in diag.get("members", []):
                member_scores.append(_score_diag_variant(
                    "member", np.asarray(vv, dtype=float), list(ee)))
            if member_scores:
                oracle = max(member_scores)
                diag_hss.setdefault("oracle", []).append(float(oracle))
                best_idx = int(np.argmax(np.asarray(member_scores, dtype=np.float64)))
                final_h = scored.get("final", 0.0)
                med_h = scored.get("medoid_pick", 0.0)
                conf_h = scored.get("confidence_pick", 0.0)
                single_h = scored.get("single", 0.0)
                diag_line = (
                    f"  | diag single={single_h:.3f} medoid={med_h:.3f} "
                    f"conf={conf_h:.3f} final={final_h:.3f} "
                    f"oracle={oracle:.3f}@{best_idx} gap={oracle-final_h:+.3f}"
                )
            if candidate_dump_f is not None and diag.get("members"):
                candidates = diag.get("candidates")
                if not candidates:
                    members_for_dump = [
                        (np.asarray(vv, dtype=float), list(ee),
                         np.asarray(ll, dtype=np.float32))
                        for vv, ee, ll in diag.get("members", [])
                    ]
                    s1_for_dump = [
                        (np.asarray(vv, dtype=float), list(ee),
                         np.asarray(ll, dtype=np.float32))
                        for vv, ee, ll in diag.get("s1_runs", [])
                    ]
                    candidates = _build_ranker_candidates(
                        members_for_dump,
                        match_tau_m=stage2_cfg.ensemble_merge_m,
                        edge_vote_frac=stage2_cfg.ensemble_edge_vote_frac,
                        medoid_idx=diag.get("medoid_idx"),
                        confidence_idx=diag.get("confidence_idx"),
                        reproj_views=_build_reprojection_views(
                            pre["raw_sample"], pre["colmap_rec"],
                        ),
                        s1_runs=s1_for_dump,
                    )
                selected_name = str(diag.get("selected_candidate", ""))
                for cand in candidates:
                    cv = np.asarray(cand["verts"], dtype=float)
                    ce = list(cand["edges"])
                    try:
                        cr = hss_fn(cv, ce, gt_v, gt_e)
                        chss = float(cr.hss)
                        cf1 = float(cr.f1)
                        ciou = float(cr.iou)
                    except Exception:
                        chss, cf1, ciou = 0.0, 0.0, 0.0
                    row = {
                        "order_id": order_id,
                        "sample_index": int(i),
                        "name": str(cand["name"]),
                        "selected": bool(str(cand["name"]) == selected_name),
                        "hss": chss,
                        "f1": cf1,
                        "iou": ciou,
                        "n_v": int(len(cv)),
                        "n_e": int(len(ce)),
                        "features": np.asarray(
                            cand["features"], dtype=np.float64).tolist(),
                    }
                    candidate_dump_f.write(json.dumps(row, separators=(",", ":")) + "\n")
                candidate_dump_f.flush()

        try:
            result = hss_fn(pred_v, pred_e, gt_v, gt_e)
        except Exception as exc:
            print(f"[sanity] {order_id}: HSS FAILED "
                  f"({type(exc).__name__}: {exc})", flush=True)
            result = HSSReturnType(hss=0.0, f1=0.0, iou=0.0)

        hss_scores.append(float(result.hss))
        f1_scores.append(float(result.f1))
        iou_scores.append(float(result.iou))
        if t_pre == t_pre: t_pre_list.append(t_pre)  # filter NaNs
        if t_enc == t_enc: t_enc_list.append(t_enc)
        if t_den == t_den: t_den_list.append(t_den)
        if t_s2 == t_s2:   t_s2_list.append(t_s2)
        if use_stage2:
            if "s2_build"     in s2_timings: t_s2_build_list.append(s2_timings["s2_build"])
            if "s2_transform" in s2_timings: t_s2_xform_list.append(s2_timings["s2_transform"])
            if "s2_enc"       in s2_timings: t_s2_enc_list.append(s2_timings["s2_enc"])
            if "s2_den"       in s2_timings: t_s2_den_list.append(s2_timings["s2_den"])

        dt = time.perf_counter() - t_total0
        if use_stage2 and t_s2 == t_s2:
            s2_steps = stage2_cfg.n_sample_steps
            s2_str = (f"  s2={t_s2:.2f}s "
                      f"(build={s2_timings.get('s2_build', float('nan')):.2f}s "
                      f"xform={s2_timings.get('s2_transform', float('nan')):.3f}s "
                      f"enc={s2_timings.get('s2_enc', float('nan')):.3f}s "
                      f"den[{s2_steps}]={s2_timings.get('s2_den', float('nan')):.3f}s)")
        else:
            s2_str = ""
        print(f"[sanity] [{i:03d}] {order_id}  pred={len(pred_v)}v/{len(pred_e)}e  "
              f"gt={len(gt_v)}v/{len(gt_e)}e  "
              f"HSS={result.hss:.3f}  F1={result.f1:.3f}  IoU={result.iou:.3f}  "
              f"| pre={t_pre:.2f}s  enc={t_enc:.3f}s  "
              f"den[{n_sample_steps}]={t_den:.3f}s{s2_str}  total={dt:.2f}s"
              f"{diag_line}", flush=True)

    def _mean_std(xs: List[float]) -> Tuple[float, float]:
        if not xs:
            return 0.0, 0.0
        arr = np.asarray(xs, dtype=np.float64)
        return float(arr.mean()), float(arr.std())

    n = len(hss_scores)
    hss_m,  hss_s  = _mean_std(hss_scores)
    f1_m,   f1_s   = _mean_std(f1_scores)
    iou_m,  iou_s  = _mean_std(iou_scores)
    print("=" * 60, flush=True)
    print(f"[sanity] samples={n}", flush=True)
    print(f"[sanity] HSS  mean={hss_m:.4f}  std={hss_s:.4f}", flush=True)
    print(f"[sanity] F1   mean={f1_m:.4f}  std={f1_s:.4f}", flush=True)
    print(f"[sanity] IoU  mean={iou_m:.4f}  std={iou_s:.4f}", flush=True)
    if diag_hss:
        print("[diag] HSS means:", flush=True)
        for name in ("single", "member", "medoid_pick", "confidence_pick",
                     "final", "oracle"):
            if name not in diag_hss:
                continue
            m, s = _mean_std(diag_hss[name])
            print(f"[diag] {name:15s} mean={m:.4f}  std={s:.4f}", flush=True)
        if "oracle" in diag_hss and "final" in diag_hss:
            gap = (np.asarray(diag_hss["oracle"], dtype=np.float64)
                   - np.asarray(diag_hss["final"], dtype=np.float64))
            print(f"[diag] oracle-final   mean={float(gap.mean()):+.4f}  "
                  f"std={float(gap.std()):.4f}", flush=True)
        for name in ("single", "medoid_pick", "confidence_pick", "final"):
            vk, ek = f"{name}_v", f"{name}_e"
            if vk in diag_counts and ek in diag_counts:
                v_m, v_s = _mean_std([float(x) for x in diag_counts[vk]])
                e_m, e_s = _mean_std([float(x) for x in diag_counts[ek]])
                print(f"[diag] {name:15s} size={v_m:.1f}±{v_s:.1f}v/"
                      f"{e_m:.1f}±{e_s:.1f}e", flush=True)
    if use_all:
        print(f"[sanity] full split evaluated: {n} samples", flush=True)
    if t_pre_list:
        pre_m, pre_s = _mean_std(t_pre_list)
        enc_m, enc_s = _mean_std(t_enc_list)
        den_m, den_s = _mean_std(t_den_list)
        per_step_m = (den_m / n_sample_steps) if n_sample_steps else 0.0
        per_step_s = (den_s / n_sample_steps) if n_sample_steps else 0.0
        print(f"[sanity] preprocess   mean={pre_m:.3f}s  std={pre_s:.3f}s", flush=True)
        print(f"[sanity] encoder      mean={enc_m:.3f}s  std={enc_s:.3f}s", flush=True)
        print(f"[sanity] denoiser[{n_sample_steps}] mean={den_m:.3f}s  std={den_s:.3f}s "
              f"({per_step_m*1000:.1f}±{per_step_s*1000:.1f} ms/step)",
              flush=True)
    if use_stage2 and t_s2_list:
        s2_m, s2_s = _mean_std(t_s2_list)
        print(f"[sanity] stage-2 total       mean={s2_m:.3f}s  std={s2_s:.3f}s",
              flush=True)
        if t_s2_build_list:
            b_m, b_s = _mean_std(t_s2_build_list)
            print(f"[sanity] stage-2 build      mean={b_m:.3f}s  std={b_s:.3f}s",
                  flush=True)
        if t_s2_xform_list:
            x_m, x_s = _mean_std(t_s2_xform_list)
            print(f"[sanity] stage-2 transform  mean={x_m:.3f}s  std={x_s:.3f}s",
                  flush=True)
        if t_s2_enc_list:
            e_m, e_s = _mean_std(t_s2_enc_list)
            print(f"[sanity] stage-2 encoder    mean={e_m:.3f}s  std={e_s:.3f}s",
                  flush=True)
        if t_s2_den_list:
            d_m, d_s = _mean_std(t_s2_den_list)
            s2_steps = stage2_cfg.n_sample_steps
            per_m = (d_m / s2_steps) if s2_steps else 0.0
            per_s = (d_s / s2_steps) if s2_steps else 0.0
            print(f"[sanity] stage-2 denoiser[{s2_steps}] mean={d_m:.3f}s  std={d_s:.3f}s "
                  f"({per_m*1000:.1f}±{per_s*1000:.1f} ms/step)",
                  flush=True)
        print(f"[sanity] stage-2 used={stage2_stats.get('stage2_used', 0)}  "
              f"stage-1 fallback={stage2_stats.get('stage1_fallback', 0)}  "
              f"empty fallback={stage2_stats.get('empty_fallback', 0)}",
              flush=True)
    if candidate_dump_f is not None:
        candidate_dump_f.close()
        print(f"[ranker/dump] closed {candidate_dump_path}", flush=True)
    print("=" * 60, flush=True)
    return hss_m


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def _build_stage2_cfg(args, params: Dict[str, Any], device: torch.device) -> Stage2Config:
    """Resolve --stage2_ckpt / env / params, load the Stage2Diffusion model,
    and assemble a ``Stage2Config``. The two-stage pipeline is mandatory;
    ``resolve_stage2_ckpt_path`` raises if no ckpt is configured."""
    ckpt_path = resolve_stage2_ckpt_path(args.stage2_ckpt, params)
    model, s2_args = load_stage2_model(ckpt_path, device)
    n_pts = (args.stage2_n_pts if args.stage2_n_pts is not None
             else s2_args.get("n_pts", 16384))
    k_verts = s2_args.get("k_verts", 64)
    n_sample_steps = (args.stage2_n_sample_steps if args.stage2_n_sample_steps is not None
                      else s2_args.get("n_sample_steps", 20))
    validity_thresh = (args.stage2_validity_thresh if args.stage2_validity_thresh is not None
                       else s2_args.get("validity_thresh", 0.0))
    # Keep stage-2 inference on the original random sampler. Earlier FPS
    # experiments added extra plumbing here, but the current submission/eval
    # path intentionally does not switch sampling modes from CLI or ckpt args.
    sampling = "random"
    # Oversample budgets: random uses a smaller candidate pool; voxel is left
    # as a code path for old experiments but is not selected by this script.
    def _default_oversample(s: str) -> int:
        if s == "voxel":
            return 16_000
        return 10_000
    n_col_oversample = (args.stage2_n_col_oversample
                        if args.stage2_n_col_oversample is not None
                        else _default_oversample(sampling))
    n_dep_oversample = (args.stage2_n_dep_oversample
                        if args.stage2_n_dep_oversample is not None
                        else _default_oversample(sampling))
    fps_max_exact_iters = None
    cfg = Stage2Config(
        model=model,
        n_pts=n_pts,
        k_verts=k_verts,
        n_sample_steps=n_sample_steps,
        validity_thresh=validity_thresh,
        above_m=args.stage2_above_m,
        below_m=args.stage2_below_m,
        side_m=args.stage2_side_m,
        n_col_oversample=n_col_oversample,
        n_dep_oversample=n_dep_oversample,
        n_depth_per_image_cap=args.stage2_n_depth_per_image_cap,
        sampling=sampling,
        density_voxel_size_m=args.stage2_density_voxel_size_m,
        density_kernel_radius=args.stage2_density_kernel_radius,
        density_kernel_axis=args.stage2_density_kernel_axis,
        density_response_power=args.stage2_density_response_power,
        density_planarity_suppression=args.stage2_density_planarity_suppression,
        density_planarity_radius=args.stage2_density_planarity_radius,
        density_planarity_min_points=args.stage2_density_planarity_min_points,
        density_min_per_voxel=args.stage2_density_min_per_voxel,
        ensemble_n=args.ensemble_n,
        ensemble_merge_m=args.ensemble_merge_m,
        ensemble_mode=args.ensemble_mode,
        ensemble_strategy=args.ensemble_strategy,
        ensemble_refine_positions=(not args.no_position_refine),
        ensemble_refine_agg=args.ensemble_refine_agg,
        ensemble_top_k=args.ensemble_top_k,
        ensemble_edge_vote=(not args.no_edge_vote),
        ensemble_edge_vote_frac=args.ensemble_edge_vote_frac,
        ensemble_vertex_vote_frac=args.ensemble_vertex_vote_frac,
        ensemble_stage2_last_k=args.ensemble_stage2_last_k,
        ensemble_selector=args.ensemble_selector,
        ensemble_ranker_path=args.ensemble_ranker_path,
        ensemble_ranker=(
            _load_ensemble_ranker(args.ensemble_ranker_path)
            if args.ensemble_selector == "ranker" else None
        ),
        ensemble_fixed_candidate=args.ensemble_fixed_candidate,
        ensemble_topk_fuse=args.ensemble_topk_fuse,
        fps_max_exact_iters=fps_max_exact_iters,
    )
    extra = ""
    if cfg.sampling == "voxel":
        extra = (f"  density(voxel={cfg.density_voxel_size_m:.2f}m "
                 f"kernel={cfg.density_kernel_axis}:r{cfg.density_kernel_radius} "
                 f"planarity_supp={cfg.density_planarity_suppression:.2f})")
    elif cfg.sampling == "fps":
        extra = (f"  fps_max_exact_iters="
                 f"{'full' if cfg.fps_max_exact_iters is None else cfg.fps_max_exact_iters}")
    print(f"[stage2] enabled  n_pts={cfg.n_pts}  steps={cfg.n_sample_steps}  "
          f"sampling={cfg.sampling}  oversample={cfg.n_col_oversample}+{cfg.n_dep_oversample}  "
          f"hull=(+{cfg.above_m:.2f}/-{cfg.below_m:.2f}/±{cfg.side_m:.2f}){extra}",
          flush=True)
    if cfg.ensemble_n > 1:
        refine_str = (cfg.ensemble_refine_agg
                      if cfg.ensemble_refine_positions else "off")
        topk_str = ("all" if cfg.ensemble_top_k == 0
                    else f"top{cfg.ensemble_top_k}")
        edge_str = (f"vote>={cfg.ensemble_edge_vote_frac:.2f}"
                    if cfg.ensemble_edge_vote else "picked")
        vdrop_str = (f"vote>={cfg.ensemble_vertex_vote_frac:.2f}"
                     if cfg.ensemble_vertex_vote_frac > 0.0 else "off")
        print(f"[ensemble] N={cfg.ensemble_n}  strategy={cfg.ensemble_strategy}  "
              f"mode={cfg.ensemble_mode}  tau={cfg.ensemble_merge_m:.2f}m  "
              f"refine={refine_str}  members={topk_str}  edges={edge_str}  "
              f"vdrop={vdrop_str}  s2_lastk={cfg.ensemble_stage2_last_k}  "
              f"selector={cfg.ensemble_selector}", flush=True)
    return cfg


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="S23DR 2026 submission script (trained model)")
    p.add_argument("--ckpt", default=None,
                   help="Checkpoint path; falls back to $S23DR_CKPT, params['ckpt'], "
                        "then ./test_checkpoint.pth")
    p.add_argument("--params", default="params.json")
    p.add_argument("--n_sample_steps", type=int, default=50,
                   help="Stage-1 diffusion ODE steps at inference (default: 50).")
    p.add_argument("--n_pts", type=int, default=None,
                   help="Stage-1 scene points at inference (default: ckpt args / 8192)")
    p.add_argument("--validity_thresh", type=float, default=0.0,
                   help="Drop predicted vertices below this validity logit "
                        "(default: 0.0)")
    p.add_argument("--n_prefetch", type=int, default=2,
                   help="Number of samples to preprocess ahead of the GPU (default: 2)")
    p.add_argument("--device", default="cuda")
    p.add_argument("--output", default="submission.json")
    p.add_argument("--limit", type=int, default=0,
                   help="If >0, only process the first N samples per split (debug)")
    # ---- stage-2 refinement (optional) -----------------------------------
    p.add_argument("--stage2_ckpt", default=None,
                   help="Stage-2 refinement checkpoint. Falls back to "
                        "$S23DR_STAGE2_CKPT then params['stage2_ckpt']. "
                        "When unset, runs stage-1 only.")
    p.add_argument("--stage2_n_pts", type=int, default=None,
                   help="Stage-2 scene points (default: ckpt args / 16384).")
    p.add_argument("--stage2_n_sample_steps", type=int, default=50,
                   help="Stage-2 ODE steps (default: 50 — 2.5x training, "
                        "inference oversampling).")
    p.add_argument("--stage2_validity_thresh", type=float, default=0.0,
                   help="Stage-2 validity logit threshold (default: 0.0).")
    p.add_argument("--stage2_above_m", type=float, default=0.5,
                   help="Hull inflation upward (metres). Default 0.5.")
    p.add_argument("--stage2_below_m", type=float, default=0.5,
                   help="Hull inflation downward (metres). Default 0.5.")
    p.add_argument("--stage2_side_m", type=float, default=0.5,
                   help="Hull inflation laterally in X/Y (metres). Default 0.5.")
    p.add_argument("--stage2_n_col_oversample", type=int, default=None,
                   help="Stage-2 COLMAP oversample before sampling drops to n_pts//2. "
                        "Default: 10k for random sampling.")
    p.add_argument("--stage2_n_dep_oversample", type=int, default=None,
                   help="Stage-2 depth oversample before sampling drops to n_pts//2. "
                        "Default: 10k for random sampling.")
    p.add_argument("--stage2_n_depth_per_image_cap", type=int, default=30_000,
                   help="Cap on per-image depth pixels considered before the hull test.")
    p.add_argument("--stage2_sampling", choices=["random"], default="random",
                   help="Stage-2 scene sampling mode. The eval/submission path "
                        "is locked to random sampling.")
    p.add_argument("--stage2_fps_max_exact_iters", type=int, default=None,
                   help=argparse.SUPPRESS)
    # ---- stage-2 density-contrast hyperparams (only used when --stage2_sampling=voxel) -
    p.add_argument("--stage2_density_voxel_size_m", type=float, default=0.25)
    p.add_argument("--stage2_density_kernel_radius", type=int, default=3)
    p.add_argument("--stage2_density_kernel_axis",
                   choices=["cube", "x", "y", "z"], default="cube")
    p.add_argument("--stage2_density_response_power", type=float, default=1.0)
    p.add_argument("--stage2_density_planarity_suppression", type=float, default=1.0)
    p.add_argument("--stage2_density_planarity_radius", type=int, default=1)
    p.add_argument("--stage2_density_planarity_min_points", type=int, default=12)
    p.add_argument("--stage2_density_min_per_voxel", type=int, default=0)
    # ---- inference-time ensembling ---------------------------------------
    p.add_argument("--ensemble_n", type=int, default=16,
                   help="Run N stage-1 random-init trajectories and combine "
                        "them per --ensemble_mode. Default 16 (paired with "
                        "--ensemble_strategy union_hull this is ~1.2x wall "
                        "since stage-2 shares one encoder pass at B=N). "
                        "Set to 1 to disable.")
    p.add_argument("--ensemble_merge_m", type=float, default=0.5,
                   help="Vertex match threshold in metres for ensemble "
                        "selection (medoid Hungarian distance, position "
                        "refinement correspondents, consensus clustering). "
                        "Default 0.5 — matches the HSS metric's vert_thresh.")
    p.add_argument("--ensemble_strategy",
                   choices=["union_hull", "per_seed"], default="union_hull",
                   help="How stage-2 is run across the N ensemble members. "
                        "'union_hull' (default) builds ONE stage-2 scene "
                        "from the union of all stage-1 vertex predictions, "
                        "encodes once, and runs a B=N denoiser batch with "
                        "per-slot init_verts. 'per_seed' runs N independent "
                        "(build → encode → denoise) passes; slower but more "
                        "robust if union-hull occasionally fails.")
    p.add_argument("--ensemble_mode",
                   choices=["medoid", "consensus", "confidence"],
                   default="medoid",
                   help="How to combine the N ensemble runs. "
                        "'medoid' (default) picks the run most similar to "
                        "the others by Hungarian-matched distance. "
                        "'confidence' picks the run with the highest mean "
                        "validity logit. 'consensus' merges vertices and "
                        "majority-votes edges (lost recall in earlier tests).")
    p.add_argument("--no_position_refine", action="store_true",
                   help="Disable position refinement. By default, after the "
                        "selector (medoid / confidence) picks a run, each "
                        "vertex of that run is replaced with the per-axis "
                        "aggregate (see --ensemble_refine_agg) of itself "
                        "and its Hungarian-matched correspondents in the "
                        "other N-1 runs (within --ensemble_merge_m). "
                        "Topology unchanged; only vertex positions denoised.")
    p.add_argument("--ensemble_refine_agg", choices=["median", "mean", "wmean"],
                   default="mean",
                   help="Aggregator for position refinement. 'mean' "
                        "(default) has lower variance under iid noise; "
                        "'median' is robust to outlier correspondents; "
                        "'wmean' is a softmax-of-logit weighted mean (each "
                        "correspondent contributes proportional to "
                        "exp(validity_logit), so high-confidence runs "
                        "dominate). 'wmean' only kicks in when validity "
                        "logits are non-trivial — falls back to plain mean "
                        "when weights are uniform or missing.")
    p.add_argument("--ensemble_top_k", type=int, default=0,
                   help="Soft-medoid: restrict the correspondent / voter set "
                        "used by position refinement and edge voting to the "
                        "K most central runs (lowest pairwise Hungarian "
                        "distance). 0 (default) = use all N. Typical values "
                        "for N=16: 5-8. The picked run (medoid) is always "
                        "included.")
    p.add_argument("--no_edge_vote", action="store_true", default=True,
                   help="Keep the picked medoid/confidence run's edges "
                        "verbatim. This is the default.")
    p.add_argument("--edge_vote", action="store_false", dest="no_edge_vote",
                   help="Enable edge majority voting: after selection each "
                        "non-picked run's edges are relabelled into the picked "
                        "run's vertex index space via the same Hungarian match "
                        "used for position refinement, and edges with >= "
                        "--ensemble_edge_vote_frac of votes survive.")
    p.add_argument("--ensemble_edge_vote_frac", type=float, default=0.5,
                   help="Edge vote fraction (default 0.5 = strict majority). "
                        "An edge (i,j) in picked-vertex space is kept iff at "
                        "least ceil(M * frac) runs vote for it, where M is "
                        "the active voter set (top_k if set, else N).")
    p.add_argument("--ensemble_vertex_vote_frac", type=float, default=0.0,
                   help="Vertex consensus drop fraction (default 0.0 = off). "
                        "When > 0, a picked-run vertex survives iff at least "
                        "ceil(M * frac) voter runs have a Hungarian "
                        "correspondent within tau. The picked run itself is "
                        "always one voter (self-vote). Tightens precision by "
                        "dropping spurious one-run-only detections — "
                        "complements --ensemble_edge_vote. Used by both "
                        "medoid/confidence modes (picked-then-refine) and "
                        "consensus mode (sets the cluster-support threshold).")
    p.add_argument("--ensemble_stage2_last_k", type=int, default=1,
                   help="Average the last K stage-2 denoiser steps' "
                        "predictions (positions + validity / edge logits) "
                        "before thresholding. 1 (default) = use only the "
                        "final step (current behaviour). 3-5 mildly smooths "
                        "ODE jitter near the trajectory endpoint at zero "
                        "extra compute.")
    p.add_argument("--ensemble_selector", choices=["legacy", "ranker", "fixed"],
                   default="legacy",
                   help="Final ensemble chooser. 'legacy' (default) keeps the "
                        "medoid/confidence/consensus path with no reranker. "
                        "'ranker' builds "
                        "raw/refined/voted/additive-edge candidates and picks "
                        "with --ensemble_ranker_path if present, otherwise a "
                        "small built-in heuristic. 'fixed' picks a named "
                        "candidate family.")
    p.add_argument("--ensemble_fixed_candidate", default="confidence_mean_add",
                   help="Candidate name used by --ensemble_selector fixed. "
                        "Default confidence_mean_add, falling back to "
                        "medoid_mean_add when the confidence candidate is "
                        "identical to the medoid and therefore absent.")
    p.add_argument("--ensemble_topk_fuse", type=int, default=3,
                   help="Top-K ranker-scored candidates to fuse (mean-refine "
                        "positions + edge-vote union) when "
                        "--ensemble_selector=ranker. 1 disables fusion and "
                        "keeps the argmax candidate.")
    p.add_argument("--ensemble_ranker_path", default="ensemble_ranker_v5.npz",
                   help="NPZ produced by --train_ensemble_ranker_jsonl. Used "
                        "only when --ensemble_selector=ranker. If missing, "
                        "ranker mode falls back to a heuristic scorer.")
    p.add_argument("--ensemble_candidate_dump", default=None,
                   help="In --sanity mode, write ranker candidate features "
                        "and GT HSS labels as JSONL for fast selector training.")
    p.add_argument("--train_ensemble_ranker_jsonl", default=None,
                   help="Train the tiny NumPy ridge ranker from a candidate "
                        "JSONL dump and exit. No model checkpoint/GPU needed.")
    p.add_argument("--ensemble_ranker_out", default="ensemble_ranker.npz",
                   help="Output path for --train_ensemble_ranker_jsonl.")
    p.add_argument("--ensemble_ranker_l2", type=float, default=1e-2,
                   help="Ridge L2 for --train_ensemble_ranker_jsonl.")
    # ---- local sanity check (no submission) ------------------------------
    p.add_argument("--sanity", action="store_true",
                   help="Local pre-submission sanity check: load default ckpt, "
                        "stream a few val samples from the public trainval "
                        "dataset, score against GT, print HSS, and exit. "
                        "Writes no submission.json.")
    p.add_argument("--sanity_n", type=int, default=0,
                   help="Number of val samples to score in --sanity mode. "
                        "0 (default) means score the entire split.")
    p.add_argument("--sanity_split", default="validation",
                   help="Split to draw sanity samples from (default 'validation').")
    p.add_argument("--sanity_dataset", default="usm3d/hoho22k_2026_trainval",
                   help="HF dataset id with GT used for --sanity scoring.")
    p.add_argument("--sanity_cache_dir", default="cache",
                   help="Local HF cache dir for sanity dataset streaming.")
    p.add_argument("--sanity_ensemble_diagnostics", action="store_true",
                   help="In --sanity mode, score ensemble alternatives: "
                        "member0, medoid pick, confidence pick, final fused, "
                        "and oracle-best member. Diagnostic only; normal "
                        "predictions are unchanged.")
    return p.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)

    if args.train_ensemble_ranker_jsonl:
        train_ensemble_ranker_jsonl(
            args.train_ensemble_ranker_jsonl,
            args.ensemble_ranker_out,
            l2=args.ensemble_ranker_l2,
        )
        return

    # In --sanity mode we don't need /tmp/data, params.json, or the
    # anonymised test shards. Just resolve a checkpoint and score against
    # the public trainval GT.
    if args.sanity:
        params: Dict[str, Any] = {}
        if Path(args.params).exists():
            with open(args.params) as f:
                params = json.load(f)
        ckpt_path = resolve_ckpt_path(args.ckpt, params)
        if args.device == "cuda" and not torch.cuda.is_available():
            if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = torch.device("mps")
            else:
                device = torch.device("cpu")
        else:
            device = torch.device(args.device)
        if device.type == "cpu":
            print("[warn] CUDA/MPS unavailable — running on CPU; this will be slow.", flush=True)
        model, ckpt_args = load_model(ckpt_path, device)

        n_pts           = (args.n_pts if args.n_pts is not None
                           else ckpt_args.get("n_pts", 8192))
        use_depth       = ckpt_args.get("use_depth", True)
        n_sample_steps  = (args.n_sample_steps if args.n_sample_steps is not None
                           else ckpt_args.get("n_sample_steps", 50))
        validity_thresh = (args.validity_thresh if args.validity_thresh is not None
                           else ckpt_args.get("validity_thresh", 0.0))

        stage2_cfg = _build_stage2_cfg(args, params, device)

        run_sanity_check(
            model=model,
            device=device,
            n_pts=n_pts,
            use_depth=use_depth,
            n_sample_steps=n_sample_steps,
            validity_thresh=validity_thresh,
            n_samples=args.sanity_n,
            split=args.sanity_split,
            dataset_id=args.sanity_dataset,
            cache_dir=args.sanity_cache_dir,
            stage2_cfg=stage2_cfg,
            ensemble_diagnostics=(
                args.sanity_ensemble_diagnostics
                or args.ensemble_candidate_dump is not None
            ),
            candidate_dump_path=args.ensemble_candidate_dump,
        )
        return

    print("------------ Loading dataset ------------", flush=True)
    with open(args.params) as f:
        params = json.load(f)
    print(params, flush=True)

    print("pwd:", flush=True);  os.system("pwd")
    os.system("ls -lahtr")
    print("/tmp/data/");          os.system("ls -lahtr /tmp/data/ 2>/dev/null")
    print("/tmp/data/data");      os.system("ls -lahtrR /tmp/data/data 2>/dev/null")

    dataset = load_test_dataset(params)
    print(dataset, flush=True)

    # ---- model: loaded ONCE -----------------------------------------------
    ckpt_path = resolve_ckpt_path(args.ckpt, params)
    if args.device == "cuda" and not torch.cuda.is_available():
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = torch.device("mps")
        else:
            device = torch.device("cpu")
    else:
        device = torch.device(args.device)
    if device.type == "cpu":
        print("[warn] CUDA/MPS unavailable — running on CPU; this will be slow.", flush=True)
    model, ckpt_args = load_model(ckpt_path, device)

    n_pts           = (args.n_pts if args.n_pts is not None
                       else ckpt_args.get("n_pts", 8192))
    use_depth       = ckpt_args.get("use_depth", True)
    n_sample_steps  = (args.n_sample_steps if args.n_sample_steps is not None
                       else ckpt_args.get("n_sample_steps", 50))
    validity_thresh = (args.validity_thresh if args.validity_thresh is not None
                       else ckpt_args.get("validity_thresh", 0.0))
    print(f"[infer] n_pts={n_pts}  use_depth={use_depth}  "
          f"n_sample_steps={n_sample_steps}  validity_thresh={validity_thresh}  "
          f"n_prefetch={args.n_prefetch}", flush=True)

    stage2_cfg = _build_stage2_cfg(args, params, device)

    # ---- inference --------------------------------------------------------
    print("------------ Running inference ------------", flush=True)
    solution: List[Dict[str, Any]] = []
    for subset_name in dataset:
        print(f"[predict] subset {subset_name}", flush=True)
        split_ds = dataset[subset_name]
        if args.limit > 0:
            split_ds = split_ds.select(range(min(args.limit, len(split_ds))))
        solution.extend(
            predict_split(
                split_ds,
                model=model,
                device=device,
                n_pts=n_pts,
                use_depth=use_depth,
                n_sample_steps=n_sample_steps,
                validity_thresh=validity_thresh,
                n_prefetch=args.n_prefetch,
                stage2_cfg=stage2_cfg,
            )
        )

    print("------------ Saving results ------------", flush=True)
    with open(args.output, "w") as f:
        json.dump(solution, f)
    print(f"[done] wrote {len(solution)} predictions → {args.output}", flush=True)


if __name__ == "__main__":
    main()