"""Scene-only DINOv3 volume + AdaLN-Zero query head — single arm.

Pre-refactor parity (minus wrist + xattn + bimanual). Forward flow::

    rgb         → DINOv3 backbone → patch (B, embed, P_dino, P_dino) + cls (B, embed)
                                ↓ 1×1 conv → upsample → 2-layer residual MLP
    F_scene  (B, feat_dim, P, P)        cls_scene (B, embed_dim)
        ↑ sample at start_pix
    eef_feat (B, feat_dim)
        ⊕ cls_scene → input_proj → d_model
    q_in (B, d_model) → expand to (B*T, d_model) + time-cond sin/cos
        → 5× AdaLN-Zero blocks
    penult (B, T, d_model)
        ├─ q_head  → (q_F, q_z, q_t)         per-T spatial query
        ├─ grip_head → (B, T, n_grip_bins)   gripper bin logits
        └─ rot_head  → (B, T, n_rot_clusters) K-means quat bin logits

    score_yx = q_F · F_scene    (B, T, P, P)
    score_z  = q_z · z_sin      (B, T, Z)
    score_t  = q_t · t_sin      (B, T, T)         — diagonal we keep
    volume   = score_yx[:,:,None] + score_z[:,:,:,None,None]
             + score_t.diagonal(...).view(B, T, 1, 1, 1)        (B, T, Z, P, P)
"""
from __future__ import annotations

import math
import os
from pathlib import Path

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


DEFAULT_DINO_REPO = str(Path(__file__).resolve().parents[3] / "dinov3")
DEFAULT_DINO_WEIGHTS = str(
    Path.home() / "yam_para/dinov3/weights/"
    "dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth"
)


def sin_pe(n: int, d: int) -> torch.Tensor:
    """Standard sinusoidal positional encoding. (n, d)."""
    pe = torch.zeros(n, d)
    pos = torch.arange(n, dtype=torch.float32).unsqueeze(1)
    div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32)
                    * -(math.log(10000.0) / d))
    pe[:, 0::2] = torch.sin(pos * div)
    pe[:, 1::2] = torch.cos(pos * div)
    return pe


class ResidualMLP2(nn.Module):
    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.conv1 = nn.Conv2d(dim, hidden, 1)
        self.conv2 = nn.Conv2d(hidden, dim, 1)
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.conv2(self.act(self.conv1(x)))


class AdaLNZeroBlock(nn.Module):
    """LayerNorm → scale/shift from cond → MLP → gate from cond. Residual."""

    def __init__(self, d_model: int, d_cond: int, mlp_ratio: int = 4):
        super().__init__()
        self.norm = nn.LayerNorm(d_model, elementwise_affine=False)
        self.modulation = nn.Linear(d_cond, 3 * d_model)
        nn.init.zeros_(self.modulation.weight)
        nn.init.zeros_(self.modulation.bias)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, mlp_ratio * d_model),
            nn.GELU(),
            nn.Linear(mlp_ratio * d_model, d_model),
        )

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        scale, shift, gate = self.modulation(cond).chunk(3, dim=-1)
        h = self.norm(x) * (1 + scale) + shift
        return x + gate * self.mlp(h)


