"""Small heads on ARROW's encoder tokens (DA3-Giant: per frame, a 3072-d token per 14x14 patch + an id token).

They decide *what to query* at decode time, so they must be far cheaper than the decoder:
  motion       per-pixel logit "this pixel moves around this frame's time" (two linears per patch, pixel-shuffled
               to 14x14 pixels)
  sky          per-pixel logit, same shape (sky pixels are never queried for geometry)
  covisibility per-pixel logit "this pixel of frame i is visible in frame j", conditioned on frame j's id token
               (FiLM on the patch feature); decides which pixels of a new frame reused geometry can cover
  descriptor   per-patch L2-normalised local descriptors + a GeM-pooled global descriptor (matching / retrieval)
All heads read the same tokens (B, F, N, C) with the patch grid (h, w); none touch the decoder.
"""

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

TOKEN_DIM = 3072
PATCH = 14


def pixel_shuffle(logits, h, w, patch=PATCH):
    """(B, F, h*w, patch*patch) per-patch pixel logits -> (B, F, h*patch, w*patch)."""
    B, Fr = logits.shape[:2]
    x = logits.view(B, Fr, h, w, patch, patch).permute(0, 1, 2, 4, 3, 5)
    return x.reshape(B, Fr, h * patch, w * patch)


class PatchMLP(nn.Module):
    def __init__(self, out_dim, hidden=256, in_dim=TOKEN_DIM):
        super().__init__()
        self.net = nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, hidden), nn.GELU(), nn.Linear(hidden, out_dim))

    def forward(self, x):
        return self.net(x)


class DenseLogitHead(nn.Module):
    """Per-pixel binary logits from patch tokens (motion, sky)."""

    def __init__(self, hidden=256, patch=PATCH):
        super().__init__()
        self.patch = patch
        self.mlp = PatchMLP(patch * patch, hidden)

    def forward(self, tokens, h, w):
        return pixel_shuffle(self.mlp(tokens), h, w, self.patch)


class CovisibilityHead(nn.Module):
    """Per-pixel logit that a pixel of frame i is seen in frame j: FiLM(patch_i | id_j) -> 14x14 logits."""

    def __init__(self, hidden=256, patch=PATCH):
        super().__init__()
        self.patch = patch
        self.source = PatchMLP(hidden, hidden)
        self.film = nn.Sequential(nn.LayerNorm(TOKEN_DIM), nn.Linear(TOKEN_DIM, 2 * hidden))
        self.out = nn.Sequential(nn.GELU(), nn.Linear(hidden, patch * patch))

    def forward(self, tokens, id_tokens, pairs, h, w):
        """tokens (B, F, N, C), id_tokens (B, F, C), pairs (P, 2) frame indices (i, j) -> (B, P, H, W)."""
        f = self.source(tokens[:, pairs[:, 0]])                     # (B, P, N, hidden)
        scale, shift = self.film(id_tokens[:, pairs[:, 1]]).chunk(2, dim=-1)
        f = f * (1 + scale[:, :, None]) + shift[:, :, None]
        return pixel_shuffle(self.out(f), h, w, self.patch)


class DescriptorHead(nn.Module):
    def __init__(self, dim=128, global_dim=256, hidden=512, gem_p=3.0):
        super().__init__()
        self.local = PatchMLP(dim, hidden)
        self.global_proj = nn.Linear(dim, global_dim)
        self.p = nn.Parameter(torch.tensor(float(gem_p)))

    def forward(self, tokens):
        local = F.normalize(self.local(tokens), dim=-1)                # (B, F, N, dim)
        pooled = local.clamp_min(1e-6).pow(self.p).mean(2).pow(1 / self.p)
        return local, F.normalize(self.global_proj(pooled), dim=-1)    # (B, F, global_dim)


class TokenHeads(nn.Module):
    """The configured subset of heads; forward returns a dict keyed by head name."""

    def __init__(self, heads=("motion", "covisibility", "descriptor", "sky"), hidden=256, descriptor_dim=128):
        super().__init__()
        self.names = tuple(heads)
        unknown = set(self.names) - {"motion", "covisibility", "descriptor", "sky"}
        if unknown:
            raise ValueError(f"Unknown heads: {sorted(unknown)}")
        self.heads = nn.ModuleDict()
        if "motion" in self.names:
            self.heads["motion"] = DenseLogitHead(hidden)
        if "sky" in self.names:
            self.heads["sky"] = DenseLogitHead(hidden)
        if "covisibility" in self.names:
            self.heads["covisibility"] = CovisibilityHead(hidden)
        if "descriptor" in self.names:
            self.heads["descriptor"] = DescriptorHead(descriptor_dim)

    def forward(self, tokens, id_tokens, h, w, pairs=None):
        out = {}
        for name in ("motion", "sky"):
            if name in self.heads:
                out[name] = self.heads[name](tokens, h, w)
        if "covisibility" in self.heads and pairs is not None and len(pairs):
            out["covisibility"] = self.heads["covisibility"](tokens, id_tokens, pairs, h, w)
        if "descriptor" in self.heads:
            out["descriptor_local"], out["descriptor_global"] = self.heads["descriptor"](tokens)
        return out