# examples/coca.py
"""Authored reduced CoCa paper reference, not original Google model code.

Primary source: https://arxiv.org/html/2205.01917v2, sections 3.1/3.2,
equations 2/3/4, section 4.1 and section 4.3 Table 8f (cascade default).
Offline CPU mechanics with random weights and synthetic token IDs. No pretrained
semantics, captions, ViT reproduction, distributed performance or benchmarks.
Contrastive reduction: sum of two directional batch means (paper Eq. 2).
Caption reduction: mean over examples of each example's valid target sum (Eq. 3).
The paper specifies contrastive/caption weights 1 and 2; original training
reduction code was not located.
"""

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, EOS = 0, 1, 2


def causal_mask(length, device):
    """Returns True where a key lies strictly after its query."""
    return torch.ones(length, length, dtype=torch.bool, device=device).triu(1)


class CausalBlock(nn.Module):
    """Applies causal self-attention, optional visual fusion, and a small FFN."""

    def __init__(self, width, heads, visual=False):
        super().__init__()
        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)
        self.norm = nn.LayerNorm(width)
        self.cross = (
            nn.MultiheadAttention(width, heads, batch_first=True)
            if visual
            else None
        )
        self.cross_norm = nn.LayerNorm(width) if visual else None
        self.ffn = nn.Sequential(
            nn.Linear(width, 2 * width), nn.GELU(), nn.Linear(2 * width, width)
        )
        self.ffn_norm = nn.LayerNorm(width)

    def forward(self, hidden, padding, visual=None):
        """Returns hidden states without any future-token dependency."""
        normalized = self.norm(hidden)
        attended, _ = self.attention(
            normalized,
            normalized,
            normalized,
            attn_mask=causal_mask(hidden.shape[1], hidden.device),
            key_padding_mask=padding,
            need_weights=False,
        )
        hidden = hidden + attended
        if self.cross is not None:
            attended, _ = self.cross(
                self.cross_norm(hidden), visual, visual, need_weights=False
            )
            hidden = hidden + attended
        return hidden + self.ffn(self.ffn_norm(hidden))


# region: visual_tokens
class ImageEncoder(nn.Module):
    """Maps 4x4 patches to width-D tokens; deliberately smaller than a ViT."""

    def __init__(self, width):
        super().__init__()
        self.patches = nn.Conv2d(3, width, kernel_size=4, stride=4)
        self.calls = 0

    def forward(self, images):
        """Evaluates the shared image encoder once per joint model forward."""
        self.calls += 1
        return self.patches(images).flatten(2).transpose(1, 2)


# endregion: visual_tokens


# region: poolers
class Pooler(nn.Module):
    """Pools keys/values with learned queries, one output per query."""

    def __init__(self, queries, width, heads):
        super().__init__()
        self.queries = nn.Parameter(torch.randn(queries, width) * 0.02)
        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)

    def forward(self, tokens):
        """Returns [B, queries, D] rather than claiming lossless compression."""
        queries = self.queries.unsqueeze(0).expand(tokens.shape[0], -1, -1)
        return self.attention(queries, tokens, tokens, need_weights=False)[0]


def cascade_poolers(model, visual):
    """Follows the original default in section 4.3/Table 8f, not Figure 2."""
    caption_tokens = model.caption_pooler(visual)
    contrastive_token = model.contrastive_pooler(caption_tokens)
    return caption_tokens, contrastive_token


# endregion: poolers