class DinoVolumeScene(nn.Module):
    """Single-arm scene-only DINOv3 + factored KV volume + grip + K-means rot heads."""

    def __init__(
        self,
        n_window: int = 16,
        n_height_bins: int = 32,
        pred_size: int = 128,
        feat_dim: int = 48,
        d_model: int = 256,
        d_cond: int = 128,
        d_sin_z: int = 48,
        d_sin_t: int = 48,
        n_blocks: int = 5,
        mlp_hidden: int = 128,
        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.n_z = n_height_bins
        self.pred_size = pred_size
        self.feat_dim = feat_dim
        self.d_model = d_model
        self.d_sin_z = d_sin_z
        self.d_sin_t = d_sin_t
        self.n_gripper_bins = n_gripper_bins
        self.n_rot_clusters = n_rot_clusters

        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                  # 384
        self.embed_dim = embed_dim

        # Per-pixel feature: 1×1 conv → bilinear upsample → 2-layer residual MLP.
        self.feat_head = nn.Conv2d(embed_dim, feat_dim, 1)
        self.refine_mlp = ResidualMLP2(feat_dim, mlp_hidden)

        # Sinusoidal positional bases for height + time.
        self.register_buffer("z_sin", sin_pe(n_height_bins, d_sin_z))
        self.register_buffer("t_sin", sin_pe(n_window, d_sin_t))

        # Input projection: (eef_feat ⊕ cls) → d_model.
        self.input_proj = nn.Linear(feat_dim + embed_dim, d_model)

        # Time-conditioning sin/cos → d_cond projector.
        self.t_cond_proj = nn.Linear(d_sin_t, d_cond)

        # 5× AdaLN-Zero MLP blocks.
        self.blocks = nn.ModuleList([
            AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4)
            for _ in range(n_blocks)
        ])
        self.final_norm = nn.LayerNorm(d_model)

        # Heads.
        d_q = feat_dim + d_sin_z + d_sin_t       # split for (q_F, q_z, q_t)
        self.q_head = nn.Linear(d_model, d_q)
        self.grip_head = nn.Linear(d_model, n_gripper_bins)
        self.rot_head = nn.Linear(d_model, n_rot_clusters)

    @property
    def n_patches_side(self) -> int:
        return self.pred_size

    def patch_features(self, rgb: torch.Tensor):
        """``rgb`` → (patch_grid (B, embed, P_dino, P_dino), cls (B, embed))."""
        B = rgb.size(0)
        out = self.backbone.forward_features(rgb)
        cls = out["x_norm_clstoken"]                         # (B, embed_dim)
        toks = out["x_norm_patchtokens"]                     # (B, P*P, embed_dim)
        P = int(round(math.sqrt(toks.shape[1])))
        return toks.transpose(1, 2).reshape(B, -1, P, P), cls

    def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor) -> dict:
        """rgb (B, 3, S, S), start_pix (B, 2) xy in S-coords."""
        B = rgb.size(0)
        T, Z, P = self.n_window, self.n_z, self.pred_size

        patch, cls = self.patch_features(rgb)
        F0 = self.feat_head(patch)
        F1 = F.interpolate(F0, size=(P, P), mode="bilinear", align_corners=False)
        F_scene = self.refine_mlp(F1)                        # (B, feat_dim, P, P)

        # Sample F_scene at the current EEF pixel (start_pix scaled to P).
        scale = P / rgb.size(-1)
        sx = (start_pix[..., 0] * scale).long().clamp(0, P - 1)
        sy = (start_pix[..., 1] * scale).long().clamp(0, P - 1)
        b_idx = torch.arange(B, device=rgb.device)
        eef_feat = F_scene[b_idx, :, sy, sx]                 # (B, feat_dim)

        q_in = self.input_proj(torch.cat([eef_feat, cls], dim=-1))   # (B, d_model)
        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)                # (T, d_cond)
        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)
        penult = h.view(B, T, self.d_model)                  # (B, T, d_model)

        q_spatial = self.q_head(penult)                      # (B, T, feat_dim + d_sin_z + d_sin_t)
        q_F = q_spatial[..., :self.feat_dim]
        q_z = q_spatial[..., self.feat_dim:self.feat_dim + self.d_sin_z]
        q_t = q_spatial[..., self.feat_dim + self.d_sin_z:]

        # Raw dot products (no normalization, no temperature) — matches the
        # pre-refactor model. Without this, all bilinear scores are bounded
        # in [-1, 1] and the per-T softmax can never sharpen as training
        # progresses, which is why the heatmap looked flat over time.
        score_yx = torch.einsum("btc,bchw->bthw", q_F, F_scene)       # (B, T, P, P)
        score_z = torch.einsum("btc,zc->btz", q_z, self.z_sin)        # (B, T, Z)
        score_t = torch.einsum("btc,tc->bt", q_t, self.t_sin)         # (B, T)  — q_t[T] · t_sin[T]

        vol = (score_yx[:, :, None, :, :]                              # (B, T, 1, P, P)
               + score_z[:, :, :, None, None]                          # (B, T, Z, 1, 1)
               + score_t[:, :, None, None, None])                     # (B, T, 1, 1, 1)

        return {
            "volume_logits": vol,
            "grip_logits": self.grip_head(penult),                    # (B, T, n_grip)
            "rot_logits": self.rot_head(penult),                      # (B, T, n_rot)
            "F_refined": F_scene,                                     # for PCA viz
            "patch_features": patch,                                  # raw DINOv3
        }


