"""Latent-to-pixel decoder — a diagnostic, never part of the control loop.

Reimplements the decoder described in LeWM appendix D ("Decoder (Visualization
Only)", arXiv:2603.19312). **It is not shipped with the model**: the released
``quentinll/lewm-pusht`` checkpoint holds 303 tensors under
``encoder/predictor/projector/pred_proj/action_encoder`` and nothing decoder
shaped, and neither ``stable-worldmodel`` nor ``lucas-maes/le-wm`` contains the
code. So it has to be trained here — see ``scripts/train_decoder.py``.

The paper's description, restated as shapes::

    z          (B, 192)                 one latent vector, the whole input
      -> to_hidden                      (B, 1, W)      memory: keys/values
    queries    (B, P, W)                P learned tokens, one per output patch
      -> depth x [cross-attn(queries <- memory); residual MLP]
      -> to_pixels                      (B, P, p*p*3)
      -> unpatchify                     (B, 3, 224, 224)

``P = (image_size / patch_size)^2``. The paper's worked example is
``(224/16)^2 = 196``. That 16 is the **decoder's own output tiling** and is
unrelated to the encoder's patch size (which is 14 for this checkpoint): the
decoder never sees patch tokens, only the single pooled vector, so it is free
to carve its output image however it likes.

### One thing worth knowing before you read the attention maps

With ``memory_tokens=1`` — exactly what the paper describes — the cross
attention has a single key. ``softmax`` over one element is identically 1, so
every query gets the same value vector and the layer reduces to

    x <- x + W_o V(z)          (the same increment added to all P queries)

i.e. FiLM-style global conditioning, not a spatial lookup. That is fine for the
job (all spatial structure lives in the learned queries), but attention weights
carry no information and plotting them would be meaningless. Set
``memory_tokens > 1`` to expand ``z`` into several distinct memory slots if you
want attention that actually selects.
"""

import torch
from torch import nn


class CrossAttentionBlock(nn.Module):
    """Pre-norm cross-attention to a fixed memory, then a residual MLP."""

    def __init__(self, dim, heads=6, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.norm_q = nn.LayerNorm(dim)
        self.norm_kv = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(
            dim, heads, dropout=dropout, batch_first=True
        )
        self.norm_mlp = nn.LayerNorm(dim)
        hidden = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x, memory):
        q = self.norm_q(x)
        kv = self.norm_kv(memory)
        x = x + self.attn(q, kv, kv, need_weights=False)[0]
        return x + self.mlp(self.norm_mlp(x))


