# examples/clip_training.py
"""Reduced random-weight CLIP mechanics; no pretrained semantic predictions.

Run `python clip_training.py --smoke` offline with PyTorch only. This uses a
32-pixel ViT, width 32, one block per encoder, a 16-dimensional joint space,
synthetic already-tokenized pairs and AdamW. It is not the WIT training recipe.
The separate clip_export.py executes authenticated official ViT-B/32 weights.
"""

import argparse
import json
import math

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


class AttentionBlock(nn.Module):
    """Pre-norm attention and QuickGELU FFN, without dropout."""

    def __init__(self, width, heads):
        super().__init__()
        if width <= 0 or heads <= 0 or width % heads:
            raise ValueError("Positive width must be divisible by head count.")
        self.norm1 = nn.LayerNorm(width)
        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)
        self.norm2 = nn.LayerNorm(width)
        self.expand = nn.Linear(width, 4 * width)
        self.contract = nn.Linear(4 * width, width)

    def forward(self, x, mask=None):
        """Preserve [batch, length, width] with an optional causal mask."""
        h = self.norm1(x)
        x = x + self.attention(h, h, h, attn_mask=mask, need_weights=False)[0]
        h = self.expand(self.norm2(x))
        return x + self.contract(h * torch.sigmoid(1.702 * h))


# region: tokens
def pad_tokens(contents, context=8, sot=62, eot=63):
    """Bracket integer content IDs; forbid specials inside the content."""
    if context < 2 or not 0 <= sot < eot or not contents:
        raise ValueError("Require nonempty batch, context >= 2 and SOT < EOT.")
    result = torch.zeros(len(contents), context, dtype=torch.long)
    for row, content in enumerate(contents):
        if len(content) + 2 > context:
            raise ValueError("Text exceeds context; no silent truncation.")
        if any(type(x) is not int or not 0 <= x < sot for x in content):
            raise ValueError(
                "Content IDs must be non-special nonnegative ints."
            )
        sequence = [sot, *content, eot]
        result[row, : len(sequence)] = torch.tensor(sequence)
    return result


# endregion: tokens


