"""MapGS: the full feed-forward model (§2.2-§2.5).

Pipeline: ViT encoder fuses the context views into image tokens; the
TokenManager builds map-anchored + free (+ dynamic) query tokens; the DETR
decoder cross-attends them to the image tokens and self-attends among them
(dynamic->static causal mask); the GaussianHead decodes each token into ``N_G``
Gaussians. Dynamic tokens are canonicalized and mapped to the target
timestamp's world frame (:mod:`mapgs.model.dynamic`). Appearance is decoded from
rendered feature maps by a light UNet (or direct RGB).

``forward`` returns the scene Gaussians (a :class:`DecodedGaussians`);
rasterization + appearance decoding are driven by the trainer/losses so that the
extra extrapolation renders (§2.6) reuse the same Gaussians.
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn

from mapgs.config import MapGSConfig
from mapgs.model.encoder import ViTEncoder
from mapgs.model.decoder import DETRDecoder
from mapgs.model.tokens import TokenManager, TokenMeta
from mapgs.model.gaussian_head import GaussianHead, DecodedGaussians
from mapgs.model.dynamic import DynamicModule
from mapgs.model.unet import FeatureUNet


class MapGS(nn.Module):
    def __init__(self, cfg: MapGSConfig):
        super().__init__()
        self.cfg = cfg
        m = cfg.model
        self.encoder = ViTEncoder(
            dim=m.embed_dim, patch=m.patch_size, depth=m.enc_depth, n_heads=m.n_heads,
            mlp_ratio=m.mlp_ratio, qk_norm=m.qk_norm, layerscale_init=m.layerscale_init,
        )
        self.token_manager = TokenManager(m)
        self.decoder = DETRDecoder(
            dim=m.embed_dim, depth=m.dec_depth, n_heads=m.n_heads, mlp_ratio=m.mlp_ratio,
            qk_norm=m.qk_norm, layerscale_init=m.layerscale_init, shared_kv=m.shared_decoder_kv,
        )
        self.head = GaussianHead(m)
        self.dynamic = DynamicModule(cfg)
        if m.feature_color and m.use_unet:
            self.unet = FeatureUNet(in_ch=m.feature_dim, out_ch=3)
        else:
            self.unet = None

    # ------------------------------------------------------------------ #
    def forward(
        self,
        images: torch.Tensor,        # [B, V, 3, H, W]
        plucker: torch.Tensor,       # [B, V, 6, H, W]
        timestep_ids: torch.Tensor,  # [B, V] long
        anchor_pos: torch.Tensor,    # [B, n_map, 3]
        anchor_type: torch.Tensor,   # [B, n_map] long
        anchor_normal: torch.Tensor, # [B, n_map, 3]
        s_t: float = 1.0,
        dynamic: Optional[dict] = None,
    ) -> DecodedGaussians:
        """Decode scene Gaussians in canonical frame. Dynamic Gaussians are
        placed at a target timestamp afterwards via :meth:`place_dynamics`."""
        image_tokens = self.encode(images, plucker, timestep_ids)
        tokens, meta, self_mask = self.build_query_tokens(
            anchor_pos, anchor_type, anchor_normal, image_tokens.shape[0],
            image_tokens.device, dynamic)
        return self.decode(tokens, image_tokens, self_mask, meta, s_t)

    # --- decomposed pieces (also used by test-time token tuning, §2.8) ---
    def encode(self, images, plucker, timestep_ids) -> torch.Tensor:
        return self.encoder(images, plucker, timestep_ids)

    def build_query_tokens(self, anchor_pos, anchor_type, anchor_normal, B, device,
                           dynamic: Optional[dict] = None):
        tokens, meta = self.token_manager.build_static(anchor_pos, anchor_type, anchor_normal)
        self_mask = None
        if dynamic is not None:
            dyn_tokens, dyn_meta = self.dynamic.build_tokens(dynamic, B, device)
            tokens = torch.cat([tokens, dyn_tokens], dim=1)
            meta = TokenMeta.cat_tokens(meta, dyn_meta)
            self_mask = self.dynamic.self_mask(meta)
        return tokens, meta, self_mask

    def decode(self, tokens, image_tokens, self_mask, meta, s_t: float) -> DecodedGaussians:
        tokens = self.decoder(tokens, image_tokens, self_mask=self_mask)
        return self.head(tokens, meta, s_t)

    def place_dynamics(self, decoded: DecodedGaussians, dynamic: Optional[dict],
                       frame_idx: int) -> DecodedGaussians:
        """Rigidly place dynamic Gaussians at ``frame_idx`` (identity if static-only)."""
        if dynamic is None:
            return decoded
        return self.dynamic.place_at(decoded, dynamic, frame_idx)

    # ------------------------------------------------------------------ #
    def set_grad_checkpoint(self, enabled: bool = True):
        self.encoder.grad_checkpoint = enabled
        self.decoder.grad_checkpoint = enabled

    def feature_to_rgb(self, feature_maps: torch.Tensor) -> torch.Tensor:
        """Decode rendered feature maps ``[V, N_c, H, W]`` to RGB ``[V, 3, H, W]``."""
        if self.unet is not None:
            return self.unet(feature_maps)
        return feature_maps[:, :3].clamp(0, 1)

    @property
    def uses_features(self) -> bool:
        return self.cfg.model.feature_color