class ReducedCoCa(nn.Module):
    """Routes causal text independently for alignment, then fuses visual tokens."""

    def __init__(
        self, width=8, heads=2, caption_queries=4, vocab=16, max_length=16
    ):
        super().__init__()
        if (
            width < 2
            or heads < 1
            or width % heads
            or caption_queries < 1
            or vocab < 4
            or max_length < 2
        ):
            raise ValueError(
                "Invalid width, heads, queries, vocabulary or length"
            )
        self.vocab, self.max_length = vocab, max_length
        self.image_encoder = ImageEncoder(width)
        self.caption_pooler = Pooler(caption_queries, width, heads)
        self.contrastive_pooler = Pooler(1, width, heads)
        self.embedding = nn.Embedding(vocab, width, padding_idx=PAD)
        self.positions = nn.Parameter(torch.randn(max_length + 1, width) * 0.02)
        self.cls = nn.Parameter(torch.randn(1, 1, width) * 0.02)
        self.unimodal = CausalBlock(width, heads)
        self.multimodal = CausalBlock(width, heads, visual=True)
        self.image_projection = nn.Linear(width, 4)
        self.text_projection = nn.Linear(width, 4)
        self.caption_head = nn.Linear(width, vocab)
        self.log_temperature = nn.Parameter(torch.tensor(math.log(0.07)))

    def validate(self, images, tokens):
        """Rejects empty, malformed, nonfinite and non-suffix-padded inputs."""
        if (
            images.ndim != 4
            or images.shape[0] < 1
            or images.shape[1] != 3
            or min(images.shape[2:]) < 4
            or any(size % 4 for size in images.shape[2:])
            or not images.is_floating_point()
            or not torch.isfinite(images).all()
        ):
            raise ValueError(
                "Images must be finite [B, 3, H, W], H/W multiples of 4"
            )
        if (
            tokens.ndim != 2
            or tokens.dtype != torch.long
            or tokens.shape[0] != images.shape[0]
            or not 2 <= tokens.shape[1] <= self.max_length
            or (tokens < 0).any()
            or (tokens >= self.vocab).any()
            or tokens.device != images.device
        ):
            raise ValueError(
                "Tokens must be in-vocabulary long [B, L], 2 <= L <= max"
            )
        for row in tokens:
            valid = row[row != PAD]
            if (
                valid.numel() < 2
                or valid[0] != BOS
                or valid[-1] != EOS
                or (valid[1:-1] < 3).any()
                or not torch.equal(row[: valid.numel()], valid)
            ):
                raise ValueError(
                    "Each row needs BOS, content, EOS, then suffix PAD"
                )

    # region: unimodal_decoder
    def encode_text(self, tokens):
        """Appends CLS so a causal query can summarize all preceding valid text."""
        hidden = self.embedding(tokens)
        cls = self.cls.expand(tokens.shape[0], -1, -1)
        hidden = torch.cat((hidden, cls), dim=1)
        hidden = hidden + self.positions[: hidden.shape[1]]
        padding = torch.cat(
            (
                tokens.eq(PAD),
                torch.zeros(
                    tokens.shape[0], 1, dtype=torch.bool, device=tokens.device
                ),
            ),
            dim=1,
        )
        # No image tensor enters this function; pads are blocked as keys.
        return self.unimodal(hidden, padding)

    # endregion: unimodal_decoder

    # region: contrastive_embeddings
    def contrastive(self, image_token, unimodal):
        """Uses trailing CLS, unit embeddings and a trainable positive temperature."""
        image = F.normalize(self.image_projection(image_token[:, 0]), dim=-1)
        text = F.normalize(self.text_projection(unimodal[:, -1]), dim=-1)
        logits = image @ text.T / self.log_temperature.exp()
        labels = torch.arange(image.shape[0], device=image.device)
        # Eq. 2 adds directional batch means: deliberately no factor of 1/2.
        loss = F.cross_entropy(logits, labels) + F.cross_entropy(
            logits.T, labels
        )
        return image, text, logits, loss

    # endregion: contrastive_embeddings

    # region: multimodal_decoder
    def caption_logits(self, unimodal, tokens, caption_tokens):
        """Fuses images after the unimodal split; excludes CLS from caption states."""
        text_states = unimodal[:, :-1]
        conditioned = self.multimodal(
            text_states, tokens.eq(PAD), caption_tokens
        )
        return self.caption_head(conditioned)

    # endregion: multimodal_decoder

    def forward(self, images, tokens):
        """Computes both objectives' states from one shared image evaluation."""
        self.validate(images, tokens)
        visual = self.image_encoder(images)
        caption_tokens, image_token = cascade_poolers(self, visual)
        unimodal = self.encode_text(tokens)
        image, text, similarities, contrastive_loss = self.contrastive(
            image_token, unimodal
        )
        logits = self.caption_logits(unimodal, tokens, caption_tokens)
        return dict(
            visual=visual,
            caption_tokens=caption_tokens,
            contrastive_token=image_token,
            unimodal=unimodal,
            image_embedding=image,
            text_embedding=text,
            similarities=similarities,
            contrastive_loss=contrastive_loss,
            caption_logits=logits,
        )


# region: caption_loss
def caption_loss(logits, tokens):
    """Sums next-token NLL per example, then averages examples, ignoring PAD."""
    # Position i predicts input token i+1; BOS and appended CLS are never targets.
    targets = tokens[:, 1:]
    predictions = logits[:, :-1]
    losses = F.cross_entropy(
        predictions.transpose(1, 2), targets, ignore_index=PAD, reduction="none"
    )
    return losses.sum(dim=1).mean()


# endregion: caption_loss


