"""Scene point cloud encoder.

Pipeline:
    (B, N=8192, ·) per-point MLP     → (B, N, d_model)
                    self-attn × n_full       (full resolution — every point sees all others)
                    cross-attn × n_pool_cross  (n_pool learned slots attend to all N features)
                    self-attn × n_pool       (pooled resolution)
    → (B, n_pool, d_model)  for cross-attention in the denoiser.

Per-point features include an RGB→16-D MLP branch (v10+): COLMAP points use the
COLMAP-stored colour, depth points inherit RGB from their nearest COLMAP
neighbour (raw images are not released), and camera tokens get the (0, 0, 0)
sentinel. Provenance is already in `type_ids` so no extra has-RGB flag is used.

The Perceiver-style cross-attention pool replaces FPS / random subsampling.
Every output slot aggregates from the full 8k-point cloud — no points discarded.
Cost: O(n_pool × N) per head per cross-attn layer, similar to one full self-attn layer.

Full-resolution self-attention uses F.scaled_dot_product_attention (flash backend).

Per-slot xyz anchors are picked per-scene by farthest-point sampling restricted
to the priority-Gestalt subset of the input cloud. This replaces the previous
zero-init learned shared anchor — every scene token now lives at a real,
structurally-relevant 3D position in *that* scene, giving the denoiser cross-
attention a meaningful spatial signal.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
from .pos_enc import fourier_3d

POS_ENC_DIM = 48  # divisible by 6; 8 freq × sin+cos × 3 axes

# Tier/ADE IDs used to bias FPS toward structurally-relevant points. Mirrors
# data/preprocess.py, duplicated here so this module does not depend on the
# data pipeline.
_TIER1_GESTALT_IDS_T = (1, 2, 3, 6, 4, 5, 12)
_TIER2_GESTALT_IDS_T = (8, 9, 10, 11, 18, 27)
_ADE_HOUSE_FOREGROUND_IDS_T = (1, 2, 9, 15, 26, 43, 49, 54)


@torch.no_grad()
def masked_fps(xyz: torch.Tensor, mask: torch.Tensor, n_pool: int) -> torch.Tensor:
    """Batched farthest-point sampling, preferring `mask=True` points first.

    Greedy FPS with the priority-aware twist: at every step, only points that
    are either (a) marked priority, or (b) in a batch that has already
    exhausted its priority points, are eligible. Distances for non-priority
    points are tracked throughout, so by the time they become eligible they
    already hold the correct min-distance to the previously selected set —
    no second pass.

    Args:
        xyz: (B, N, 3) point coordinates.
        mask: (B, N) bool — True for priority candidates.
        n_pool: number of indices to select per batch.
    Returns:
        (B, n_pool) int64 indices into `xyz`.
    """
    B, N, _ = xyz.shape
    device = xyz.device

    INF = 1e10
    NEG = -1e10

    # `dist[b, i]` = min distance from i to the selected-so-far set. Updated
    # for ALL points each iteration (priority and non-priority alike).
    dist = torch.full(xyz.shape[:-1], INF, device=device, dtype=xyz.dtype)
    indices = torch.zeros(B, n_pool, dtype=torch.long, device=device)
    batch_idx = torch.arange(B, device=device)
    n_priority = mask.sum(dim=1)  # (B,)

    for k in range(n_pool):
        priority_remaining = (n_priority > k).unsqueeze(1)   # (B, 1)
        eligible = mask | (~priority_remaining)              # (B, N)

        score = torch.where(eligible, dist, torch.full_like(dist, NEG))
        if k == 0:
            # Initial priorities are all at INF; break ties randomly so
            # different scenes start FPS at different anchor points.
            score = score + torch.rand_like(score) * 1e-3
        farthest = score.argmax(dim=1)
        indices[:, k] = farthest

        last_xyz = xyz[batch_idx, farthest]                   # (B, 3)
        new_dist = (xyz - last_xyz[:, None, :]).norm(dim=-1)
        dist = torch.minimum(dist, new_dist)
        dist[batch_idx, farthest] = NEG  # never re-pick this point

    return indices

class _MHA(nn.Module):
    """Multi-head self-attention via F.scaled_dot_product_attention (flash backend)."""

    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def _shape(self, x: torch.Tensor) -> torch.Tensor:
        B, L, _ = x.shape
        return x.view(B, L, self.n_heads, self.d_head).transpose(1, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        q = self._shape(self.q_proj(x))
        k = self._shape(self.k_proj(x))
        v = self._shape(self.v_proj(x))
        out = F.scaled_dot_product_attention(q, k, v)
        B, H, L, D = out.shape
        out = out.transpose(1, 2).contiguous().view(B, L, H * D)
        return self.out_proj(out)


class _SelfAttnBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.attn = _MHA(d_model, n_heads)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.ff(self.norm2(x))
        return x


class _CrossAttnBlock(nn.Module):
    """Single cross-attention layer: query slots attend to scene tokens.

    Queries (the learned slots) are normalised before attention; key/value
    (scene features) are normalised independently.  A feed-forward follows.
    """

    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.norm_q  = nn.LayerNorm(d_model)
        self.norm_kv = nn.LayerNorm(d_model)
        self.norm_ff = nn.LayerNorm(d_model)
        self.attn = _MHA(d_model, n_heads)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
        # Cross-attention: queries read from scene features.
        # Reuse _MHA but feed it q/k/v from two sources via the proj matrices.
        # _MHA.forward only accepts a single tensor, so we call the projections manually.
        nq  = self.norm_q(q)
        nkv = self.norm_kv(kv)
        B, Lq, _ = nq.shape
        B, Lkv, _ = nkv.shape

        def _shape(x, L):
            return x.view(B, L, self.attn.n_heads, self.attn.d_head).transpose(1, 2)

        qh = _shape(self.attn.q_proj(nq),  Lq)
        kh = _shape(self.attn.k_proj(nkv), Lkv)
        vh = _shape(self.attn.v_proj(nkv), Lkv)
        out = F.scaled_dot_product_attention(qh, kh, vh)
        out = out.transpose(1, 2).contiguous().view(B, Lq, -1)
        q = q + self.attn.out_proj(out)
        q = q + self.ff(self.norm_ff(q))
        return q


class SceneEncoder(nn.Module):
    """Per-point MLP → full self-attn → cross-attn pool → pooled self-attn.

    Pooling uses n_pool learned query slots (Perceiver-style) that cross-attend
    to all N scene features — no points are discarded.
    """

    N_TYPES   = 3
    N_GESTALT = 29   # 0..27 = real gestalt classes, 28 = unobserved sentinel
    N_ADE20K  = 151  # 0 = unknown/unprojected, 1..150 = ADE20K classes

    RGB_DIM = 16  # output width of the per-point RGB MLP

    def __init__(
        self,
        d_model: int = 256,
        n_heads: int = 8,
        d_ff: int = 1024,
        n_full_layers: int = 2,
        n_pool: int = 1024,
        n_pool_cross_layers: int = 1,
        n_pool_layers: int = 2,
        use_rgb: bool = True,
        predict_query_xyz: bool = False,
        n_query: int = 64,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_pool  = n_pool
        self.use_rgb = use_rgb
        self.predict_query_xyz = bool(predict_query_xyz)
        self.n_query = int(n_query)

        self.type_emb    = nn.Embedding(self.N_TYPES,   16)
        # Gestalt: 12-dim embedding, looked up twice (top-1 and top-2) and
        # concatenated → 24 dim. The mixing weight `gestalt_w1` is appended as
        # a raw scalar feature instead of being used to blend the embeddings,
        # so the network can decide how to combine the two labels.
        self.gestalt_emb = nn.Embedding(self.N_GESTALT, 12)
        self.ade_emb     = nn.Embedding(self.N_ADE20K,  8)

        # Per-point RGB → RGB_DIM MLP. Only built when `use_rgb=True` (e.g.
        # cache_version=sem_v10). For sem_v7 caches the RGB branch is absent
        # entirely, keeping the point_mlp input width at 102.
        if self.use_rgb:
            self.rgb_mlp = nn.Sequential(
                nn.Linear(3, self.RGB_DIM),
                nn.GELU(),
                nn.Linear(self.RGB_DIM, self.RGB_DIM),
            )

        # Scalar features: geom_conf, sem_conf, gestalt_w1.
        in_dim = 3 + POS_ENC_DIM + 16 + (12 + 12) + 8 + 3
        if self.use_rgb:
            in_dim += self.RGB_DIM
        self.point_mlp = nn.Sequential(
            nn.Linear(in_dim, d_model),
            nn.LayerNorm(d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
            nn.LayerNorm(d_model),
        )

        self.full_blocks = nn.ModuleList(
            [_SelfAttnBlock(d_model, n_heads, d_ff) for _ in range(n_full_layers)]
        )

        # Learned query content — one d_model vector per output slot. The
        # slot's *spatial anchor* is now picked per scene via FPS over the
        # priority-Gestalt subset of the input (see `forward`), so it is no
        # longer a Parameter.
        self.pool_queries = nn.Parameter(torch.randn(1, n_pool, d_model) * 0.02)
        self.anchor_pos_proj = nn.Linear(POS_ENC_DIM, d_model)

        # Tier/ADE LUTs (registered as buffers so they follow the module's device).
        tier1 = torch.tensor(_TIER1_GESTALT_IDS_T, dtype=torch.long)
        tier2 = torch.tensor(_TIER2_GESTALT_IDS_T, dtype=torch.long)
        house_ade = torch.tensor(_ADE_HOUSE_FOREGROUND_IDS_T, dtype=torch.long)
        self.register_buffer("tier1_ids", tier1, persistent=False)
        self.register_buffer("tier2_ids", tier2, persistent=False)
        self.register_buffer("house_ade_ids", house_ade, persistent=False)

        self.pool_cross_blocks = nn.ModuleList(
            [_CrossAttnBlock(d_model, n_heads, d_ff) for _ in range(n_pool_cross_layers)]
        )
        self.pool_blocks = nn.ModuleList(
            [_SelfAttnBlock(d_model, n_heads, d_ff) for _ in range(n_pool_layers)]
        )
        self.out_norm = nn.LayerNorm(d_model)

        # Optional vertex-query head: K learned query slots cross-attend to the
        # pooled scene tokens, then project to xyz. Provides a learned
        # replacement for FPS/weighted-sampled query-point initialisation in
        # the denoiser. Training can optionally warm-start these outputs with
        # priority-FPS anchor supervision before relying on denoiser losses only.
        if self.predict_query_xyz:
            self.query_embed = nn.Parameter(
                torch.randn(1, self.n_query, d_model) * 0.02
            )
            self.query_cross = _CrossAttnBlock(d_model, n_heads, d_ff)
            self.query_self  = _SelfAttnBlock(d_model, n_heads, d_ff)
            self.query_norm  = nn.LayerNorm(d_model)
            self.query_xyz_head = nn.Linear(d_model, 3)
            nn.init.zeros_(self.query_xyz_head.weight)
            nn.init.zeros_(self.query_xyz_head.bias)

    def forward(
        self,
        xyz: torch.Tensor,          # (B, N, 3)
        type_ids: torch.Tensor,     # (B, N) int
        gestalt_ids: torch.Tensor,  # (B, N) int
        ade_ids: torch.Tensor,      # (B, N) int, 0=unknown
        gestalt_id2: torch.Tensor = None,  # (B, N) int
        gestalt_w1: torch.Tensor = None,   # (B, N) float
        scene_geom_conf: torch.Tensor = None,  # (B, N) float in [0, 1]
        scene_sem_conf: torch.Tensor = None,   # (B, N) float in [0, 1]
        scene_rgb: torch.Tensor = None,        # (B, N, 3) uint8 [0, 255] or float
    ) -> Tuple[torch.Tensor, ...]:
        # Returns:
        #   if predict_query_xyz: (scene_feats (B, n_pool, d_model),
        #                          scene_xyz   (B, n_pool, 3),
        #                          query_xyz   (B, n_query, 3))
        #   else:                 (scene_feats (B, n_pool, d_model),
        #                          scene_xyz   (B, n_pool, 3))
        B, N, _ = xyz.shape

        # Per-point features.
        pos   = fourier_3d(xyz, POS_ENC_DIM)
        t_e   = self.type_emb(type_ids)

        # Gestalt: separate embeddings for top-1 and top-2, concatenated raw
        # (no blend). When gestalt_id2 is missing we look up the same sentinel
        # row used for "unobserved", so the MLP sees a fixed marker.
        if gestalt_id2 is None:
            gestalt_id2 = torch.full_like(gestalt_ids, -1)
        g1_safe = gestalt_ids.masked_fill(gestalt_ids < 0, self.N_GESTALT - 1)
        g2_safe = gestalt_id2.masked_fill(gestalt_id2 < 0, self.N_GESTALT - 1)
        g_e1 = self.gestalt_emb(g1_safe)
        g_e2 = self.gestalt_emb(g2_safe)
        gestalt_feat = torch.cat([g_e1, g_e2], dim=-1)   # (B, N, 24)

        a_e = self.ade_emb(ade_ids)

        if gestalt_w1 is None:
            gestalt_w1 = torch.ones(B, N, device=xyz.device, dtype=xyz.dtype)
        if scene_geom_conf is None:
            scene_geom_conf = torch.ones(B, N, device=xyz.device, dtype=xyz.dtype)
        if scene_sem_conf is None:
            scene_sem_conf = torch.ones(B, N, device=xyz.device, dtype=xyz.dtype)
        scalars = torch.stack([
            scene_geom_conf.to(xyz.dtype),
            scene_sem_conf.to(xyz.dtype),
            gestalt_w1.to(xyz.dtype).clamp(0.0, 1.0),
        ], dim=-1)  # (B, N, 3)

        feats = [xyz, pos, t_e, gestalt_feat, a_e, scalars]
        if self.use_rgb:
            if scene_rgb is None:
                rgb_in = torch.zeros(B, N, 3, device=xyz.device, dtype=xyz.dtype)
            else:
                # Map [0, 255] uint8 → [-1, 1] float in the encoder's dtype.
                rgb_in = scene_rgb.to(xyz.dtype) / 127.5 - 1.0
            feats.append(self.rgb_mlp(rgb_in))  # (B, N, RGB_DIM)

        x = self.point_mlp(torch.cat(feats, dim=-1))  # (B, N, d_model)

        # Full-resolution self-attention: every point sees all others.
        for block in self.full_blocks:
            x = block(x)

        # Spatial anchors for the pool slots: Tier 1 Gestalt points are always
        # priority; Tier 2 Gestalt points count only when ADE says
        # house/foreground, matching preprocessing.
        tier1_mask = (gestalt_ids[..., None] == self.tier1_ids).any(dim=-1)
        tier2_mask = (gestalt_ids[..., None] == self.tier2_ids).any(dim=-1)
        house_mask = (ade_ids[..., None] == self.house_ade_ids).any(dim=-1)
        priority_mask = tier1_mask | (tier2_mask & house_mask)                    # (B, N)
        anchor_idx = masked_fps(xyz, priority_mask, self.n_pool)                   # (B, n_pool)
        xyz_pooled = torch.gather(
            xyz, 1, anchor_idx.unsqueeze(-1).expand(-1, -1, 3)
        )                                                                          # (B, n_pool, 3)

        # Cross-attention pooling: learned slots are conditioned on their
        # per-scene spatial anchors before reading from the full cloud. Without
        # this, pooled features and `xyz_pooled` are only index-paired after the
        # fact, which makes the denoiser attend to features whose content is not
        # tied to the advertised anchor position.
        q = self.pool_queries.expand(B, -1, -1)
        q = q + self.anchor_pos_proj(fourier_3d(xyz_pooled, POS_ENC_DIM))
        for block in self.pool_cross_blocks:
            q = block(q, x)

        # Pooled self-attention: slots refine their joint representation.
        for block in self.pool_blocks:
            q = block(q)

        scene_feats = self.out_norm(q)

        if not self.predict_query_xyz:
            return scene_feats, xyz_pooled

        # K learned vertex queries cross-attend to the pooled scene tokens,
        # then a small MLP projects each slot to xyz. Direct prediction —
        # no offsets relative to scene anchors.
        qv = self.query_embed.expand(B, -1, -1)
        qv = self.query_cross(qv, scene_feats)
        qv = self.query_self(qv)
        qv = self.query_norm(qv)
        query_xyz = self.query_xyz_head(qv)

        return scene_feats, xyz_pooled, query_xyz