"""Motion Tracks baseline — CLS → (T, 3) but as (u, v, z) instead of (x, y, z).

Architectural twin of :class:`DinoCLSXYZBaseline`: same DINOv3 CLS,
same AdaLN-Zero stack, same auxiliary grip / rot heads. The position
head still produces ``(B, T, 3)``, but the three channels are now:

  - **u, v**  pixel coordinates of the EE in the scene camera at
              ``img_size`` resolution.
  - **z**     world-frame EE height (metres), same as the
              ``--baseline_cls_xyz`` target.

Training: project the ground-truth world XYZ through ``K_in`` and
``T_w2c`` (scene cam, scaled to ``img_size``) and pair it with world-z;
MSE on the resulting (u, v, z) tuple.

Inference: unproject (u, v) at world height z via the same factorized
solver our v4 / v3 models use:
  - cast a ray from cam_pos through pixel (u, v),
  - intersect with the world plane z = predicted_z,
  - return the world-XYZ. (Helper in :func:`uvz_to_world_xyz`.)

The point of this baseline is to mirror "Motion Tracks" (RT-Tracks /
P3PO-style methods) — regress 2D image tracks then lift to 3D — while
keeping every other design knob identical to the regression baseline,
so any delta is purely from "XYZ space vs UVZ space."
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F

from .model import AdaLNZeroBlock, DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS, sin_pe


class DinoCLSUVZBaseline(nn.Module):
    """CLS-only baseline that regresses (u, v, z) instead of (x, y, z).

    Output shape is unchanged ``(B, T, 3)`` so downstream pipelines (the
    loss bundle, the deploy lift) just need to know which interpretation
    to apply. ``model_kind = 'baseline_cls_uvz'`` in the ckpt is the
    discriminator.
    """

    def __init__(
        self,
        n_window: int = 32,
        d_model: int = 256,
        d_cond: int = 128,
        d_sin_t: int = 48,
        n_blocks: int = 5,
        n_gripper_bins: int = 32,
        n_rot_clusters: int = 64,
        dino_repo: str = DEFAULT_DINO_REPO,
        dino_weights: str = DEFAULT_DINO_WEIGHTS,
    ):
        super().__init__()
        self.n_window = n_window
        self.d_model = d_model
        self.n_gripper_bins = n_gripper_bins
        self.n_rot_clusters = n_rot_clusters

        import os
        self.backbone = torch.hub.load(
            os.environ.get("DINO_REPO_DIR", dino_repo),
            "dinov3_vits16plus",
            source="local",
            weights=os.environ.get("DINO_WEIGHTS_PATH", dino_weights),
        )
        embed_dim = self.backbone.embed_dim
        self.embed_dim = embed_dim

        self.register_buffer("t_sin", sin_pe(n_window, d_sin_t))
        self.t_cond_proj = nn.Linear(d_sin_t, d_cond)
        self.input_proj = nn.Linear(embed_dim, d_model)
        self.blocks = nn.ModuleList([
            AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks)
        ])
        self.final_norm = nn.LayerNorm(d_model)

        self.uvz_head = nn.Linear(d_model, 3)        # (u, v, z) per timestep
        self.grip_head = nn.Linear(d_model, n_gripper_bins)
        self.rot_head = nn.Linear(d_model, n_rot_clusters)

    def _encode(self, rgb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        import math
        out = self.backbone.forward_features(rgb)
        cls = out["x_norm_clstoken"]
        toks = out["x_norm_patchtokens"]
        P = int(round(math.sqrt(toks.shape[1])))
        patches = toks.transpose(1, 2).reshape(toks.shape[0], -1, P, P)
        return cls, patches

    def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor = None) -> dict:
        B = rgb.size(0); T = self.n_window
        cls, patches = self._encode(rgb)
        q_in = self.input_proj(cls)
        q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1)
        cond_t = self.t_cond_proj(self.t_sin)
        cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1)
        h = q_in_bt
        for blk in self.blocks:
            h = blk(h, cond_bt)
        h = self.final_norm(h).view(B, T, self.d_model)
        return {
            "uvz_pred": self.uvz_head(h),                 # (B, T, 3): (u, v, z)
            "grip_logits": self.grip_head(h),
            "rot_logits": self.rot_head(h),
            "cls": cls,
            "patch_features": patches,                     # (B, D, P, P) for DINO PCA
        }


def world_xyz_to_uvz(target_xyz: torch.Tensor, K_in: torch.Tensor,
                      T_w2c: torch.Tensor) -> torch.Tensor:
    """Project (B, T, 3) world XYZ → (B, T, 3) (u, v, z) using the
    SCENE cam pose. ``K_in`` is at ``img_size`` scale.

    Returns z = world Z unchanged. ``u``/``v`` are pixel coordinates
    (not normalised). Out-of-frustum or behind-camera points get
    pixel coords that lie outside the image; downstream loss masks
    them with the same chunk-level mask the deterministic baseline
    uses (i.e. no special handling — the network just learns the
    extrapolation if any).
    """
    B, T, _ = target_xyz.shape
    ones = torch.ones(B, T, 1, device=target_xyz.device, dtype=target_xyz.dtype)
    xyz_h = torch.cat([target_xyz, ones], dim=-1)                     # (B, T, 4)
    cam = torch.einsum("bij,btj->bti", T_w2c, xyz_h)[..., :3]         # (B, T, 3)
    z_cam = cam[..., 2:3].clamp(min=1e-3)
    uv = torch.einsum("bij,btj->bti", K_in, cam) / z_cam              # (B, T, 3)
    uv = uv[..., :2]                                                   # (B, T, 2)
    z_world = target_xyz[..., 2:3]                                     # (B, T, 1)
    return torch.cat([uv, z_world], dim=-1)


def uvz_to_world_xyz(uvz: torch.Tensor, K_in: torch.Tensor,
                      T_w2c: torch.Tensor) -> torch.Tensor:
    """Inverse of :func:`world_xyz_to_uvz`. Lifts predicted (u, v, z)
    to world XYZ by ray-casting from the camera centre through pixel
    (u, v) and intersecting the world plane Z = z.

    Shapes:
      uvz   : (..., 3) — (u, v, world_z)
      K_in  : (..., 3, 3) broadcastable with the leading uvz dims
      T_w2c : (..., 4, 4)

    Returns world XYZ with the same leading dims as ``uvz``.
    """
    *lead, _ = uvz.shape
    u = uvz[..., 0]
    v = uvz[..., 1]
    z_w = uvz[..., 2]
    K_inv = torch.inverse(K_in)
    T_c2w = torch.inverse(T_w2c)
    cam_o = T_c2w[..., :3, 3]                                          # (..., 3)
    R = T_c2w[..., :3, :3]                                             # (..., 3, 3)

    pix_h = torch.stack([u, v, torch.ones_like(u)], dim=-1)            # (..., 3)
    # Ray in cam frame; broadcast K_inv to ray pixels.
    ray_cam = torch.einsum("...ij,...j->...i", K_inv.expand_as(R), pix_h)
    ray_world = torch.einsum("...ij,...j->...i", R, ray_cam)
    rz = ray_world[..., 2:3]
    rz = torch.where(rz.abs() < 1e-6, torch.full_like(rz, 1e-6), rz)
    lam = (z_w.unsqueeze(-1) - cam_o[..., 2:3]) / rz
    return cam_o.unsqueeze(-2).expand_as(ray_world) * 0 + cam_o.unsqueeze(-2) \
              if False else (cam_o.unsqueeze(-2).broadcast_to(ray_world.shape)
                              + lam * ray_world)


def baseline_motion_tracks_losses(
        out: dict,
        target_xyz: torch.Tensor,
        target_grip: torch.Tensor,
        target_rot: torch.Tensor,
        K_in: torch.Tensor,
        T_w2c: torch.Tensor,
        img_size: int = 504) -> dict:
    """Loss bundle: MSE on projected (u, v, z); CE on grip / rot.

    ``u`` and ``v`` are normalised by ``img_size`` BEFORE the MSE so
    they sit on the same scale as ``z`` (metres). This keeps the loss
    balanced without a separate weight per channel.
    """
    target_uvz = world_xyz_to_uvz(target_xyz, K_in, T_w2c)              # (B, T, 3)
    pred_uvz = out["uvz_pred"]
    # Normalise pixel coords by img_size so u, v ∈ [0, 1] roughly,
    # comparable in magnitude to world-z (~[0, 0.5] m).
    pred_n = torch.stack([
        pred_uvz[..., 0] / img_size,
        pred_uvz[..., 1] / img_size,
        pred_uvz[..., 2],
    ], dim=-1)
    target_n = torch.stack([
        target_uvz[..., 0] / img_size,
        target_uvz[..., 1] / img_size,
        target_uvz[..., 2],
    ], dim=-1)
    uvz_loss = F.mse_loss(pred_n, target_n)

    B, T = target_grip.shape
    grip_loss = F.cross_entropy(
        out["grip_logits"].reshape(B * T, -1), target_grip.reshape(B * T))
    rot_loss = F.cross_entropy(
        out["rot_logits"].reshape(B * T, -1), target_rot.reshape(B * T))
    return {
        "loss/xyz": uvz_loss,                  # keyed "xyz" for log-compat
        "loss/grip": grip_loss,
        "loss/rot": rot_loss,
        "loss/total": uvz_loss + grip_loss + rot_loss,
    }