class ImageEncoder(nn.Module):
    """Small ViT with the extra pre-transformer normalization used by CLIP."""

    def __init__(self, size=32, patch=8, width=32, dim=16, heads=4):
        super().__init__()
        if min(size, patch, dim) <= 0 or size % patch:
            raise ValueError("Image size must be divisible by patch size.")
        self.size = size
        self.patch = nn.Conv2d(3, width, patch, stride=patch, bias=False)
        self.cls = nn.Parameter(torch.randn(width) / math.sqrt(width))
        self.positions = nn.Parameter(
            torch.randn(1 + (size // patch) ** 2, width) / math.sqrt(width)
        )
        self.pre = nn.LayerNorm(width)
        self.block = AttentionBlock(width, heads)
        self.post = nn.LayerNorm(width)
        self.projection = nn.Parameter(
            torch.randn(width, dim) / math.sqrt(width)
        )

    # region: image
    def forward(self, images):
        """Map [B,3,H,W] through CLS to a learned joint-space projection."""
        if images.ndim != 4 or images.shape[1:] != (3, self.size, self.size):
            raise ValueError("Unexpected image shape.")
        patches = self.patch(images).flatten(2).transpose(1, 2)
        cls = self.cls.expand(images.shape[0], 1, -1)
        tokens = torch.cat([cls, patches], dim=1) + self.positions
        hidden = self.block(self.pre(tokens))
        return self.post(hidden[:, 0]) @ self.projection

    # endregion: image


class TextEncoder(nn.Module):
    """Causal text transformer with EOT pooling and a learned projection."""

    def __init__(self, vocab=64, context=8, width=32, dim=16, heads=4):
        super().__init__()
        if vocab < 3 or context < 2 or dim <= 0:
            raise ValueError("Invalid vocabulary, context or output dimension.")
        self.eot = vocab - 1
        self.context = context
        self.embedding = nn.Embedding(vocab, width)
        self.positions = nn.Parameter(torch.randn(context, width) * 0.01)
        self.block = AttentionBlock(width, heads)
        self.final = nn.LayerNorm(width)
        self.projection = nn.Parameter(
            torch.randn(width, dim) / math.sqrt(width)
        )
        self.register_buffer("mask", causal_mask(context))

    # region: text
    def hidden(self, tokens):
        """Future token keys receive minus infinity before attention softmax."""
        if tokens.ndim != 2 or tokens.shape[1] != self.context:
            raise ValueError("Tokens must have the configured context length.")
        if (
            tokens.dtype != torch.long
            or ((tokens < 0) | (tokens > self.eot)).any()
        ):
            raise ValueError("Tokens must be in-vocabulary int64 IDs.")
        x = self.embedding(tokens) + self.positions
        return self.final(self.block(x, self.mask))

    def forward(self, tokens):
        """EOT is the largest ID; require exactly one per sequence."""
        if tokens.ndim != 2 or (tokens == self.eot).sum(1).ne(1).any():
            raise ValueError("Each sequence must contain exactly one EOT.")
        hidden = self.hidden(tokens)
        eot_positions = tokens.argmax(dim=-1)
        pooled = hidden[torch.arange(tokens.shape[0]), eot_positions]
        return pooled @ self.projection

    # endregion: text


def causal_mask(context):
    """Return zero on/below diagonal and minus infinity above it."""
    if context < 1:
        raise ValueError("Context must be positive.")
    return torch.full((context, context), float("-inf")).triu(1)


# region: normalize
def unit_features(features):
    """Normalize each nonzero finite vector; reject undefined cosine inputs."""
    if features.ndim != 2 or min(features.shape) < 1:
        raise ValueError("Features must be a nonempty matrix.")
    norms = features.norm(dim=-1, keepdim=True)
    if (
        not torch.isfinite(features).all()
        or not torch.isfinite(norms).all()
        or (norms <= 0).any()
    ):
        raise ValueError("Every embedding must be finite and nonzero.")
    return features / norms


# endregion: normalize


# region: scale
def similarity_logits(
    image_features, text_features, logit_scale, training=False
):
    """Original training caps scale at 100; released eval uses raw exp(scale)."""
    image = unit_features(image_features)
    text = unit_features(text_features)
    if image.shape[1] != text.shape[1]:
        raise ValueError("Image/text embedding dimensions differ.")
    if logit_scale.numel() != 1 or not torch.isfinite(logit_scale).all():
        raise ValueError("Logit scale must be one finite scalar.")
    scale = logit_scale.exp()
    if training:
        scale = scale.clamp(max=100)
    if not torch.isfinite(scale):
        raise ValueError("Exponentiated scale is not finite.")
    return scale * (image @ text.T)


# endregion: scale


# region: loss
def symmetric_loss(logits):
    """The diagonal contains paired positives in both directions."""
    if (
        logits.ndim != 2
        or logits.shape[0] != logits.shape[1]
        or not logits.numel()
    ):
        raise ValueError("Paired loss requires a nonempty square logit matrix.")
    if not torch.isfinite(logits).all():
        raise ValueError("Logits must be finite.")
    labels = torch.arange(logits.shape[0], device=logits.device)
    image_loss = F.cross_entropy(logits, labels)
    text_loss = F.cross_entropy(logits.T, labels)
    return (image_loss + text_loss) / 2


# endregion: loss


class PairedEncoders(nn.Module):
    """Two jointly trainable encoders and a log-parameterized scale."""

    def __init__(self):
        super().__init__()
        self.image = ImageEncoder()
        self.text = TextEncoder()
        self.logit_scale = nn.Parameter(torch.tensor(math.log(1 / 0.07)))


# region: update
def train_step(model, optimizer, images, tokens):
    """Differentiate both encoders, both projections and the learned scale."""
    model.train()
    optimizer.zero_grad(set_to_none=True)
    image_features = model.image(images)
    text_features = model.text(tokens)
    logits = similarity_logits(
        image_features, text_features, model.logit_scale, training=True
    )
    loss = symmetric_loss(logits)
    loss.backward()
    optimizer.step()
    # Keep the parameter itself inside the original training scale bound.
    with torch.no_grad():
        model.logit_scale.clamp_(max=math.log(100))
    return loss.detach(), logits.detach()


# endregion: update


# region: prototypes
def class_prototypes(prompt_features):
    """Normalize prompts, average within each class, then renormalize."""
    if prompt_features.ndim != 3 or min(prompt_features.shape) < 1:
        raise ValueError("Require [classes, templates, embedding dimension].")
    classes, templates, dim = prompt_features.shape
    normalized = unit_features(prompt_features.reshape(-1, dim))
    averaged = normalized.reshape(classes, templates, dim).mean(1)
    return unit_features(averaged)


# endregion: prototypes


# region: inference
def candidate_scores(image_features, prototypes, logit_scale):
    """Return candidate-conditional, uncalibrated softmax and cosine scores."""
    cosine = unit_features(image_features) @ unit_features(prototypes).T
    logits = similarity_logits(image_features, prototypes, logit_scale)
    return {"cosine": cosine, "logits": logits, "softmax": logits.softmax(-1)}


# endregion: inference


def smoke():
    """Execute every annotated region using deterministic synthetic pairs."""
    torch.set_num_threads(2)
    torch.manual_seed(12)
    model = PairedEncoders()
    images = torch.randn(3, 3, 32, 32)
    tokens = pad_tokens([[1, 4], [2, 5, 6], [3, 7]])
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss, logits = train_step(model, optimizer, images, tokens)
    model.eval()
    with torch.no_grad():
        features = model.text(tokens)
        prototypes = class_prototypes(features[:, None, :])
        scores = candidate_scores(
            model.image(images), prototypes, model.logit_scale
        )
    return {
        "scope": "random-weight reduced paired-encoder training mechanics",
        "seed": 12,
        "loss": float(loss),
        "tokens": tokens.tolist(),
        "trainingLogits": logits.tolist(),
        "imageShape": list(images.shape),
        "prototypeShape": list(prototypes.shape),
        "candidateSoftmax": scores["softmax"].tolist(),
        "finiteGradients": all(
            p.grad is not None and torch.isfinite(p.grad).all().item()
            for p in model.parameters()
        ),
    }


def main():
    """Expose a small, network-free command-line check."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    args = parser.parse_args()
    if not args.smoke:
        parser.error("Choose --smoke.")
    print(json.dumps(smoke(), indent=2))


if __name__ == "__main__":
    main()