# region: joint_loss
def joint_loss(output, tokens, contrastive_weight=1.0, caption_weight=2.0):
    """Weights explicitly declared reductions; these weights follow section 4.1."""
    weights = (contrastive_weight, caption_weight)
    if any(not math.isfinite(w) or w < 0 for w in weights) or not any(weights):
        raise ValueError(
            "Loss weights must be finite, nonnegative, and not all zero"
        )
    contrastive = output["contrastive_loss"]
    caption = caption_loss(output["caption_logits"], tokens)
    total = contrastive_weight * contrastive + caption_weight * caption
    return total, contrastive, caption


def train_step(model, images, tokens, optimizer):
    """Runs one joint forward, one backward, then one optimizer update."""
    optimizer.zero_grad(set_to_none=True)
    output = model(images, tokens)
    total, contrastive, caption = joint_loss(output, tokens)
    total.backward()
    optimizer.step()
    return output, (total.detach(), contrastive.detach(), caption.detach())


# endregion: joint_loss


def synthetic_batch(seed=7):
    """Returns a local seeded batch containing unequal caption lengths."""
    generator = torch.Generator().manual_seed(seed)
    images = torch.randn(3, 3, 8, 8, generator=generator)
    tokens = torch.tensor(
        [[BOS, 3, 4, EOS, PAD], [BOS, 5, 6, 7, EOS], [BOS, 8, EOS, PAD, PAD]]
    )
    return images, tokens


def smoke(seed=7):
    """Measures topology and gradients at fixed parameters, then a real update."""
    torch.manual_seed(seed)
    model = ReducedCoCa()
    images, tokens = synthetic_batch(seed)
    output = model(images, tokens)
    total, contrastive, caption = joint_loss(output, tokens)
    parameters = tuple(model.parameters())
    con_grad = torch.autograd.grad(
        contrastive, parameters, retain_graph=True, allow_unused=True
    )
    cap_grad = torch.autograd.grad(
        caption, parameters, retain_graph=True, allow_unused=True
    )
    total_grad = torch.autograd.grad(
        total, parameters, retain_graph=True, allow_unused=True
    )
    errors = []
    for parameter, con, cap, joint in zip(
        parameters, con_grad, cap_grad, total_grad
    ):
        zero = torch.zeros_like(parameter)
        expected = (con if con is not None else zero) + 2 * (
            cap if cap is not None else zero
        )
        errors.append(
            float(
                ((joint if joint is not None else zero) - expected).abs().max()
            )
        )
    forward_count = model.image_encoder.calls
    with torch.no_grad():
        other = model(images + 1.0, tokens)
        future = tokens.clone()
        future[:2, 2] = 10
        changed = model(images, future)
    before = model.image_encoder.patches.weight.detach().clone()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
    optimizer.zero_grad(set_to_none=True)
    total.backward()  # Same already-computed forward graph; no image reevaluation.
    norm = float(model.image_encoder.patches.weight.grad.norm())
    optimizer.step()
    return dict(
        boundary="authored reduced paper reference; random CPU mechanics only",
        source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        seed=seed,
        reductions=dict(
            contrastive="sum of directional batch means",
            caption="batch mean of valid target sums",
        ),
        weights=dict(contrastive=1.0, caption=2.0),
        shapes={
            key: list(output[key].shape)
            for key in (
                "visual",
                "caption_tokens",
                "contrastive_token",
                "unimodal",
                "image_embedding",
                "text_embedding",
                "caption_logits",
                "similarities",
            )
        },
        losses=dict(
            total=float(total.detach()),
            contrastive=float(contrastive.detach()),
            caption=float(caption.detach()),
        ),
        valid_caption_targets=int(tokens[:, 1:].ne(PAD).sum()),
        joint_image_encoder_forwards=forward_count,
        unimodal_image_delta=float(
            (output["text_embedding"] - other["text_embedding"]).abs().max().detach()
        ),
        conditioned_caption_delta=float(
            (output["caption_logits"] - other["caption_logits"]).abs().max().detach()
        ),
        causal_unimodal_prefix_delta=float(
            (output["unimodal"][:, :2] - changed["unimodal"][:, :2]).abs().max().detach()
        ),
        causal_caption_prefix_delta=float(
            (output["caption_logits"][:, :2] - changed["caption_logits"][:, :2])
            .abs()
            .max().detach()
        ),
        weighted_gradient_max_error=max(errors),
        image_gradient_norm=norm,
        image_weight_change=float(
            (model.image_encoder.patches.weight.detach() - before).norm()
        ),
    )


def main():
    """Prints finite measured JSON only when explicitly invoked."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not args.smoke:
        parser.error("Pass --smoke for seeded offline CPU verification")
    torch.set_num_threads(1)
    print(json.dumps(smoke(args.seed), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()