"""Scene-graph dynamics for rigid instances (§2.5, PointForward formulation).

Each instance ``i`` has a per-frame box trajectory. The earliest box defines the
canonical frame; dynamic tokens decode Gaussians in that canonical world
placement. To render timestamp ``t'`` we apply the rigid map back:

    p_world(t') = R_{t'} R_1^T (p_canon - c_1) + c_{t'}              (inverse of PF Eq.2)

rotating the Gaussian's quaternion by the same ``dR = R_{t'} R_1^T`` and scaling
opacity by the lifespan ``o' = o * exp(-1/2 ((t'-t)/(sigma+1))^2)`` (PF Eq.7).

Decoding is done once in the canonical frame, so a single forward pass serves
all target timestamps. Dynamic tokens attend **causally to static tokens only**
(no token attends to a dynamic key), realized as a key mask.
"""

from __future__ import annotations

from typing import Tuple

import torch
import torch.nn as nn

from mapgs.config import MapGSConfig
from mapgs.geometry.transforms import rotmat_to_quat, quat_multiply, normalize_quat
from mapgs.model.gaussian_head import DecodedGaussians
from mapgs.model.tokens import TokenMeta
from mapgs.render.gaussians import GROUP_DYNAMIC


def matrix_to_rot6d(R: torch.Tensor) -> torch.Tensor:
    """First two columns of a rotation, ``[..., 3, 3]`` -> ``[..., 6]`` (continuous repr)."""
    return R[..., :, :2].reshape(*R.shape[:-2], 6)


def place_dynamic_gaussians(g, box_centers, box_rots, canon_idx, frame_idx, lifespan_sigma=2.0):
    """Rigidly place a single scene's dynamic Gaussians at ``frame_idx`` (§2.5).

    Operates on a :class:`Gaussians` whose ``group``/``instance_id`` tag dynamic
    primitives. ``box_centers [I,F,3]``, ``box_rots [I,F,3,3]``, ``canon_idx [I]``.
    """
    dyn = g.group == GROUP_DYNAMIC
    if not bool(dyn.any()):
        return g
    device = g.means.device
    box_centers = box_centers.to(device)
    box_rots = box_rots.to(device)
    canon_idx = canon_idx.to(device)
    inst = g.instance_id[dyn].clamp(min=0)
    F = box_centers.shape[1]
    f = min(int(frame_idx), F - 1)
    c1 = box_centers[inst, canon_idx[inst]]
    R1 = box_rots[inst, canon_idx[inst]]
    cf = box_centers[inst, f]
    Rf = box_rots[inst, f]
    dR = Rf @ R1.transpose(-1, -2)
    pw = torch.einsum("nij,nj->ni", dR, g.means[dyn] - c1) + cf
    dq = rotmat_to_quat(dR)
    qw = normalize_quat(quat_multiply(dq, g.quats[dyn]))
    lifespan = torch.exp(-0.5 * ((f - canon_idx[inst].float()) / (lifespan_sigma + 1)) ** 2)

    means = g.means.clone(); means[dyn] = pw
    quats = g.quats.clone(); quats[dyn] = qw
    opac = g.opacities.clone(); opac[dyn] = g.opacities[dyn] * lifespan
    return g.with_overrides(means=means, quats=quats, opacities=opac)


