"""WireframeDiffusion: flow-matching wrapper around SceneEncoder + VertexDenoiser.

State is xyz only (3-d). The K vertex slots are unordered; we use Hungarian matching
between fresh noise samples and the GT vertices for each scene to build a permutation-
invariant flow target. Validity and pairwise edges are predicted by separate
logit heads; only xyz is integrated through the ODE.
"""
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
from concurrent.futures import ThreadPoolExecutor

import numpy as np
from scipy.optimize import linear_sum_assignment

GESTALT_TIER1_IDS = (1, 2, 3, 6, 4, 5, 12)
GESTALT_TIER2_IDS = (8, 9, 10, 11, 18, 27)
ADE_HOUSE_FOREGROUND_IDS = (1, 2, 9, 15, 26, 43, 49, 54)

# scipy.linear_sum_assignment is C and releases the GIL — threads parallelise.
# One global pool keeps thread reuse cheap across training iterations.
_HUNGARIAN_POOL = ThreadPoolExecutor(max_workers=int(os.environ.get("S23DR_HUNGARIAN_THREADS", "8")))


def _batched_hungarian(costs: list[np.ndarray]) -> list[tuple[np.ndarray, np.ndarray]]:
    """Run linear_sum_assignment on a list of cost matrices in parallel."""
    if not costs:
        return []
    return list(_HUNGARIAN_POOL.map(linear_sum_assignment, costs))




