# examples/flamingo.py
"""Offline reduced Flamingo bridges on synthetic features and token IDs.

Primary source: https://arxiv.org/html/2204.14198v2, Sections 2.1--2.4,
Figure 4 and Appendix A.1.1--A.1.3 (Figure 5 pseudocode).
This is authored paper reference code, not an original DeepMind release.
Random frozen linear/causal blocks stand in for NFNet and Chinchilla. No
checkpoints, network, semantic answers, native timings, or benchmark claims.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

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

PAD, BOS, IMAGE, EOS = 0, 1, 2, 3


class Attention(nn.Module):
    """Small multi-head attention with exactly zero all-masked outputs."""

    def __init__(self, width, heads=2):
        super().__init__()
        if width < 1 or heads < 1 or width % heads:
            raise ValueError("width must be positive and divisible by heads")
        self.heads = heads
        self.head_width = width // heads
        self.query = nn.Linear(width, width, bias=False)
        self.key = nn.Linear(width, width, bias=False)
        self.value = nn.Linear(width, width, bias=False)
        self.output = nn.Linear(width, width, bias=False)

    def forward(self, query, context, allowed=None):
        """Attend with a boolean [B, Q, K] visibility matrix; True permits."""
        batch, queries, width = query.shape
        keys = context.shape[1]
        if keys == 0:
            return torch.zeros_like(query)

        def split(tensor):
            return tensor.reshape(
                batch, -1, self.heads, self.head_width
            ).transpose(1, 2)

        q = split(self.query(query))
        k = split(self.key(context))
        v = split(self.value(context))
        scores = q @ k.transpose(-1, -2) / math.sqrt(self.head_width)
        if allowed is None:
            weights = scores.softmax(dim=-1)
        else:
            mask = allowed[:, None]
            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
            weights = scores.softmax(dim=-1) * mask
            weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(
                torch.finfo(weights.dtype).tiny
            )
        result = (weights @ v).transpose(1, 2).reshape(batch, queries, width)
        return self.output(result)


def feed_forward(width):
    """Return a reduced pre-normalized dense residual branch."""
    return nn.Sequential(
        nn.LayerNorm(width),
        nn.Linear(width, width * 2),
        nn.GELU(),
        nn.Linear(width * 2, width),
    )


class PerceiverResampler(nn.Module):
    """Map each image/video's frame/spatial features to a fixed latent count."""

    # region: latents
    def __init__(self, width=8, num_latents=4, max_frames=4):
        super().__init__()
        if num_latents < 1 or max_frames < 1:
            raise ValueError("latent and frame counts must be positive")
        self.latents = nn.Parameter(torch.randn(num_latents, width) * 0.02)
        self.temporal = nn.Parameter(torch.randn(max_frames, width) * 0.02)
        self.feature_norm = nn.LayerNorm(width)
        self.latent_norm = nn.LayerNorm(width)
        self.attention = Attention(width)
        self.dense = feed_forward(width)

    # endregion: latents

    # region: resampler
    def forward(self, features):
        """Resample [B, M, F, S, D] to [B, M, L, D], separately per media."""
        batch, media, frames, spatial, width = features.shape
        if frames < 1 or spatial < 1 or frames > self.temporal.shape[0]:
            raise ValueError("require nonempty features within max_frames")
        if media == 0:
            return features.new_empty(batch, 0, self.latents.shape[0], width)
        x = features + self.temporal[:frames][None, None, :, None, :]
        x = self.feature_norm(
            x.reshape(batch * media, frames * spatial, width)
        )
        z = self.latents[None].expand(batch * media, -1, -1)
        context = torch.cat((x, self.latent_norm(z)), dim=1)
        z = z + self.attention(self.latent_norm(z), context)
        z = z + self.dense(z)
        return z.reshape(batch, media, self.latents.shape[0], width)

    # endregion: resampler


# region: media_mask
def media_mask(tokens, num_media, num_latents):
    """Permit only the last image/video marker, including its marker position.

    IMAGE stands for a supplied visual input inserted at that position; 0 in
    phi denotes text before any image. Each input occupies L keys/values.
    """
    if tokens.ndim != 2 or tokens.dtype != torch.long:
        raise ValueError("tokens must be a rank-two int64 tensor")
    if num_media < 0 or num_latents < 1:
        raise ValueError("invalid media or latent count")
    phi = tokens.eq(IMAGE).long().cumsum(dim=1)
    if torch.any(phi > num_media):
        raise ValueError("more image markers than supplied visual inputs")
    key_media = torch.arange(1, num_media + 1, device=tokens.device)
    key_media = key_media.repeat_interleave(num_latents)
    return (phi[:, :, None] == key_media) & (phi[:, :, None] > 0)


# endregion: media_mask


