# examples/mae_training.py
"""Native PyTorch MAE mechanics, offline smoke, and official-weight inference.

Architecture follows facebookresearch/mae at efb2a8062c206524e35e47d04501ed4f544c0ae8.
The CPU smoke uses a reduced random model; export uses the ViT-L visualization
release with raw ImageNet-normalized pixel targets, not normalized-patch targets.
"""

import argparse
import json
import math

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


def sincos_positions(grid_size, width):
    """Returns fixed 2-D positions in official x-then-y channel order."""
    if width % 4:
        raise ValueError('Position width must be divisible by four.')
    axis = torch.arange(grid_size, dtype=torch.float64)
    yy, xx = torch.meshgrid(axis, axis, indexing='ij')
    frequency = 1 / (10000 ** (torch.arange(width // 4).double() / (width // 4)))
    parts = []
    for coordinates in (xx, yy):
        phase = coordinates.flatten()[:, None] * frequency[None, :]
        parts.extend([phase.sin(), phase.cos()])
    return torch.cat([torch.zeros(1, width), torch.cat(parts, dim=1).float()])[None]


class PatchEmbed(nn.Module):
    """Non-overlapping RGB projection with checkpoint-compatible names."""

    def __init__(self, img_size, patch_size, in_chans, embed_dim):
        super().__init__()
        self.patch_size = (patch_size, patch_size)
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_chans, embed_dim, patch_size, patch_size)

    def forward(self, images):
        return self.proj(images).flatten(2).transpose(1, 2)


class Attention(nn.Module):
    """Multihead self-attention without attention dropout."""

    def __init__(self, width, heads):
        super().__init__()
        self.heads = heads
        self.qkv = nn.Linear(width, width * 3)
        self.proj = nn.Linear(width, width)

    def forward(self, tokens):
        batch, length, width = tokens.shape
        qkv = self.qkv(tokens).reshape(batch, length, 3, self.heads, width // self.heads)
        query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        output = F.scaled_dot_product_attention(query, key, value)
        return self.proj(output.transpose(1, 2).reshape(batch, length, width))


class Mlp(nn.Module):
    """Checkpoint-compatible GELU feed-forward network."""

    def __init__(self, width, hidden):
        super().__init__()
        self.fc1 = nn.Linear(width, hidden)
        self.fc2 = nn.Linear(hidden, width)

    def forward(self, tokens):
        return self.fc2(F.gelu(self.fc1(tokens)))


class Block(nn.Module):
    """Pre-norm transformer block, matching the released timm layout."""

    def __init__(self, dim, num_heads, mlp_ratio=4, qkv_bias=True,
                 qk_scale=None, norm_layer=None):
        super().__init__()
        if not qkv_bias or qk_scale is not None:
            raise ValueError('Only the official MAE attention configuration is supported.')
        norm_layer = norm_layer or (lambda width: nn.LayerNorm(width, eps=1e-6))
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads)
        self.norm2 = norm_layer(dim)
        self.mlp = Mlp(dim, int(dim * mlp_ratio))

    def forward(self, tokens):
        tokens = tokens + self.attn(self.norm1(tokens))
        return tokens + self.mlp(self.norm2(tokens))


class MAE(nn.Module):
    """Visible-only encoder and full-grid decoder with fixed positions.

    Args:
        image_size: Square input side length, divisible by patch_size.
        patch_size: Square patch side length.
        width: Encoder embedding width.
        depth: Number of encoder blocks.
        heads: Encoder attention head count.
        decoder_width: Decoder embedding width.
        decoder_depth: Number of decoder blocks.
        decoder_heads: Decoder attention head count.
        norm_pix_loss: Normalize each target patch using sample variance.
    """

    def __init__(self, image_size=32, patch_size=8, width=32, depth=2,
                 heads=4, decoder_width=24, decoder_depth=1,
                 decoder_heads=4, norm_pix_loss=True):
        super().__init__()
        dimensions = (image_size, patch_size, width, depth, heads,
                      decoder_width, decoder_depth, decoder_heads)
        if any(not isinstance(value, int) or value <= 0 for value in dimensions):
            raise ValueError('All architecture dimensions must be positive integers.')
        if image_size % patch_size or width % heads or decoder_width % decoder_heads:
            raise ValueError('Image/patch sizes and embedding/head counts must divide evenly.')
        if width % 4 or decoder_width % 4:
            raise ValueError('Position embedding widths must be divisible by four.')
        self.image_size = image_size
        self.patch_embed = PatchEmbed(image_size, patch_size, 3, width)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, width))
        self.pos_embed = nn.Parameter(sincos_positions(image_size // patch_size, width),
                                      requires_grad=False)
        self.blocks = nn.ModuleList([Block(width, heads) for _ in range(depth)])
        self.norm = nn.LayerNorm(width, eps=1e-6)
        self.decoder_embed = nn.Linear(width, decoder_width)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_width))
        self.decoder_pos_embed = nn.Parameter(
            sincos_positions(image_size // patch_size, decoder_width), requires_grad=False)
        self.decoder_blocks = nn.ModuleList(
            [Block(decoder_width, decoder_heads) for _ in range(decoder_depth)])
        self.decoder_norm = nn.LayerNorm(decoder_width, eps=1e-6)
        self.decoder_pred = nn.Linear(decoder_width, patch_size ** 2 * 3)
        self.norm_pix_loss = norm_pix_loss
        self.apply(self._initialize)
        nn.init.xavier_uniform_(self.patch_embed.proj.weight.flatten(1))
        nn.init.normal_(self.cls_token, std=.02)
        nn.init.normal_(self.mask_token, std=.02)

    @staticmethod
    def _initialize(module):
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            nn.init.zeros_(module.bias)

    # region: patches
    def patchify(self, images):
        """Maps BCHW into raster-ordered patches with pixel-then-channel order."""
        p = self.patch_embed.patch_size[0]
        if images.ndim != 4 or images.shape[1] != 3:
            raise ValueError('Expected [B, 3, H, H] RGB images.')
        if images.shape[2:] != (self.image_size, self.image_size):
            raise ValueError('Image dimensions must match the configured square size.')
        b, _, h, _ = images.shape
        grid = h // p
        return images.reshape(b, 3, grid, p, grid, p).permute(
            0, 2, 4, 3, 5, 1).reshape(b, grid * grid, p * p * 3)

    def unpatchify(self, patches):
        """Inverts patchify exactly, preserving pixel/channel order."""
        p = self.patch_embed.patch_size[0]
        grid = self.image_size // p
        if patches.ndim != 3 or patches.shape[1:] != (grid * grid, p * p * 3):
            raise ValueError('Patch count or patch vector width is invalid.')
        b = patches.shape[0]
        return patches.reshape(b, grid, grid, p, p, 3).permute(
            0, 5, 1, 3, 2, 4).reshape(b, 3, grid * p, grid * p)
    # endregion: patches

    # region: masking
    def random_masking(self, tokens, mask_ratio):
        """Samples independent permutations, keeping only visible tokens."""
        if not math.isfinite(mask_ratio) or not 0 < mask_ratio < 1:
            raise ValueError('Mask ratio must be finite and strictly between zero and one.')
        b, length, width = tokens.shape
        keep = int(length * (1 - mask_ratio))
        if not 0 < keep < length:
            raise ValueError('Mask ratio must leave at least one visible and one masked patch.')
        order = torch.rand(b, length, device=tokens.device).argsort(dim=1)
        restore = order.argsort(dim=1)
        visible = tokens.gather(1, order[:, :keep, None].expand(-1, -1, width))
        mask = tokens.new_ones(b, length)
        mask[:, :keep] = 0
        return visible, mask.gather(1, restore), restore
    # endregion: masking

    # region: encoder
    def forward_encoder(self, images, mask_ratio):
        self.patchify(images)  # Validate the input before convolution.
        tokens = self.patch_embed(images) + self.pos_embed[:, 1:]
        visible, mask, restore = self.random_masking(tokens, mask_ratio)
        cls = (self.cls_token + self.pos_embed[:, :1]).expand(images.shape[0], -1, -1)
        latent = torch.cat([cls, visible], dim=1)
        for block in self.blocks:
            latent = block(latent)
        return self.norm(latent), mask, restore
    # endregion: encoder

    # region: restore
    def restore_tokens(self, latent, restore):
        projected = self.decoder_embed(latent)
        missing = restore.shape[1] + 1 - projected.shape[1]
        masks = self.mask_token.expand(projected.shape[0], missing, -1)
        shuffled = torch.cat([projected[:, 1:], masks], dim=1)
        patches = shuffled.gather(1, restore[..., None].expand(-1, -1, projected.shape[2]))
        return torch.cat([projected[:, :1], patches], dim=1)
    # endregion: restore

    # region: decoder
    def forward_decoder(self, latent, restore):
        tokens = self.restore_tokens(latent, restore) + self.decoder_pos_embed
        for block in self.decoder_blocks:
            tokens = block(tokens)
        return self.decoder_pred(self.decoder_norm(tokens))[:, 1:]
    # endregion: decoder

    # region: loss
    def forward_loss(self, images, prediction, mask):
        target = self.patchify(images)
        if prediction.shape != target.shape or mask.shape != target.shape[:2]:
            raise ValueError('Prediction and mask dimensions must match target patches.')
        if not torch.all((mask == 0) | (mask == 1)) or mask.sum() == 0:
            raise ValueError('Mask must be binary with at least one removed patch.')
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            variance = target.var(dim=-1, keepdim=True, correction=1)
            target = (target - mean) / (variance + 1e-6).sqrt()
        patch_error = (prediction - target).square().mean(dim=-1)
        return (patch_error * mask).sum() / mask.sum()
    # endregion: loss

    def forward(self, images, mask_ratio=.75):
        latent, mask, restore = self.forward_encoder(images, mask_ratio)
        prediction = self.forward_decoder(latent, restore)
        return self.forward_loss(images, prediction, mask), prediction, mask

    # region: reconstruction
    def reconstruct(self, images, prediction, mask):
        """Produces a visible-patch composite only for raw-pixel target models."""
        if self.norm_pix_loss:
            raise ValueError('Display requires a model trained on raw pixel targets.')
        expanded_mask = mask[..., None]
        composite = self.patchify(images) * (1 - expanded_mask) + prediction * expanded_mask
        return self.unpatchify(composite)
    # endregion: reconstruction


# region: update
def train_step(model, images, optimizer):
    """Takes one reconstruction update; both encoder and decoder learn."""
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss, prediction, mask = model(images, mask_ratio=.75)
    loss.backward()
    optimizer.step()
    return loss.detach(), prediction.detach(), mask.detach()
# endregion: update


def smoke():
    """Runs deterministic reduced random-weight training without downloads."""
    torch.set_num_threads(2)
    torch.manual_seed(7)
    model = MAE()
    images = torch.rand(2, 3, 32, 32)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss, prediction, mask = train_step(model, images, optimizer)
    return {'scope': 'reduced random-weight CPU update; no pretrained semantics',
            'seed': 7, 'loss': loss.item(), 'predictionShape': list(prediction.shape),
            'maskedPatchesPerImage': mask.sum(1).tolist(),
            'encoderTokensIncludingCls': 5, 'decoderTokensIncludingCls': 17}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--smoke', action='store_true', help='Run an offline CPU update.')
    args = parser.parse_args()
    if not args.smoke:
        parser.error('Use --smoke; pretrained exports use mae_export.py.')
    print(json.dumps(smoke(), indent=2))


if __name__ == '__main__':
    main()