class DynamicModule(nn.Module):
    def __init__(self, cfg: MapGSConfig):
        super().__init__()
        self.n_dyn = cfg.model.tokens.n_dyn_per_instance
        self.dim = cfg.model.embed_dim
        self.fps = cfg.data.fps
        self.lifespan_sigma = 2.0  # frames; opacity temporal spread (Eq.7)
        self.template = nn.Parameter(torch.randn(self.n_dyn, self.dim) * 0.01)
        self.box_mlp = nn.Sequential(
            nn.Linear(3 + 3 + 6, self.dim),  # center + size + rot6d
            nn.GELU(),
            nn.Linear(self.dim, self.dim),
        )
        self.dyn_type_embed = nn.Parameter(torch.zeros(1, 1, self.dim))

    # ------------------------------------------------------------------ #
    def build_tokens(self, dynamic: dict, B: int, device, target_time=None) -> Tuple[torch.Tensor, TokenMeta]:
        centers = dynamic["box_centers"].to(device)  # [B, I, F, 3]
        rots = dynamic["box_rots"].to(device)        # [B, I, F, 3, 3]
        size = dynamic["box_size"].to(device)        # [B, I, 3]
        inst_valid = dynamic["valid"].to(device)     # [B, I] bool
        canon_idx = dynamic["canon_idx"].to(device)  # [B, I] long
        Bc, I, F, _ = centers.shape

        bidx = torch.arange(B, device=device)[:, None].expand(B, I)
        iidx = torch.arange(I, device=device)[None].expand(B, I)
        c1 = centers[bidx, iidx, canon_idx]           # [B, I, 3]
        R1 = rots[bidx, iidx, canon_idx]              # [B, I, 3, 3]
        box_feat = torch.cat([c1, size, matrix_to_rot6d(R1)], dim=-1)  # [B, I, 12]
        cond = self.box_mlp(box_feat)                 # [B, I, C]

        tokens = self.template[None, None] + cond[:, :, None, :] + self.dyn_type_embed
        tokens = tokens.reshape(B, I * self.n_dyn, self.dim)

        group = torch.full((B, I * self.n_dyn), GROUP_DYNAMIC, dtype=torch.long, device=device)
        anchor_pos = torch.zeros(B, I * self.n_dyn, 3, device=device)
        instance_id = iidx[:, :, None].expand(B, I, self.n_dyn).reshape(B, I * self.n_dyn).contiguous()
        valid = inst_valid[:, :, None].expand(B, I, self.n_dyn).reshape(B, I * self.n_dyn).contiguous()
        anchor_type = torch.full((B, I * self.n_dyn), -1, dtype=torch.long, device=device)
        meta = TokenMeta(group=group, anchor_pos=anchor_pos, instance_id=instance_id,
                         valid=valid, anchor_type=anchor_type)
        return tokens, meta

    # ------------------------------------------------------------------ #
    def self_mask(self, meta: TokenMeta) -> torch.Tensor:
        """Key mask ``[B, 1, 1, T]`` (True = attendable): valid & non-dynamic.

        No token attends to a dynamic key (motion depends on static structure,
        not vice versa, and dynamic-dynamic interactions are excluded). Dynamic
        queries still attend to all static keys.
        """
        keep_key = meta.valid & (meta.group != GROUP_DYNAMIC)  # [B, T]
        return keep_key[:, None, None, :]

    # ------------------------------------------------------------------ #
    def place_at(self, decoded: DecodedGaussians, dynamic: dict, frame_idx: int) -> DecodedGaussians:
        """Rigid canonical->world placement of dynamic Gaussians at ``frame_idx``."""
        device = decoded.means.device
        centers = dynamic["box_centers"].to(device)
        rots = dynamic["box_rots"].to(device)
        canon_idx = dynamic["canon_idx"].to(device)
        B, I, F, _ = centers.shape

        bidx = torch.arange(B, device=device)[:, None].expand(B, I)
        iidx = torch.arange(I, device=device)[None].expand(B, I)
        c1 = centers[bidx, iidx, canon_idx]                  # [B, I, 3]
        R1 = rots[bidx, iidx, canon_idx]                     # [B, I, 3, 3]
        fr = torch.full((B, I), int(frame_idx), device=device, dtype=torch.long).clamp(0, F - 1)
        ct = centers[bidx, iidx, fr]                         # [B, I, 3]
        Rt = rots[bidx, iidx, fr]                            # [B, I, 3, 3]
        dR = Rt @ R1.transpose(-1, -2)                       # [B, I, 3, 3]
        dquat = rotmat_to_quat(dR)                           # [B, I, 4]
        t_canon = canon_idx.to(torch.float32)
        lifespan = torch.exp(-0.5 * ((frame_idx - t_canon) / (self.lifespan_sigma + 1)) ** 2)  # [B,I]

        means = decoded.means.clone()
        quats = decoded.quats.clone()
        opac = decoded.opacities.clone()

        is_dyn = decoded.group == GROUP_DYNAMIC                # [B, M]
        for b in range(B):
            m = is_dyn[b]
            if not m.any():
                continue
            inst = decoded.instance_id[b][m].clamp(min=0)      # [Md]
            p = decoded.means[b][m]                            # [Md, 3]
            dRm = dR[b][inst]                                  # [Md, 3, 3]
            c1m = c1[b][inst]
            ctm = ct[b][inst]
            pw = torch.einsum("nij,nj->ni", dRm, p - c1m) + ctm
            means[b][m] = pw
            qw = quat_multiply(dquat[b][inst], decoded.quats[b][m])
            quats[b][m] = normalize_quat(qw)
            opac[b][m] = decoded.opacities[b][m] * lifespan[b][inst]

        return DecodedGaussians(
            means=means, colors=decoded.colors, scales=decoded.scales,
            opacities=opac, quats=quats, group=decoded.group,
            instance_id=decoded.instance_id, valid=decoded.valid, lane=decoded.lane,
        )