class GatedBridge(nn.Module):
    """New cross-attention and dense branches with independent zero gates."""

    def __init__(self, width=8):
        super().__init__()
        self.norm = nn.LayerNorm(width)
        self.attention = Attention(width)
        self.dense = feed_forward(width)
        self.alpha_xattn = nn.Parameter(torch.zeros(()))
        self.alpha_dense = nn.Parameter(torch.zeros(()))

    # region: gated_xattn
    def cross_attention(self, text, visual, allowed):
        """Keep pre-image rows visually unchanged, even when the gate opens."""
        batch, media, latents, width = visual.shape
        context = visual.reshape(batch, media * latents, width)
        delta = self.attention(self.norm(text), context, allowed)
        return text + self.alpha_xattn.tanh() * delta

    # endregion: gated_xattn

    # region: gated_dense
    def forward(self, text, visual, allowed):
        """Apply cross-attention, then a separately gated dense residual."""
        text = self.cross_attention(text, visual, allowed)
        return text + self.alpha_dense.tanh() * self.dense(text)

    # endregion: gated_dense


class FrozenTextBlock(nn.Module):
    """Random causal text stand-in; freezing weights still allows input grads."""

    def __init__(self, width):
        super().__init__()
        self.norm = nn.LayerNorm(width)
        self.attention = Attention(width)
        self.dense = feed_forward(width)

    def forward(self, text, tokens):
        length = tokens.shape[1]
        causal = torch.ones(
            length, length, dtype=torch.bool, device=tokens.device
        ).tril()
        allowed = causal[None] & tokens.ne(PAD)[:, None, :]
        normalized = self.norm(text)
        text = text + self.attention(normalized, normalized, allowed)
        return text + self.dense(text)


class ReducedFlamingo(nn.Module):
    """One new bridge before one frozen causal block, on supplied tensors."""

    def __init__(self, width=8, feature_width=6, vocab_size=16, num_latents=4):
        super().__init__()
        if feature_width < 1 or vocab_size <= EOS:
            raise ValueError(
                "require positive feature width and special-token vocabulary"
            )
        self.vision = nn.Linear(feature_width, width)
        self.embedding = nn.Embedding(vocab_size, width)
        self.positions = nn.Embedding(64, width)
        self.text_block = FrozenTextBlock(width)
        self.vocabulary = nn.Linear(width, vocab_size, bias=False)
        self.resampler = PerceiverResampler(width, num_latents)
        self.bridge = GatedBridge(width)
        for module in self.frozen_modules():
            module.requires_grad_(False)

    def frozen_modules(self):
        """Return all backbone/embedding/readout stand-ins that never update."""
        return (
            self.vision,
            self.embedding,
            self.positions,
            self.text_block,
            self.vocabulary,
        )

    def validate(self, features, tokens):
        """Reject incompatible, nonfinite, or ambiguous padded input shapes."""
        if tokens.ndim != 2 or tokens.dtype != torch.long:
            raise ValueError("tokens must be [B, T] int64")
        if tokens.shape[0] < 1 or not 2 <= tokens.shape[1] <= 64:
            raise ValueError("require B >= 1 and 2 <= T <= 64")
        if torch.any(tokens < 0) or torch.any(
            tokens >= self.embedding.num_embeddings
        ):
            raise ValueError("token ID outside vocabulary")
        if torch.any(tokens[:, 0] != BOS):
            raise ValueError("sequences must start with BOS")
        padding = tokens.eq(PAD)
        if torch.any(padding[:, :-1] & ~padding[:, 1:]):
            raise ValueError("PAD must be a suffix")
        if features.ndim != 5 or features.shape[0] != tokens.shape[0]:
            raise ValueError("features must be [B, M, F, S, feature_width]")
        if (
            features.shape[-1] != self.vision.in_features
            or not features.is_floating_point()
        ):
            raise ValueError("wrong feature width or nonfloating features")
        if (
            features.device != tokens.device
            or features.device != self.vision.weight.device
        ):
            raise ValueError("model and inputs must share a device")
        if (
            features.dtype != self.vision.weight.dtype
            or not torch.isfinite(features).all()
        ):
            raise ValueError("features must be finite and match model dtype")
        valid = tokens[:, 1:].ne(PAD) & tokens[:, 1:].ne(IMAGE)
        if not valid.any(dim=1).all():
            raise ValueError("each sequence needs a next-text-token target")

    def text_input(self, tokens):
        """Add a frozen position embedding; no real tokenizer is involved."""
        positions = torch.arange(tokens.shape[1], device=tokens.device)
        return self.embedding(tokens) + self.positions(positions)[None]

    def text_only(self, tokens):
        """Compute the exact frozen baseline used in zero-gate checks."""
        return self.vocabulary(
            self.text_block(self.text_input(tokens), tokens)
        )

    # region: frozen_forward
    def forward(self, features, tokens):
        """Keep backbone weights frozen, while differentiating through the LM."""
        self.validate(features, tokens)
        allowed = media_mask(
            tokens, features.shape[1], self.resampler.latents.shape[0]
        )
        with torch.no_grad():
            encoded = self.vision(features)
        visual = self.resampler(encoded)
        hidden = self.bridge(self.text_input(tokens), visual, allowed)
        # Do not use no_grad here: loss gradients must reach the bridge gates.
        logits = self.vocabulary(self.text_block(hidden, tokens))
        labels = tokens[:, 1:].clone()
        labels[(labels == PAD) | (labels == IMAGE)] = -100
        nll = F.cross_entropy(
            logits[:, :-1].transpose(1, 2),
            labels,
            ignore_index=-100,
            reduction="none",
        )
        loss = nll.sum(dim=1).mean()
        return {
            "logits": logits,
            "visual": visual,
            "allowed": allowed,
            "loss": loss,
        }

    # endregion: frozen_forward