class DinoVolumeScenePerVoxel(nn.Module):
    """Per-voxel feature volume — no factorized scoring, no eef sampling.

    Compared to :class:`DinoVolumeScene`:

    - The per-pixel head outputs ``feat_dim * n_height_bins`` channels per
      pixel; reshape gives a **per-voxel** feature volume
      ``F (B, C, Z, P, P)`` — every (height, pixel) cell carries its own
      learned 32-d embedding (no sin/cos PE for height).
    - Input projection is **CLS-only** (no per-pixel ``eef_feat`` lookup at
      ``start_pix``); the per-T global descriptor comes entirely from the
      backbone's CLS token plus FiLM time-conditioning.
    - Per-T query is a single ``C``-dim vector; volume logits come from one
      direct dot product over the voxel grid:
      ``vol = einsum("btc, bczhw -> btzhw", q, F)``.

    Other heads (``grip``, ``rot``) are unchanged.
    """

    def __init__(
        self,
        n_window: int = 16,
        n_height_bins: int = 32,
        pred_size: int = 128,
        feat_dim: int = 32,
        d_model: int = 256,
        d_cond: int = 128,
        d_sin_t: int = 48,
        n_blocks: int = 5,
        mlp_hidden: int = 128,
        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.n_z = n_height_bins
        self.pred_size = pred_size
        self.feat_dim = feat_dim
        self.d_model = d_model
        self.d_sin_t = d_sin_t
        self.n_gripper_bins = n_gripper_bins
        self.n_rot_clusters = n_rot_clusters

        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

        # Per-pixel head: outputs C*Z channels (one C-dim embedding per height bin).
        self.feat_head = nn.Conv2d(embed_dim, feat_dim * n_height_bins, 1)
        self.refine_mlp = ResidualMLP2(feat_dim * n_height_bins, mlp_hidden)

        # Sin/cos PE for time only (used by AdaLN-Zero conditioning, NOT
        # by score factorization — there is no factorization here).
        self.register_buffer("t_sin", sin_pe(n_window, d_sin_t))

        # CLS-only input projection.
        self.input_proj = nn.Linear(embed_dim, d_model)
        self.t_cond_proj = nn.Linear(d_sin_t, d_cond)

        self.blocks = nn.ModuleList([
            AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks)
        ])
        self.final_norm = nn.LayerNorm(d_model)

        # Per-T query: just C-dim. Other heads same as DinoVolumeScene.
        self.q_head = nn.Linear(d_model, feat_dim)
        self.grip_head = nn.Linear(d_model, n_gripper_bins)
        self.rot_head = nn.Linear(d_model, n_rot_clusters)

    @property
    def n_patches_side(self) -> int:
        return self.pred_size

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

    def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor = None) -> dict:
        """``start_pix`` is accepted but ignored (this model uses CLS only)."""
        B = rgb.size(0)
        T, Z, P, C = self.n_window, self.n_z, self.pred_size, self.feat_dim

        patch, cls = self.patch_features(rgb)
        F0 = self.feat_head(patch)                                   # (B, C*Z, P_dino, P_dino)
        F1 = F.interpolate(F0, size=(P, P), mode="bilinear", align_corners=False)
        F2 = self.refine_mlp(F1)                                     # (B, C*Z, P, P)
        F_vol = F2.view(B, C, Z, P, P)                                # per-voxel features

        q_in = self.input_proj(cls)                                  # (B, d_model)
        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)                        # (T, d_cond)
        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)
        penult = h.view(B, T, self.d_model)

        q = self.q_head(penult)                                      # (B, T, C)
        vol = torch.einsum("btc,bczhw->btzhw", q, F_vol)             # (B, T, Z, P, P)

        return {
            "volume_logits": vol,
            "grip_logits": self.grip_head(penult),
            "rot_logits": self.rot_head(penult),
            "F_refined": F2,                                          # (B, C*Z, P, P) — flat for PCA
            "patch_features": patch,
        }


def all_losses(
    out: dict,
    target_yx: torch.Tensor,        # (B, T, 2)
    target_z: torch.Tensor,         # (B, T)
    target_grip: torch.Tensor,      # (B, T)
    target_rot: torch.Tensor,       # (B, T)
) -> dict:
    """Returns dict with per-head CE + total weighted sum."""
    vol = out["volume_logits"]                                       # (B, T, Z, P, P)
    B, T, Z, P, _ = vol.shape
    flat_target = (target_z * (P * P)
                   + target_yx[..., 0] * P
                   + target_yx[..., 1]).reshape(B * T)
    loss_vol = F.cross_entropy(vol.reshape(B * T, Z * P * P), flat_target)
    loss_grip = F.cross_entropy(
        out["grip_logits"].reshape(B * T, -1), target_grip.reshape(B * T))
    loss_rot = F.cross_entropy(
        out["rot_logits"].reshape(B * T, -1), target_rot.reshape(B * T))
    return {
        "loss/volume": loss_vol,
        "loss/grip": loss_grip,
        "loss/rot": loss_rot,
        "loss/total": loss_vol + loss_grip + loss_rot,
    }