class LatentDecoder(nn.Module):
    """Decode one latent vector into a 224x224 RGB image.

    Args:
        latent_dim: Width of the latent being decoded (192 here).
        hidden_dim: Decoder working width.
        image_size / patch_size: Output tiling; ``P = (image_size/patch_size)^2``
            learned query tokens, each emitting one ``patch_size^2 * 3`` patch.
        depth: Number of cross-attention + MLP blocks.
        heads: Attention heads.
        memory_tokens: Slots the latent is expanded into. 1 reproduces the
            paper (and makes attention degenerate — see the module docstring).
        out_range: ``'unit'`` clamps nothing and expects targets in ``[0, 1]``
            (the paper's "linearly projected to pixels"); ``'sigmoid'`` bounds
            the output instead. Linear is the faithful default.
    """

    def __init__(
        self,
        latent_dim=192,
        hidden_dim=384,
        image_size=224,
        patch_size=16,
        depth=4,
        heads=6,
        mlp_ratio=4.0,
        dropout=0.0,
        memory_tokens=1,
        out_range='unit',
    ):
        super().__init__()
        assert image_size % patch_size == 0, (
            f'image_size {image_size} not divisible by patch_size {patch_size}'
        )
        self.image_size = image_size
        self.patch_size = patch_size
        self.grid = image_size // patch_size
        self.num_patches = self.grid ** 2
        self.memory_tokens = memory_tokens
        self.out_range = out_range

        # The latent is the decoder's entire input: it becomes the key/value
        # memory that every query reads from.
        self.to_hidden = nn.Linear(latent_dim, hidden_dim * memory_tokens)

        # One learned query per output patch. These carry all the spatial
        # structure — the latent itself has no spatial layout at all.
        self.queries = nn.Parameter(
            torch.randn(1, self.num_patches, hidden_dim) * 0.02
        )

        self.blocks = nn.ModuleList(
            CrossAttentionBlock(hidden_dim, heads, mlp_ratio, dropout)
            for _ in range(depth)
        )
        self.norm = nn.LayerNorm(hidden_dim)
        self.to_pixels = nn.Linear(hidden_dim, patch_size * patch_size * 3)

    @torch.no_grad()
    def init_output_at(self, mean_pixel):
        """Start the decoder at "always predict the mean colour".

        Default init emits roughly zero-mean patches against targets in
        ``[0, 1]``, so step 0 sits ~300x worse than simply predicting the mean
        image and the first few hundred steps are spent rediscovering the
        background. Zeroing the output weight and parking the bias on the mean
        colour starts training at that baseline instead, so every subsequent
        step is spent on structure. Gradients are unaffected — the weight is
        zero, not frozen.

        Args:
            mean_pixel: Per-channel mean in ``[0, 1]``, shape ``(3,)``.
        """
        mean_pixel = torch.as_tensor(mean_pixel, dtype=torch.float32).view(3)
        patch = mean_pixel.repeat(self.patch_size * self.patch_size)
        self.to_pixels.weight.zero_()
        self.to_pixels.bias.copy_(patch)

    def unpatchify(self, patches):
        """``(B, P, p*p*3)`` -> ``(B, 3, H, W)``, row-major over the grid."""
        B = patches.size(0)
        g, p = self.grid, self.patch_size
        x = patches.view(B, g, g, p, p, 3)
        x = x.permute(0, 5, 1, 3, 2, 4)          # (B, 3, gy, py, gx, px)
        return x.reshape(B, 3, g * p, g * p)

    def patchify(self, images):
        """``(B, 3, H, W)`` -> ``(B, P, p*p*3)``; inverse of :meth:`unpatchify`."""
        B = images.size(0)
        g, p = self.grid, self.patch_size
        x = images.view(B, 3, g, p, g, p)
        x = x.permute(0, 2, 4, 3, 5, 1)          # (B, gy, gx, py, px, c)
        return x.reshape(B, self.num_patches, p * p * 3)

    def forward(self, z):
        """Decode ``(B, D)`` or ``(B, T, D)`` latents to images.

        A ``(B, T, D)`` input returns ``(B, T, 3, H, W)``, which is what you
        want for a plan rollout: one frame per horizon step.
        """
        squeeze_time = z.dim() == 2
        if squeeze_time:
            z = z.unsqueeze(1)
        B, T, _ = z.shape
        z = z.reshape(B * T, -1)

        memory = self.to_hidden(z).view(B * T, self.memory_tokens, -1)
        x = self.queries.expand(B * T, -1, -1)
        for block in self.blocks:
            x = block(x, memory)

        images = self.unpatchify(self.to_pixels(self.norm(x)))
        if self.out_range == 'sigmoid':
            images = torch.sigmoid(images)

        images = images.view(B, T, 3, self.image_size, self.image_size)
        return images.squeeze(1) if squeeze_time else images


def reconstruction_loss(pred, target):
    """Plain pixel MSE.

    The paper never writes the reconstruction objective down — it only says a
    decoder was "trained to reconstruct pixel observations" and, in the
    appendix G ablation, that adding "a reconstruction loss" to LeWM training
    *hurt* control (PushT 96.0 -> 86.0 SR). MSE is the default reading and the
    one used here; nothing downstream depends on the choice, since this decoder
    never touches the world model's gradients.
    """
    return torch.nn.functional.mse_loss(pred, target)


def load_decoder(path, device='cpu'):
    """Rebuild a decoder from a training checkpoint."""
    ckpt = torch.load(path, map_location=device, weights_only=False)
    decoder = LatentDecoder(**ckpt['config'])
    decoder.load_state_dict(ckpt['state_dict'])
    return decoder.to(device).eval(), ckpt