def synthetic_batch():
    """Return deterministic supplied features and two interleaved text sequences."""
    generator = torch.Generator().manual_seed(19)
    features = torch.randn(2, 2, 2, 3, 6, generator=generator)
    tokens = torch.tensor(
        [
            [BOS, 4, IMAGE, 5, IMAGE, 6, EOS],
            [BOS, IMAGE, 7, 8, IMAGE, EOS, PAD],
        ]
    )
    return features, tokens


def scalar(tensor):
    """Detach diagnostic scalars without autograd conversion warnings."""
    return float(tensor.detach())


def smoke(seed=7):
    """Run actual finite CPU forwards, backward, and a bridge-only SGD update."""
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    model = ReducedFlamingo()
    features, tokens = synthetic_batch()
    frozen_before = [
        p.detach().clone()
        for m in model.frozen_modules()
        for p in m.parameters()
    ]
    zero = model(features, tokens)
    baseline = model.text_only(tokens)
    zero_error = scalar((zero["logits"] - baseline).abs().max())
    zero["loss"].backward()
    gate_gradients = {
        "xattn": scalar(model.bridge.alpha_xattn.grad),
        "dense": scalar(model.bridge.alpha_dense.grad),
    }
    branch_grad = max(
        scalar(p.grad.abs().max())
        for n, p in model.named_parameters()
        if p.requires_grad and "alpha_" not in n and p.grad is not None
    )
    frozen_gradients_absent = all(
        p.grad is None for m in model.frozen_modules() for p in m.parameters()
    )
    optimizer = torch.optim.SGD(
        (p for p in model.parameters() if p.requires_grad), lr=0.01
    )
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    opened = model(features, tokens)
    changed = features.clone()
    changed[:, 1] = changed[:, 1] + 1.0
    other = model(changed, tokens)
    # Inspect the local bridge, before text self-attention mixes earlier history.
    x = model.text_input(tokens)
    local = model.bridge.cross_attention(
        x, opened["visual"], opened["allowed"]
    )
    local_other = model.bridge.cross_attention(
        x, other["visual"], other["allowed"]
    )
    first_image_rows = tokens.eq(IMAGE).long().cumsum(1).eq(1)
    masked_error = scalar(
        (local[first_image_rows] - local_other[first_image_rows]).abs().max()
    )
    no_media = features[:, :0]
    no_image_tokens = torch.tensor([[BOS, 4, EOS], [BOS, 5, EOS]])
    no_image = model(no_media, no_image_tokens)
    opened["loss"].backward()
    opened_resampler_gradient = sum(
        scalar(p.grad.norm())
        for p in model.resampler.parameters()
        if p.grad is not None
    )
    frozen_after = [
        p.detach() for m in model.frozen_modules() for p in m.parameters()
    ]
    result = {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "device": "cpu",
        "boundary": "random reduced bridges; no pretrained semantics",
        "reductions": {
            "text_nll": "sum valid shifted text targets per sequence; mean batch",
            "excluded_targets": [PAD, IMAGE],
            "dataset_mixture": "omitted",
        },
        "shapes": {
            "features": list(features.shape),
            "visual": list(zero["visual"].shape),
            "logits": list(zero["logits"].shape),
            "mask": list(zero["allowed"].shape),
        },
        "valid_targets": int(
            (tokens[:, 1:].ne(PAD) & tokens[:, 1:].ne(IMAGE)).sum()
        ),
        "loss_zero_gates": scalar(zero["loss"]),
        "loss_after_step": scalar(opened["loss"]),
        "zero_gate_baseline_max_error": zero_error,
        "zero_gate_gradients": gate_gradients,
        "zero_gate_branch_gradient_max": branch_grad,
        "opened_visual_influence_max": scalar(
            (opened["logits"] - other["logits"]).abs().max()
        ),
        "masked_media_local_max_error": masked_error,
        "no_image_finite": bool(torch.isfinite(no_image["logits"]).all()),
        "opened_resampler_gradient_norm_sum": opened_resampler_gradient,
        "frozen_gradients_absent": frozen_gradients_absent,
        "frozen_parameters_unchanged": all(
            torch.equal(a, b) for a, b in zip(frozen_before, frozen_after)
        ),
        "gates_after_step": {
            "xattn": scalar(model.bridge.alpha_xattn),
            "dense": scalar(model.bridge.alpha_dense),
        },
    }
    if not torch.isfinite(opened["logits"]).all() or not math.isfinite(
        result["loss_after_step"]
    ):
        raise RuntimeError("nonfinite smoke output")
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    print(json.dumps(smoke(args.seed), indent=2, allow_nan=False))