# examples/vit.py
"""Readable ViT with state keys compatible with torchvision 0.27.0.

Run --smoke for random-weight mechanics or --predict --image PHOTO for ViT-B/16.
This is an educational CPU reference, not a reproduction of JFT pre-training.
"""

import argparse
from collections import OrderedDict
import json
import math
from pathlib import Path

from PIL import Image
import torch
from torch import nn
from torch.nn import functional as F
from torchvision.models import ViT_B_16_Weights
from torchvision.transforms import functional as TF


# region: attention
class SelfAttention(nn.Module):
    """Expose the Q/K/V and head axes; keep released checkpoint key names."""

    def __init__(self, hidden_dim, num_heads, dropout=0.0):
        super().__init__()
        if hidden_dim % num_heads:
            raise ValueError("Hidden width must be divisible by head count")
        self.num_heads = num_heads
        self.in_proj_weight = nn.Parameter(
            torch.empty(3 * hidden_dim, hidden_dim)
        )
        self.in_proj_bias = nn.Parameter(torch.zeros(3 * hidden_dim))
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)
        nn.init.xavier_uniform_(self.in_proj_weight)
        nn.init.zeros_(self.out_proj.bias)

    def forward(self, x):
        batch, tokens, width = x.shape
        head_dim = width // self.num_heads
        qkv = F.linear(x, self.in_proj_weight, self.in_proj_bias)
        qkv = qkv.reshape(batch, tokens, 3, self.num_heads, head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(head_dim)
        weights = self.dropout(scores.softmax(dim=-1))
        attended = (weights @ v).transpose(1, 2).reshape(batch, tokens, width)
        return self.out_proj(attended)


# endregion: attention


# region: encoder
class EncoderBlock(nn.Module):
    """Pre-normalize each branch; add its result to the unnormalized residual."""

    def __init__(self, width, heads, mlp_dim, dropout, attention_dropout):
        super().__init__()
        self.ln_1 = nn.LayerNorm(width, eps=1e-6)
        self.self_attention = SelfAttention(width, heads, attention_dropout)
        self.dropout = nn.Dropout(dropout)
        self.ln_2 = nn.LayerNorm(width, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(width, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, width),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        x = x + self.dropout(self.self_attention(self.ln_1(x)))
        return x + self.mlp(self.ln_2(x))


# endregion: encoder


class Encoder(nn.Module):
    """Named blocks, learned positions, and final tokenwise normalization."""

    def __init__(
        self, tokens, depth, heads, width, mlp_dim, dropout, attention_dropout
    ):
        super().__init__()
        self.pos_embedding = nn.Parameter(
            torch.empty(1, tokens, width).normal_(std=0.02)
        )
        self.dropout = nn.Dropout(dropout)
        self.layers = nn.Sequential(
            OrderedDict(
                (
                    f"encoder_layer_{i}",
                    EncoderBlock(
                        width, heads, mlp_dim, dropout, attention_dropout
                    ),
                )
                for i in range(depth)
            )
        )
        self.ln = nn.LayerNorm(width, eps=1e-6)

    def forward(self, tokens):
        return self.ln(self.layers(self.dropout(tokens)))


class VisionTransformer(nn.Module):
    """Patch projection, learned CLS/positions, encoder, and supervised head."""

    def __init__(
        self,
        image_size=224,
        patch_size=16,
        num_layers=12,
        num_heads=12,
        hidden_dim=768,
        mlp_dim=3072,
        num_classes=1000,
        dropout=0.0,
        attention_dropout=0.0,
        representation_size=None,
    ):
        super().__init__()
        dimensions = (
            image_size,
            patch_size,
            num_layers,
            num_heads,
            hidden_dim,
            mlp_dim,
            num_classes,
        )
        if any(not isinstance(n, int) or n <= 0 for n in dimensions):
            raise ValueError("Model dimensions must be positive integers")
        if image_size % patch_size or hidden_dim % num_heads:
            raise ValueError(
                "Image and hidden dimensions must divide into patches and heads"
            )
        self.image_size = image_size
        self.patch_size = patch_size
        self.conv_proj = nn.Conv2d(
            3, hidden_dim, kernel_size=patch_size, stride=patch_size
        )
        self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))
        tokens = (image_size // patch_size) ** 2 + 1
        self.encoder = Encoder(
            tokens,
            num_layers,
            num_heads,
            hidden_dim,
            mlp_dim,
            dropout,
            attention_dropout,
        )
        head = OrderedDict()
        if representation_size is not None:
            if (
                not isinstance(representation_size, int)
                or representation_size <= 0
            ):
                raise ValueError(
                    "Representation width must be a positive integer"
                )
            head["pre_logits"] = nn.Linear(hidden_dim, representation_size)
            head["act"] = nn.Tanh()
        head["head"] = nn.Linear(representation_size or hidden_dim, num_classes)
        self.heads = nn.Sequential(head)
        # The paper replaces the downstream classifier with a zero-initialized head.
        nn.init.zeros_(self.heads.head.weight)
        nn.init.zeros_(self.heads.head.bias)

    # region: patches
    def patch_tokens(self, images):
        if images.ndim != 4 or images.shape[1] != 3:
            raise ValueError("Expected RGB images with shape [B, 3, H, W]")
        if tuple(images.shape[-2:]) != (self.image_size, self.image_size):
            raise ValueError(
                "Use the configured image size; resize positions when fine-tuning"
            )
        # Conv2d uses channel/row/column order within each flattened patch.
        grid = self.conv_proj(images)
        return grid.flatten(2).transpose(1, 2)  # [B, N, D], row-major patches

    # endregion: patches

    # region: tokens
    def embed_tokens(self, patches):
        cls = self.class_token.expand(patches.shape[0], -1, -1)
        tokens = torch.cat([cls, patches], dim=1)
        return tokens + self.encoder.pos_embedding  # [B, N+1, D]

    # endregion: tokens

    # region: readout
    def forward(self, images):
        patches = self.patch_tokens(images)
        tokens = self.embed_tokens(patches)
        encoded = self.encoder(tokens)  # L blocks, then final LayerNorm
        image_feature = encoded[:, 0]  # CLS, not the mean of patch tokens
        return self.heads(image_feature)  # Raw class logits

    # endregion: readout


# region: training
def training_step(model, images, labels, optimizer):
    """One supervised update; scheduling and real datasets belong to the caller."""
    model.train()
    optimizer.zero_grad(set_to_none=True)
    logits = model(images)
    loss = F.cross_entropy(logits, labels)  # Do not apply softmax first.
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    return loss.detach().item()


# endregion: training


# region: positions
def resize_positions(positions, old_grid, new_grid):
    """Bilinearly resize patch positions, leaving the CLS position untouched."""
    if any(
        len(grid) != 2 or any(not isinstance(n, int) or n <= 0 for n in grid)
        for grid in (old_grid, new_grid)
    ):
        raise ValueError(
            "Position grids must contain two positive integer sizes"
        )
    if positions.ndim != 3 or positions.shape[0] != 1:
        raise ValueError("Expected positions with shape [1, N+1, D]")
    if positions.shape[1] != old_grid[0] * old_grid[1] + 1:
        raise ValueError(
            "Old grid does not match the number of patch positions"
        )
    cls_position, patch_positions = positions[:, :1], positions[:, 1:]
    grid = patch_positions.reshape(1, *old_grid, positions.shape[-1])
    grid = grid.permute(0, 3, 1, 2)
    grid = F.interpolate(
        grid, size=new_grid, mode="bilinear", align_corners=True
    )
    resized = grid.flatten(2).transpose(1, 2)
    return torch.cat([cls_position, resized], dim=1)


# endregion: positions


def checkpoint_keys(state):
    """Migrate the released V1 MLP names, as torchvision's loader also does."""
    migrated = {}
    for key, value in state.items():
        key = key.replace(".mlp.linear_1.", ".mlp.0.")
        key = key.replace(".mlp.linear_2.", ".mlp.3.")
        if key in migrated:
            raise ValueError(
                "Checkpoint contains both legacy and current MLP keys"
            )
        migrated[key] = value
    return migrated


# region: inference
def load_pretrained():
    """Load an explicit ImageNet-1k checkpoint into this authored architecture."""
    weights = ViT_B_16_Weights.IMAGENET1K_V1
    model = VisionTransformer()
    state = weights.get_state_dict(progress=True, check_hash=True)
    model.load_state_dict(checkpoint_keys(state), strict=True)
    return model.cpu().float().eval(), weights


def prepare_image(image, weights):
    """Return exactly what the checkpoint sees, before and after normalization."""
    transform = weights.transforms()
    resized = TF.resize(
        image.convert("RGB"),
        transform.resize_size,
        interpolation=transform.interpolation,
        antialias=transform.antialias,
    )
    crop = TF.center_crop(resized, transform.crop_size)
    pixels = TF.convert_image_dtype(TF.pil_to_tensor(crop), torch.float32)
    batch = TF.normalize(pixels, transform.mean, transform.std).unsqueeze(0)
    return crop, batch


@torch.inference_mode()
def predict(model, weights, image):
    crop, batch = prepare_image(image, weights)
    logits = model(batch)
    scores, indices = logits.softmax(dim=-1)[0].topk(5)
    top5 = [
        {
            "index": int(i),
            "label": weights.meta["categories"][int(i)],
            "score": float(s),
        }
        for s, i in zip(scores, indices)
    ]
    return crop, batch, logits, top5


# endregion: inference


def smoke():
    """Small deterministic mechanics check; random pixels have no useful labels."""
    torch.manual_seed(17)
    model = VisionTransformer(32, 4, 2, 4, 64, 128, num_classes=10)
    images, labels = torch.randn(2, 3, 32, 32), torch.tensor([1, 7])
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    losses = [training_step(model, images, labels, optimizer) for _ in range(2)]
    positions = resize_positions(model.encoder.pos_embedding, (8, 8), (12, 12))
    return {
        "mode": "random-weight mechanics",
        "patch_shape": list(model.patch_tokens(images).shape),
        "token_shape": list(model.embed_tokens(model.patch_tokens(images)).shape),
        "logit_shape": list(model(images).shape),
        "losses": losses,
        "resized_position_shape": list(positions.shape),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--smoke", action="store_true")
    mode.add_argument("--predict", action="store_true")
    parser.add_argument("--image", type=Path)
    parser.add_argument("--output-dir", type=Path)
    args = parser.parse_args()
    torch.set_num_threads(1)
    if args.smoke:
        result = smoke()
    else:
        if args.image is None:
            parser.error("--predict requires --image")
        model, weights = load_pretrained()
        with Image.open(args.image) as image:
            crop, batch, _, top5 = predict(model, weights, image)
        result = {
            "weights": str(weights),
            "input_shape": list(batch.shape),
            "top5": top5,
        }
        if args.output_dir:
            args.output_dir.mkdir(parents=True, exist_ok=True)
            crop.save(args.output_dir / "checkpoint-input.png")
            (args.output_dir / "prediction.json").write_text(
                json.dumps(result, indent=2) + "\n"
            )
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()