class WireframeDiffusion(nn.Module):
    def __init__(
        self,
        scene_encoder: nn.Module,
        denoiser: nn.Module,
        loss_flow_weight: float = 1.0,
        loss_endpoint_weight: float = 0.5,
        loss_validity_weight: float = 0.2,
        loss_edge_weight: float = 0.2,
        loss_huber_beta: float = 0.05,
        focal_gamma: float = 2.0,
        focal_alpha: float = 0.25,
        real_slot_weight: float = 1.0,
        null_slot_weight: float = 0.1,
        noise_sigma_xyz: float = 1.0,
        init_from_scene: bool = False,
        init_from_encoder_head: bool = False,
        scene_init_jitter: float = 0.05,
        xyz_clip: float = 4.0,
        pred_clip: float = 10.0,
        validity_logit_clip: float = 20.0,
        loss_iou_weight: float = 0.0,
        loss_iou_t_min: float = 0.8,
        loss_iou_gate_m: float = 0.5,
        loss_iou_step_m: float = 0.05,
        loss_iou_max_samples: int = 512,
        loss_iou_norm: str = "l1",
        loss_flow_norm: str = "smooth_l1",
        loss_endpoint_norm: str = "smooth_l1",
        loss_soft_vertex_f1_weight: float = 0.0,
        loss_soft_edge_f1_weight: float = 0.0,
        loss_count_weight: float = 0.0,
        loss_hss_radius_m: float = 0.5,
        loss_hss_temp_m: float = 0.10,
        encoder_head_supervision_weight: float = 0.0,
        encoder_head_supervision_frac: float = 0.15,
        encoder_head_blend_frac: float = 0.0,
    ):
        super().__init__()
        self.scene_encoder = scene_encoder
        self.denoiser = denoiser
        self.loss_flow_weight = float(loss_flow_weight)
        self.loss_endpoint_weight = float(loss_endpoint_weight)
        self.loss_validity_weight = float(loss_validity_weight)
        self.loss_edge_weight = float(loss_edge_weight)
        self.loss_huber_beta = float(loss_huber_beta)
        self.focal_gamma = float(focal_gamma)
        self.focal_alpha = float(focal_alpha)
        self.real_slot_weight = float(real_slot_weight)
        self.null_slot_weight = float(null_slot_weight)
        self.noise_sigma_xyz = float(noise_sigma_xyz)
        self.init_from_scene = bool(init_from_scene)
        self.init_from_encoder_head = bool(init_from_encoder_head)
        self.scene_init_jitter = float(scene_init_jitter)
        self.xyz_clip = float(xyz_clip)
        self.pred_clip = float(pred_clip)
        self.validity_logit_clip = float(validity_logit_clip)
        self.loss_iou_weight = float(loss_iou_weight)
        self.loss_iou_t_min = float(loss_iou_t_min)
        self.loss_iou_gate_m = float(loss_iou_gate_m)
        self.loss_iou_step_m = float(loss_iou_step_m)
        self.loss_iou_max_samples = int(loss_iou_max_samples)
        self.loss_iou_norm = self._validate_norm("loss_iou_norm", loss_iou_norm)
        self.loss_flow_norm = self._validate_norm("loss_flow_norm", loss_flow_norm)
        self.loss_endpoint_norm = self._validate_norm("loss_endpoint_norm", loss_endpoint_norm)
        self.loss_soft_vertex_f1_weight = float(loss_soft_vertex_f1_weight)
        self.loss_soft_edge_f1_weight = float(loss_soft_edge_f1_weight)
        self.loss_count_weight = float(loss_count_weight)
        self.loss_hss_radius_m = float(loss_hss_radius_m)
        self.loss_hss_temp_m = float(loss_hss_temp_m)
        self.encoder_head_supervision_weight = float(encoder_head_supervision_weight)
        self.encoder_head_supervision_frac = float(encoder_head_supervision_frac)
        self.encoder_head_blend_frac = float(encoder_head_blend_frac)
        self.last_loss_terms: dict[str, torch.Tensor] = {}

    _ALLOWED_NORMS = ("l1", "l2", "smooth_l1")

    @classmethod
    def _validate_norm(cls, name: str, value: str) -> str:
        v = str(value).lower()
        if v not in cls._ALLOWED_NORMS:
            raise ValueError(f"{name} must be one of {cls._ALLOWED_NORMS}; got {value!r}")
        return v

    def _vec_loss(
        self,
        pred: torch.Tensor,    # (..., D)
        target: torch.Tensor,  # (..., D)
        norm: str,
    ) -> torch.Tensor:
        """Per-element distance loss summed over the last dim. Shared by flow,
        endpoint, and iou losses so they all use the same norm-selection logic.

          - l1:        Euclidean displacement |Δ| = sqrt(Δ·Δ + eps); grad bounded by 1
          - l2:        squared displacement Δ·Δ; smooth, punishes outliers more
          - smooth_l1: component-wise Huber (uses self.loss_huber_beta), summed over D;
                       quadratic for small Δ, linear for large Δ
        """
        if norm == "l1":
            diff = pred - target
            return (diff * diff).sum(dim=-1).add(1e-12).sqrt()
        if norm == "l2":
            diff = pred - target
            return (diff * diff).sum(dim=-1)
        # smooth_l1
        return F.smooth_l1_loss(
            pred, target, reduction="none", beta=self.loss_huber_beta,
        ).sum(dim=-1)

    def _weighted_mean(self, values: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
        denom = weights.sum().clamp_min(1.0)
        return (values * weights).sum() / denom

    def _hss_norm_thresholds(
        self,
        bbox_scale: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        scale = bbox_scale.to(dtype=torch.float32).clamp_min(1e-6)
        radius = (self.loss_hss_radius_m / scale).view(-1, 1)
        temp = (self.loss_hss_temp_m / scale).view(-1, 1).clamp_min(1e-6)
        return radius, temp

    def _soft_vertex_f1_loss(
        self,
        x1_pred: torch.Tensor,       # (B, K, 3)
        gt_xyz_full: torch.Tensor,   # (B, K, 3)
        real_mask: torch.Tensor,     # (B, K)
        valid_prob: torch.Tensor,    # (B, K)
        bbox_scale: torch.Tensor,    # (B,)
    ) -> torch.Tensor:
        """Differentiable proxy for HSS corner F1 at the metric radius."""
        zero = x1_pred.sum() * 0.0
        B, K, _ = x1_pred.shape
        if K == 0:
            return zero

        radius, temp = self._hss_norm_thresholds(bbox_scale)
        dist = torch.cdist(x1_pred.float(), gt_xyz_full.float())
        match = torch.sigmoid((radius[:, None, :] - dist) / temp[:, None, :])
        real = real_mask.float()
        match = match * real[:, None, :]

        pred_match = match.max(dim=2).values
        gt_match = (match * valid_prob[:, :, None]).max(dim=1).values

        pred_count = valid_prob.sum(dim=1)
        gt_count = real.sum(dim=1)
        tp_pred = (valid_prob * pred_match).sum(dim=1)
        tp_gt = gt_match.sum(dim=1)

        precision = tp_pred / pred_count.clamp_min(1e-6)
        recall = tp_gt / gt_count.clamp_min(1e-6)
        soft_f1 = (2.0 * precision * recall) / (precision + recall).clamp_min(1e-6)
        has_gt = gt_count > 0
        if not has_gt.any():
            return zero
        return (1.0 - soft_f1[has_gt]).mean()

    def _soft_edge_f1_loss(
        self,
        x1_pred: torch.Tensor,       # (B, K, 3)
        target_xyz: torch.Tensor,    # (B, K, 3)
        edge_logit: torch.Tensor,    # (B, K, K)
        edge_target: torch.Tensor,   # (B, K, K)
        valid_prob: torch.Tensor,    # (B, K)
        bbox_scale: torch.Tensor,    # (B,)
    ) -> torch.Tensor:
        """Soft edge F1 using edge logits, validity, and endpoint proximity."""
        device = x1_pred.device
        zero = x1_pred.sum() * 0.0
        B, K, _ = x1_pred.shape
        if K < 2:
            return zero

        tri = torch.triu(torch.ones(K, K, device=device, dtype=torch.bool), diagonal=1)
        radius, temp = self._hss_norm_thresholds(bbox_scale)

        endpoint_dist = (x1_pred.float() - target_xyz.float()).norm(dim=-1)
        pair_dist = 0.5 * (endpoint_dist[:, :, None] + endpoint_dist[:, None, :])
        close = torch.sigmoid((radius[:, :, None] - pair_dist) / temp[:, :, None])

        edge_prob = torch.sigmoid(
            edge_logit.float().clamp(-self.validity_logit_clip, self.validity_logit_clip)
        )
        edge_prob = edge_prob * valid_prob[:, :, None] * valid_prob[:, None, :]

        pred = edge_prob[:, tri]
        target = edge_target.float()[:, tri]
        close = close[:, tri]
        tp = (pred * target * close).sum(dim=1)
        pred_count = pred.sum(dim=1)
        gt_count = target.sum(dim=1)

        precision = tp / pred_count.clamp_min(1e-6)
        recall = tp / gt_count.clamp_min(1e-6)
        soft_f1 = (2.0 * precision * recall) / (precision + recall).clamp_min(1e-6)
        has_gt = gt_count > 0
        if not has_gt.any():
            return zero
        return (1.0 - soft_f1[has_gt]).mean()

    def _valid_count_loss(
        self,
        valid_prob: torch.Tensor,
        real_mask: torch.Tensor,
    ) -> torch.Tensor:
        K = valid_prob.shape[1]
        pred_count = valid_prob.sum(dim=1) / max(1, K)
        gt_count = real_mask.float().sum(dim=1) / max(1, K)
        return F.smooth_l1_loss(pred_count, gt_count, reduction="mean")

    def _edge_targets_from_matching(
        self,
        batch: dict,
        gt_to_slot: torch.Tensor,
        slot_real: torch.Tensor,
    ) -> torch.Tensor:
        """Map raw GT edge endpoints onto the matched denoising slots.

        For v11 caches the raw wireframe is filtered by `wf_vertex_kept_mask` /
        `wf_edge_kept_mask` (small chimney components dropped). `gt_to_slot` is
        keyed by the KEPT-vertex index, so raw edges must be dropped if either
        endpoint is discarded and the remaining endpoints remapped from raw to
        kept-vertex index space via cumsum(keep_v) - 1.
        """
        B, K = slot_real.shape
        device = slot_real.device
        edge_target = torch.zeros(B, K, K, device=device, dtype=torch.float32)
        wf_edges = batch.get("wf_edges_raw")
        if not isinstance(wf_edges, (list, tuple)):
            return edge_target

        keep_v_list = batch.get("wf_vertex_kept_mask") or [None] * B
        keep_e_list = batch.get("wf_edge_kept_mask") or [None] * B

        for b, edges in enumerate(wf_edges[:B]):
            if edges is None:
                continue
            edge_t = torch.as_tensor(edges, device=device, dtype=torch.long)
            if edge_t.numel() == 0:
                continue
            edge_t = edge_t.view(-1, 2)

            keep_v = keep_v_list[b] if b < len(keep_v_list) else None
            keep_e = keep_e_list[b] if b < len(keep_e_list) else None

            if keep_v is not None:
                keep_v_t = torch.as_tensor(keep_v, device=device, dtype=torch.bool)
                if keep_e is not None:
                    keep_e_t = torch.as_tensor(keep_e, device=device, dtype=torch.bool)
                    if keep_e_t.numel() == edge_t.shape[0]:
                        edge_t = edge_t[keep_e_t]
                if edge_t.numel() == 0:
                    continue
                # raw -> kept index remap; -1 marks discarded raw verts
                raw_to_kept = torch.cumsum(keep_v_t.long(), dim=0) - 1
                raw_to_kept = torch.where(
                    keep_v_t, raw_to_kept, torch.full_like(raw_to_kept, -1)
                )
                n_raw = raw_to_kept.shape[0]
                in_raw = (
                    (edge_t[:, 0] >= 0) & (edge_t[:, 1] >= 0)
                    & (edge_t[:, 0] < n_raw) & (edge_t[:, 1] < n_raw)
                )
                edge_t = edge_t[in_raw]
                if edge_t.numel() == 0:
                    continue
                e0 = raw_to_kept.index_select(0, edge_t[:, 0])
                e1 = raw_to_kept.index_select(0, edge_t[:, 1])
                kept_ok = (e0 >= 0) & (e1 >= 0)
                e0 = e0[kept_ok]
                e1 = e1[kept_ok]
            else:
                # Pre-v11 cache: raw indices are already the kept indices.
                e0 = edge_t[:, 0]
                e1 = edge_t[:, 1]

            in_range = (e0 >= 0) & (e1 >= 0) & (e0 < K) & (e1 < K)
            e0 = e0[in_range]
            e1 = e1[in_range]
            if e0.numel() == 0:
                continue
            s0 = gt_to_slot[b].index_select(0, e0)
            s1 = gt_to_slot[b].index_select(0, e1)
            ok = (s0 >= 0) & (s1 >= 0) & (s0 != s1)
            s0 = s0[ok]
            s1 = s1[ok]
            if s0.numel() == 0:
                continue
            edge_target[b, s0, s1] = 1.0
            edge_target[b, s1, s0] = 1.0
        return edge_target

    def _edge_iou_loss(
        self,
        x1_pred: torch.Tensor,        # (B, K, 3) predicted clean endpoints
        target_xyz: torch.Tensor,     # (B, K, 3) matched GT endpoints
        edge_target: torch.Tensor,    # (B, K, K) binary GT edge mask
        slot_real: torch.Tensor,      # (B, K)
        t: torch.Tensor,              # (B,) flow time in [0, 1]
        bbox_scale: torch.Tensor,     # (B,) world meters per normalized unit
    ) -> torch.Tensor:
        """Parametric correspondence loss between predicted and GT edges.

        For each gated edge:
          1. N = ceil(L_gt / step)  (per-edge, capped for safety)
          2. Sample N points uniformly along the GT edge at u_i = i/(N-1) ∈ [0, 1]
          3. Sample N points along the PREDICTED edge at the SAME u_i values
          4. Per-edge loss = step · Σ |pred_pt_i − gt_pt_i|  (Riemann sum → length-weighted)

        Length-weighting is automatic: more samples for longer edges → larger sum.
        Multiplying by `step` (constant per scene) gives the integral interpretation
        ∫₀^L |d(s)| ds, comparable across scenes.

        Two gates suppress meaningless gradients: per-edge endpoint distance must be
        < gate_m (world m), and t must be > t_min (only late denoising steps).

        Differentiability: pred sample positions are p0 + u·(p1-p0) with u a fixed
        grid (no grad), so gradients flow through both endpoints. Distance uses
        sqrt(d² + eps) for a finite gradient at d=0 (Euclidean / L1, not L2).
        """
        device = x1_pred.device
        zero = x1_pred.sum() * 0.0  # keeps the autograd graph alive when nothing passes
        K = x1_pred.shape[1]

        tri = torch.triu(torch.ones(K, K, device=device, dtype=torch.bool), diagonal=1)
        pos = edge_target.bool() & tri[None] & slot_real[:, :, None] & slot_real[:, None, :]
        if not pos.any():
            return zero

        idx = pos.nonzero(as_tuple=False)              # (E_all, 3)
        b_idx, i_idx, j_idx = idx[:, 0], idx[:, 1], idx[:, 2]

        p0 = x1_pred[b_idx, i_idx]                     # (E_all, 3) pred endpoints
        p1 = x1_pred[b_idx, j_idx]
        g0 = target_xyz[b_idx, i_idx]                  # (E_all, 3) GT endpoints
        g1 = target_xyz[b_idx, j_idx]

        # Per-scene thresholds in normalized coords (world m / scene scale).
        scale_safe = bbox_scale.clamp_min(1e-6)
        step_norm_all = (self.loss_iou_step_m / scale_safe)[b_idx]    # (E_all,)
        gate_per_edge = (self.loss_iou_gate_m / scale_safe)[b_idx]
        e0 = (p0 - g0).norm(dim=-1)
        e1 = (p1 - g1).norm(dim=-1)
        keep = (e0 < gate_per_edge) & (e1 < gate_per_edge) & (t[b_idx] > self.loss_iou_t_min)
        if not keep.any():
            return zero

        p0, p1, g0, g1 = p0[keep], p1[keep], g0[keep], g1[keep]
        step_norm = step_norm_all[keep]                # (E,)
        E = p0.shape[0]

        # Sample count per edge from GT length (GT defines the reference spacing).
        # Detached: count is integer with no grad path.
        Lg = (g1 - g0).norm(dim=-1).detach()
        N = (Lg / step_norm).ceil().clamp(min=2, max=self.loss_iou_max_samples).long()

        # Flatten (edge, sample_idx) into one (total,) tensor; one CPU↔GPU sync.
        total = int(N.sum().item())
        edge_id = torch.repeat_interleave(torch.arange(E, device=device), N)
        cum = torch.zeros(E + 1, device=device, dtype=torch.long)
        cum[1:] = N.cumsum(0)
        local = torch.arange(total, device=device) - cum[edge_id]
        N_per = N[edge_id].float()
        # Endpoint-inclusive parameterization: u[0]=0, u[N-1]=1.
        u = local.float() / (N_per - 1).clamp_min(1.0)

        # Same parametric u on both segments → one-to-one correspondence.
        pred_pts = p0[edge_id] + u.unsqueeze(-1) * (p1 - p0)[edge_id]
        gt_pts   = g0[edge_id] + u.unsqueeze(-1) * (g1 - g0)[edge_id]

        # Per-sample loss in the configured norm (l1 / l2 / smooth_l1).
        d = self._vec_loss(pred_pts, gt_pts, self.loss_iou_norm)

        sum_d = torch.zeros(E, device=device, dtype=d.dtype).scatter_add_(0, edge_id, d)
        # Per-edge integral ≈ step · Σ d_i ≈ ∫₀^L |d(s)| ds. step is detached → grad
        # flows only through the endpoint coordinates, as intended.
        per_edge = step_norm * sum_d
        return per_edge.mean()

    def _scene_guided_xyz_init(self, batch: dict, k_verts: int, device: torch.device) -> torch.Tensor | None:
        """Sample K starting xyz from structurally useful scene points."""
        scene_xyz = batch.get("scene_xyz")
        if not isinstance(scene_xyz, torch.Tensor) or scene_xyz.ndim != 3:
            return None
        scene_xyz = torch.nan_to_num(scene_xyz, nan=0.0, posinf=0.0, neginf=0.0).to(device)
        B, N, _ = scene_xyz.shape
        if N <= 0:
            return None

        weights = torch.ones(B, N, device=device, dtype=torch.float32)

        type_ids = batch.get("scene_type_ids")
        if isinstance(type_ids, torch.Tensor) and type_ids.shape[:2] == (B, N):
            type_ids = type_ids.to(device)
            weights = torch.where(type_ids == 2, weights * 0.25, weights)  # camera centres
            weights = torch.where(type_ids == 1, weights * 1.5, weights)   # depth samples

        gestalt_ids = batch.get("scene_gestalt_ids")
        if isinstance(gestalt_ids, torch.Tensor) and gestalt_ids.shape[:2] == (B, N):
            gestalt_ids = gestalt_ids.to(device)
            tier1 = torch.zeros(B, N, device=device, dtype=torch.bool)
            tier2 = torch.zeros(B, N, device=device, dtype=torch.bool)
            for gid in GESTALT_TIER1_IDS:
                tier1 |= gestalt_ids == gid
            for gid in GESTALT_TIER2_IDS:
                tier2 |= gestalt_ids == gid

            house = torch.zeros(B, N, device=device, dtype=torch.bool)
            ade_ids = batch.get("scene_ade_ids")
            if isinstance(ade_ids, torch.Tensor) and ade_ids.shape[:2] == (B, N):
                ade_ids = ade_ids.to(device)
                for aid in ADE_HOUSE_FOREGROUND_IDS:
                    house |= ade_ids == aid

            priority = tier1 | (tier2 & house)
            weights = torch.where(priority, weights * 8.0, weights)
            weights = torch.where(gestalt_ids >= 0, weights * 1.5, weights)

        confs = []
        geom_conf = batch.get("scene_geom_conf")
        sem_conf = batch.get("scene_sem_conf")
        if isinstance(geom_conf, torch.Tensor) and geom_conf.shape[:2] == (B, N):
            confs.append(geom_conf.to(device).float().clamp(0.0, 1.0))
        if isinstance(sem_conf, torch.Tensor) and sem_conf.shape[:2] == (B, N):
            confs.append(sem_conf.to(device).float().clamp(0.0, 1.0))
        if confs:
            conf = torch.stack(confs, dim=-1).mean(dim=-1)
            weights = weights * (0.25 + conf)

        idx = torch.multinomial(weights.clamp_min(1e-6), k_verts, replacement=True)
        base = torch.gather(scene_xyz, dim=1, index=idx[..., None].expand(-1, -1, 3))
        if self.scene_init_jitter > 0:
            base = base + self.scene_init_jitter * torch.randn_like(base)
        return base.clamp(-self.xyz_clip, self.xyz_clip)

    @torch.no_grad()
    def _priority_fps_xyz_target(
        self,
        batch: dict,
        k_verts: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Deterministic priority-FPS anchors used as a query-head teacher."""
        scene_xyz = batch.get("scene_xyz")
        if not isinstance(scene_xyz, torch.Tensor) or scene_xyz.ndim != 3:
            return None

        xyz = torch.nan_to_num(
            scene_xyz.to(device=device, dtype=torch.float32),
            nan=0.0,
            posinf=self.xyz_clip,
            neginf=-self.xyz_clip,
        ).clamp(-self.xyz_clip, self.xyz_clip)
        B, N, _ = xyz.shape
        if N <= 0:
            return None

        priority = torch.zeros(B, N, device=device, dtype=torch.bool)
        gestalt_ids = batch.get("scene_gestalt_ids")
        ade_ids = batch.get("scene_ade_ids")
        if isinstance(gestalt_ids, torch.Tensor) and gestalt_ids.shape[:2] == (B, N):
            gids = gestalt_ids.to(device)
            tier1 = torch.zeros(B, N, device=device, dtype=torch.bool)
            tier2 = torch.zeros(B, N, device=device, dtype=torch.bool)
            for gid in GESTALT_TIER1_IDS:
                tier1 |= gids == gid
            for gid in GESTALT_TIER2_IDS:
                tier2 |= gids == gid

            house = torch.zeros(B, N, device=device, dtype=torch.bool)
            if isinstance(ade_ids, torch.Tensor) and ade_ids.shape[:2] == (B, N):
                aids = ade_ids.to(device)
                for aid in ADE_HOUSE_FOREGROUND_IDS:
                    house |= aids == aid
            priority = tier1 | (tier2 & house)

        inf = torch.full((B, N), 1e10, device=device, dtype=xyz.dtype)
        neg = torch.full_like(inf, -1e10)
        dist = inf.clone()
        idx = torch.zeros(B, k_verts, device=device, dtype=torch.long)
        batch_idx = torch.arange(B, device=device)
        n_priority = priority.sum(dim=1)

        for k in range(k_verts):
            priority_remaining = (n_priority > k).unsqueeze(1)
            eligible = priority | (~priority_remaining)
            score = torch.where(eligible, dist, neg)
            farthest = score.argmax(dim=1)
            idx[:, k] = farthest

            last_xyz = xyz[batch_idx, farthest]
            new_dist = (xyz - last_xyz[:, None, :]).norm(dim=-1)
            dist = torch.minimum(dist, new_dist)
            dist[batch_idx, farthest] = -1e10

        return torch.gather(xyz, 1, idx.unsqueeze(-1).expand(-1, -1, 3))

    def _encoder_head_supervision_scale(
        self,
        batch: dict,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor | None:
        if self.encoder_head_supervision_weight <= 0.0:
            return None
        if self.encoder_head_supervision_frac <= 0.0:
            return None

        progress = batch.get("_train_progress", 0.0)
        if isinstance(progress, torch.Tensor):
            progress_t = progress.to(device=device, dtype=dtype).reshape(())
        else:
            progress_t = torch.tensor(float(progress), device=device, dtype=dtype)
        weight = progress_t.new_tensor(self.encoder_head_supervision_weight)
        zero = progress_t.new_zeros(())
        return torch.where(progress_t < self.encoder_head_supervision_frac, weight, zero)

    def _encoder_head_blend_alpha(
        self,
        batch: dict,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        if self.encoder_head_blend_frac <= 0.0:
            return torch.ones((), device=device, dtype=dtype)

        progress = batch.get("_train_progress")
        if progress is None:
            return torch.ones((), device=device, dtype=dtype)
        if isinstance(progress, torch.Tensor):
            progress_t = progress.to(device=device, dtype=dtype).reshape(())
        else:
            progress_t = torch.tensor(float(progress), device=device, dtype=dtype)
        return (progress_t / self.encoder_head_blend_frac).clamp(0.0, 1.0)

    def _safe_encoder_query_xyz(self, query_xyz: torch.Tensor) -> torch.Tensor:
        query_xyz = torch.nan_to_num(
            query_xyz.float(),
            nan=0.0,
            posinf=self.xyz_clip,
            neginf=-self.xyz_clip,
        )
        return self.xyz_clip * torch.tanh(query_xyz / self.xyz_clip)

    def _fallback_x0(self, batch: dict, K: int, device: torch.device, B: int) -> torch.Tensor:
        x0 = torch.randn(B, K, 3, device=device) * self.noise_sigma_xyz
        if self.init_from_scene:
            xyz_init = self._scene_guided_xyz_init(batch, K, device)
            if xyz_init is not None:
                x0 = xyz_init
        return torch.nan_to_num(
            x0,
            nan=0.0,
            posinf=self.xyz_clip,
            neginf=-self.xyz_clip,
        ).clamp(-self.xyz_clip, self.xyz_clip)

    def _encode_scene(self, batch: dict):
        # Pass through the encoder's native return. When the encoder's
        # `predict_query_xyz` flag is off this is a 2-tuple
        # `(scene_feats, scene_xyz)`; when it's on, a 3-tuple
        # `(scene_feats, scene_xyz, query_xyz)`.
        return self.scene_encoder(
            batch["scene_xyz"],
            batch["scene_type_ids"],
            batch["scene_gestalt_ids"],
            batch["scene_ade_ids"],
            gestalt_id2=batch.get("scene_gestalt_id2"),
            gestalt_w1=batch.get("scene_gestalt_w1"),
            scene_geom_conf=batch.get("scene_geom_conf"),
            scene_sem_conf=batch.get("scene_sem_conf"),
            scene_rgb=batch.get("scene_rgb"),
        )

    def _encode_scene_with_query(self, batch: dict):
        # Internal helper for paths (training, sampling) that need to thread
        # the optional `query_xyz` through `_init_x0`. Always yields a 3-tuple
        # with `query_xyz=None` when the encoder does not predict it.
        out = self._encode_scene(batch)
        if len(out) == 3:
            return out
        scene_feats, scene_xyz = out
        return scene_feats, scene_xyz, None

    def forward(self, batch: dict) -> torch.Tensor:
        """Default forward = compute_loss. Required so DDP's gradient-sync
        hooks fire on `model(batch)`. Inference paths (`sample`,
        `predict_wireframe`) are called directly by name; under DDP, callers
        should `.module.sample(...)` to bypass the wrapper.
        """
        return self.compute_loss(batch)

    def _init_x0(
        self,
        batch: dict,
        K: int,
        device: torch.device,
        B: int,
        query_xyz: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if self.init_from_encoder_head and query_xyz is not None:
            # Learned deterministic query xyz. During early training we can
            # blend from the old stable initializer to the head; inference has
            # no progress marker and therefore uses alpha=1.
            head_x0 = self._safe_encoder_query_xyz(query_xyz)
            if self.encoder_head_blend_frac <= 0.0 or "_train_progress" not in batch:
                return head_x0
            alpha = self._encoder_head_blend_alpha(batch, device, head_x0.dtype)
            base_x0 = self._fallback_x0(batch, K, device, B)
            return ((1.0 - alpha) * base_x0 + alpha * head_x0).clamp(
                -self.xyz_clip, self.xyz_clip
            )
        return self._fallback_x0(batch, K, device, B)

    # ------------------------------------------------------------------
    # Training
    # ------------------------------------------------------------------

    def compute_loss(self, batch: dict) -> torch.Tensor:
        scene_feats, scene_xyz, query_xyz = self._encode_scene_with_query(batch)

        verts_gt = torch.nan_to_num(batch["verts_gt"], nan=0.0, posinf=0.0, neginf=0.0)
        gt_xyz_full = verts_gt[..., :3].clamp(-self.xyz_clip, self.xyz_clip)  # (B, K, 3)
        real_mask = verts_gt[..., 3] > 0                                       # (B, K)
        B, K, _ = gt_xyz_full.shape
        device = gt_xyz_full.device

        x0 = self._init_x0(batch, K, device, B, query_xyz=query_xyz)
        x0 = torch.nan_to_num(
            x0,
            nan=0.0,
            posinf=self.xyz_clip,
            neginf=-self.xyz_clip,
        ).clamp(-self.xyz_clip, self.xyz_clip)
        t = torch.rand(B, device=device)

        # OT-style Hungarian: for each scene, match the K noise slots to the
        # n_real GT vertices on xyz cost. Unmatched slots keep target=x0 (zero
        # velocity) and validity target 0. linear_sum_assignment runs on CPU
        # and releases the GIL, so we dispatch all B problems to a thread pool.
        target_xyz = x0.clone()
        slot_real = torch.zeros(B, K, dtype=torch.bool, device=device)
        gt_to_slot = torch.full((B, K), -1, dtype=torch.long, device=device)

        x0_np = x0.detach().float().cpu().numpy()
        gt_np = gt_xyz_full.detach().float().cpu().numpy()
        n_real_per_scene = real_mask.sum(dim=-1).tolist()

        cost_mats: list[np.ndarray] = []
        active: list[int] = []          # batch indices that have GT to match
        active_n_real: list[int] = []
        for b in range(B):
            n_real = int(n_real_per_scene[b])
            if n_real == 0:
                continue
            diff = x0_np[b][:, None, :] - gt_np[b, :n_real][None, :, :]
            cost_mats.append((diff * diff).sum(-1))    # (K, n_real)
            active.append(b)
            active_n_real.append(n_real)

        for b, n_real, (row, col) in zip(active, active_n_real, _batched_hungarian(cost_mats)):
            row_t = torch.from_numpy(row).to(device).long()
            col_t = torch.from_numpy(col).to(device).long()
            target_xyz[b, row_t] = gt_xyz_full[b, :n_real].index_select(0, col_t)
            slot_real[b, row_t] = True
            gt_to_slot[b, col_t] = row_t

        # Linear-flow targets
        xt = (1.0 - t[:, None, None]) * x0 + t[:, None, None] * target_xyz
        v_target = target_xyz - x0
        v_pred, valid_logit, edge_logit = self.denoiser(xt, t, scene_feats, scene_xyz)
        v_pred = torch.nan_to_num(v_pred, nan=0.0, posinf=self.pred_clip, neginf=-self.pred_clip)

        slot_weights = torch.where(
            slot_real,
            torch.full_like(slot_real, self.real_slot_weight, dtype=torch.float32),
            torch.full_like(slot_real, self.null_slot_weight, dtype=torch.float32),
        )

        # 1) Flow loss on velocity, in the configured norm.
        flow_elem = self._vec_loss(v_pred.float(), v_target.float(), self.loss_flow_norm)
        flow_loss = self._weighted_mean(flow_elem, slot_weights)

        # 2) Endpoint loss: x1 = xt + (1 - t) * v, in the configured norm.
        x1_pred = xt.float() + (1.0 - t[:, None, None]) * v_pred.float()
        endpoint_elem = self._vec_loss(x1_pred, target_xyz.float(), self.loss_endpoint_norm)
        endpoint_loss = self._weighted_mean(endpoint_elem, slot_weights)

        # 3) Focal BCE on validity logit (decoupled head)
        valid_target = slot_real.float()
        valid_logit = valid_logit.float().clamp(
            -self.validity_logit_clip, self.validity_logit_clip,
        )
        bce = F.binary_cross_entropy_with_logits(valid_logit, valid_target, reduction="none")
        prob = torch.sigmoid(valid_logit)
        pt = prob * valid_target + (1.0 - prob) * (1.0 - valid_target)
        focal = (1.0 - pt).clamp_min(1e-6).pow(self.focal_gamma)
        if 0.0 <= self.focal_alpha <= 1.0:
            alpha_t = valid_target * self.focal_alpha + (1.0 - valid_target) * (1.0 - self.focal_alpha)
            focal = focal * alpha_t
        validity_loss = (focal * bce).mean()

        # 4) Simple pairwise BCE on matched real slots. Edges are sparse, so a
        # bounded positive weight keeps the all-negative solution from winning.
        edge_target = self._edge_targets_from_matching(batch, gt_to_slot, slot_real)
        tri = torch.triu(torch.ones(K, K, device=device, dtype=torch.bool), diagonal=1)
        pair_mask = tri[None] & slot_real[:, :, None] & slot_real[:, None, :]
        if pair_mask.any():
            edge_logits = edge_logit.float()[pair_mask].clamp(
                -self.validity_logit_clip, self.validity_logit_clip,
            )
            edge_targets = edge_target[pair_mask]
            pos = edge_targets.sum()
            neg = edge_targets.numel() - pos
            pos_weight = (neg / pos.clamp_min(1.0)).clamp(1.0, 20.0)
            edge_loss = F.binary_cross_entropy_with_logits(
                edge_logits,
                edge_targets,
                pos_weight=pos_weight,
            )
        else:
            edge_loss = edge_logit.sum() * 0.0

        total = (
            self.loss_flow_weight * flow_loss
            + self.loss_endpoint_weight * endpoint_loss
            + self.loss_validity_weight * validity_loss
            + self.loss_edge_weight * edge_loss
        )

        bbox_scale = batch.get("bbox_scale")
        if isinstance(bbox_scale, torch.Tensor):
            bbox_scale = bbox_scale.to(device=device, dtype=torch.float32)
        else:
            bbox_scale = torch.ones(B, device=device)

        iou_loss = edge_logit.sum() * 0.0
        if self.loss_iou_weight > 0.0:
            iou_loss = self._edge_iou_loss(
                x1_pred=x1_pred,
                target_xyz=target_xyz.float(),
                edge_target=edge_target,
                slot_real=slot_real,
                t=t,
                bbox_scale=bbox_scale,
            )
            total = total + self.loss_iou_weight * iou_loss

        soft_vertex_f1_loss = edge_logit.sum() * 0.0
        if self.loss_soft_vertex_f1_weight > 0.0:
            soft_vertex_f1_loss = self._soft_vertex_f1_loss(
                x1_pred=x1_pred,
                gt_xyz_full=gt_xyz_full.float(),
                real_mask=real_mask,
                valid_prob=prob,
                bbox_scale=bbox_scale,
            )
            total = total + self.loss_soft_vertex_f1_weight * soft_vertex_f1_loss

        soft_edge_f1_loss = edge_logit.sum() * 0.0
        if self.loss_soft_edge_f1_weight > 0.0:
            soft_edge_f1_loss = self._soft_edge_f1_loss(
                x1_pred=x1_pred,
                target_xyz=target_xyz.float(),
                edge_logit=edge_logit,
                edge_target=edge_target,
                valid_prob=prob,
                bbox_scale=bbox_scale,
            )
            total = total + self.loss_soft_edge_f1_weight * soft_edge_f1_loss

        count_loss = edge_logit.sum() * 0.0
        if self.loss_count_weight > 0.0:
            count_loss = self._valid_count_loss(prob, real_mask)
            total = total + self.loss_count_weight * count_loss

        head_sup_scale = self._encoder_head_supervision_scale(batch, device, total.dtype)
        if head_sup_scale is not None and query_xyz is not None:
            fps_target = self._priority_fps_xyz_target(batch, K, device)
            if fps_target is not None and fps_target.shape == query_xyz.shape:
                query_safe = torch.nan_to_num(
                    query_xyz.float(),
                    nan=0.0,
                    posinf=self.xyz_clip,
                    neginf=-self.xyz_clip,
                ).clamp(-self.xyz_clip, self.xyz_clip)
                head_sup_elem = self._vec_loss(
                    query_safe,
                    fps_target.to(query_safe),
                    self.loss_endpoint_norm,
                )
                total = total + head_sup_scale * head_sup_elem.mean()

        self.last_loss_terms = {
            "total": total.detach(),
            "flow": flow_loss.detach(),
            "endpoint": endpoint_loss.detach(),
            "validity": validity_loss.detach(),
            "edge": edge_loss.detach(),
            "iou": iou_loss.detach(),
            "soft_vertex_f1": soft_vertex_f1_loss.detach(),
            "soft_edge_f1": soft_edge_f1_loss.detach(),
            "count": count_loss.detach(),
        }
        return total

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    @torch.no_grad()
    def sample(
        self,
        batch: dict,
        n_steps: int = 50,
        validity_thresh: float = 0.0,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        scene_feats, scene_xyz, query_xyz = self._encode_scene_with_query(batch)
        B = scene_feats.shape[0]
        K = self.denoiser.k_verts
        device = scene_feats.device
        dt = 1.0 / n_steps

        x = self._init_x0(batch, 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 = self.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

        valid_mask = last_logit > validity_thresh
        return x, valid_mask, last_edge_logit

    def predict_wireframe(
        self,
        batch: dict,
        n_steps: int = 50,
        validity_thresh: float = 0.0,
    ) -> Tuple[np.ndarray, list]:
        xyz_norm, valid, edge_logit = self.sample(batch, n_steps=n_steps, validity_thresh=validity_thresh)
        xyz_norm = torch.nan_to_num(
            xyz_norm[0].float(),
            nan=0.0,
            posinf=self.xyz_clip,
            neginf=-self.xyz_clip,
        ).clamp(-self.xyz_clip, self.xyz_clip)
        valid = valid[0]
        edge_logit = 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(device=xyz_valid.device, dtype=torch.float32)
        scale = batch["bbox_scale"][0].to(device=xyz_valid.device, dtype=torch.float32)

        bbox_R = batch.get("bbox_R")
        if isinstance(bbox_R, torch.Tensor):
            R = bbox_R[0].to(device=xyz_valid.device, dtype=torch.float32)
            verts_world = (xyz_valid * scale) @ R + center
        else:
            verts_world = xyz_valid * scale + center
        verts_world = verts_world.float().cpu().numpy()

        edges: list[tuple[int, int]] = []
        n_valid = int(valid_idx.numel())
        if n_valid >= 2:
            sub_logits = edge_logit.index_select(0, valid_idx).index_select(1, valid_idx)
            tri = torch.triu(torch.ones(n_valid, n_valid, device=sub_logits.device, dtype=torch.bool), diagonal=1)
            pairs = torch.nonzero((sub_logits > 0.0) & tri, as_tuple=False)
            edges = [(int(i), int(j)) for i, j in pairs.detach().cpu().tolist()]

        return verts_world, edges