# examples/swin_transformer_training.py
"""Execute reduced Swin V1 mechanics without checkpoints or network access.

Window/bias/mask/merge order follows Microsoft Swin V1 revision
f82860bfb5225915aca09c3227159ee9e1df874d (MIT, LICENSE-swin.txt).
Dynamic bottom/right padding and per-axis shift disablement follow the later
TorchVision V1 port; original fixed-resolution PatchMerging requires even axes.
Smoke uses random weights, width 24, depths (2,2,2,2), windows 4 and five classes.
It omits augmentation, dropout/stochastic depth and the paper training schedule.
"""

import argparse
import json

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


def validate_grid(x):
    """Require a nonempty channels-last feature grid."""
    if x.ndim != 4 or min(x.shape) <= 0:
        raise ValueError("Expected nonempty [B,H,W,C] features")


# region: patches
class PatchEmbed(nn.Module):
    """Project non-overlapping 4x4 RGB patches, then normalize channels."""

    def __init__(self, dim):
        super().__init__()
        if dim <= 0:
            raise ValueError("Embedding width must be positive")
        self.projection = nn.Conv2d(3, dim, kernel_size=4, stride=4)
        self.norm = nn.LayerNorm(dim)

    def forward(self, images):
        if images.ndim != 4 or images.shape[1] != 3:
            raise ValueError("Expected [B,3,H,W] RGB images")
        if min(images.shape) <= 0 or any(n % 4 for n in images.shape[-2:]):
            raise ValueError("Image axes must be positive multiples of four")
        return self.norm(self.projection(images).permute(0, 2, 3, 1))


# endregion: patches


# region: windows
def window_partition(x, size):
    """Reorder [B,H,W,C] into [B*nW,M*M,C] windows without mixing pixels."""
    validate_grid(x)
    b, h, w, c = x.shape
    if size <= 0 or h % size or w % size:
        raise ValueError("Window size must divide both grid axes")
    grouped = x.reshape(b, h // size, size, w // size, size, c)
    return grouped.permute(0, 1, 3, 2, 4, 5).reshape(-1, size * size, c)


def window_reverse(windows, size, height, width):
    """Invert partition exactly for a known divisible spatial grid."""
    if size <= 0 or min(height, width) <= 0 or height % size or width % size:
        raise ValueError("Reverse axes must be positive multiples of window size")
    count = (height // size) * (width // size)
    if windows.ndim != 3 or windows.shape[1] != size * size:
        raise ValueError("Expected [B*nW,M*M,C] windows")
    if min(windows.shape) <= 0 or windows.shape[0] % count:
        raise ValueError("Window count cannot reconstruct this grid")
    grouped = windows.reshape(
        -1, height // size, width // size, size, size, windows.shape[-1]
    )
    return grouped.permute(0, 1, 3, 2, 4, 5).reshape(
        -1, height, width, windows.shape[-1]
    )


# endregion: windows


# region: relative
def relative_position_index(size):
    """Map signed query-minus-key offsets to (2M-1)^2 table entries."""
    if size <= 0:
        raise ValueError("Window size must be positive")
    coordinates = torch.stack(
        torch.meshgrid(torch.arange(size), torch.arange(size), indexing="ij")
    )
    coordinates = coordinates.flatten(1)
    offsets = coordinates[:, :, None] - coordinates[:, None, :]
    offsets = offsets.permute(1, 2, 0) + size - 1
    return offsets[..., 0] * (2 * size - 1) + offsets[..., 1]


def relative_bias(table, index):
    """Gather an independent learned bias for each head and ordered pair."""
    n = index.shape[0]
    return table[index.reshape(-1)].reshape(n, n, -1).permute(2, 0, 1)


# endregion: relative


# region: mask
def shifted_grid(x, size, shift):
    """Pad then disable a shift independently along each single-window axis."""
    validate_grid(x)
    if size <= 0 or not 0 <= shift < size:
        raise ValueError("Require 0 <= shift < positive window size")
    height, width = x.shape[1:3]
    x = F.pad(x, (0, 0, 0, (-width) % size, 0, (-height) % size))
    sy = shift if x.shape[1] > size else 0
    sx = shift if x.shape[2] > size else 0
    return torch.roll(x, shifts=(-sy, -sx), dims=(1, 2)), (sy, sx)


def region_mask(height, width, size, shifts, device=None, dtype=torch.float32):
    """Use finite -100 between cyclic regions, zero within each region."""
    sy, sx = shifts
    regions = torch.zeros((1, height, width, 1), device=device)
    ys = (
        (slice(0, -size), slice(-size, -sy), slice(-sy, None)) if sy else (slice(None),)
    )
    xs = (
        (slice(0, -size), slice(-size, -sx), slice(-sx, None)) if sx else (slice(None),)
    )
    label = 0
    for y in ys:
        for x in xs:
            regions[:, y, x, :] = label
            label += 1
    labels = window_partition(regions, size).squeeze(-1)
    different = labels[:, :, None] != labels[:, None, :]
    return different.to(dtype) * -100.0


# endregion: mask


class WindowAttention(nn.Module):
    """V1 scaled dot-product attention with relative bias and cyclic mask."""

    def __init__(self, dim, heads, size=4, shift=0):
        super().__init__()
        if min(dim, heads, size) <= 0 or dim % heads or not 0 <= shift < size:
            raise ValueError("Invalid width, heads, window size or shift")
        self.heads, self.size, self.shift = heads, size, shift
        self.qkv = nn.Linear(dim, 3 * dim)
        self.projection = nn.Linear(dim, dim)
        self.bias_table = nn.Parameter(torch.empty((2 * size - 1) ** 2, heads))
        nn.init.trunc_normal_(self.bias_table, std=0.02)
        self.register_buffer("index", relative_position_index(size))

    # region: attention
    def forward(self, x):
        height, width = x.shape[1:3]
        shifted, shifts = shifted_grid(x, self.size, self.shift)
        ph, pw = shifted.shape[1:3]
        windows = window_partition(shifted, self.size)
        batch_windows, n, c = windows.shape
        qkv = self.qkv(windows).reshape(
            batch_windows, n, 3, self.heads, c // self.heads
        )
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        logits = (q * ((c // self.heads) ** -0.5)) @ k.transpose(-2, -1)
        logits = logits + relative_bias(self.bias_table, self.index)
        mask = region_mask(ph, pw, self.size, shifts, x.device, logits.dtype)
        nw = mask.shape[0]
        logits = logits.reshape(-1, nw, self.heads, n, n) + mask[None, :, None]
        weights = logits.reshape(batch_windows, self.heads, n, n).softmax(-1)
        attended = (weights @ v).transpose(1, 2).reshape(batch_windows, n, c)
        restored = window_reverse(self.projection(attended), self.size, ph, pw)
        restored = torch.roll(restored, shifts=shifts, dims=(1, 2))
        return restored[:, :height, :width, :].contiguous()

    # endregion: attention


# region: block
class SwinBlock(nn.Module):
    """Pre-normalized V1 attention/residual and GELU MLP/residual."""

    def __init__(self, dim, heads, size=4, shift=0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attention = WindowAttention(dim, heads, size, shift)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
        )

    def forward(self, x):
        x = x + self.attention(self.norm1(x))
        return x + self.mlp(self.norm2(x))


# endregion: block


# region: merge
def merge_neighbors(x):
    """Concatenate x00,x10,x01,x11; dynamic port pads odd spatial axes."""
    validate_grid(x)
    height, width = x.shape[1:3]
    x = F.pad(x, (0, 0, 0, width % 2, 0, height % 2))
    return torch.cat(
        (x[:, 0::2, 0::2], x[:, 1::2, 0::2], x[:, 0::2, 1::2], x[:, 1::2, 1::2]), dim=-1
    )


class PatchMerging(nn.Module):
    """Normalize the four-neighbor vector before its 4C-to-2C projection."""

    def __init__(self, dim):
        super().__init__()
        if dim <= 0:
            raise ValueError("Channel width must be positive")
        self.norm = nn.LayerNorm(4 * dim)
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)

    def forward(self, x):
        return self.reduction(self.norm(merge_neighbors(x)))


# endregion: merge


# region: classifier
class SwinMechanics(nn.Module):
    """Four reduced stages, alternating regular and shifted blocks."""

    def __init__(
        self, dim=24, depths=(2, 2, 2, 2), heads=(3, 6, 12, 24), size=4, classes=5
    ):
        super().__init__()
        if len(depths) != 4 or len(heads) != 4 or min(depths) <= 0 or classes <= 0:
            raise ValueError("Require four nonempty stages and positive classes")
        self.patch_embed = PatchEmbed(dim)
        self.stages, self.merges = nn.ModuleList(), nn.ModuleList()
        for stage, depth in enumerate(depths):
            width = dim * 2**stage
            self.stages.append(
                nn.Sequential(
                    *[
                        SwinBlock(
                            width,
                            heads[stage],
                            size,
                            0 if block % 2 == 0 else size // 2,
                        )
                        for block in range(depth)
                    ]
                )
            )
            if stage < 3:
                self.merges.append(PatchMerging(width))
        self.norm = nn.LayerNorm(8 * dim)
        self.head = nn.Linear(8 * dim, classes)

    def forward(self, images):
        x = self.patch_embed(images)
        shapes = []
        for stage, blocks in enumerate(self.stages):
            x = blocks(x)
            shapes.append(list(x.shape))
            if stage < 3:
                x = self.merges[stage](x)
        logits = self.head(self.norm(x).mean(dim=(1, 2)))
        return logits, shapes


def train_step(model, optimizer, images, labels):
    """Update the entire random-weight hierarchy with supervised cross entropy."""
    logits, shapes = model(images)
    if labels.shape != (images.shape[0],) or labels.dtype != torch.long:
        raise ValueError("Require one integer class label per image")
    if labels.numel() == 0 or labels.min() < 0 or labels.max() >= logits.shape[1]:
        raise ValueError("Class labels are outside the classifier vocabulary")
    optimizer.zero_grad(set_to_none=True)
    loss = F.cross_entropy(logits, labels)
    loss.backward()
    optimizer.step()
    return loss.detach(), shapes


# endregion: classifier


def smoke():
    """Run one deterministic CPU training step, explicitly without semantics."""
    torch.set_num_threads(2)
    torch.manual_seed(13)
    model = SwinMechanics()
    images = torch.randn(2, 3, 64, 64)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss, shapes = train_step(model, optimizer, images, torch.tensor([0, 4]))
    return {
        "scope": "reduced random-weight Swin V1 mechanics, not pretrained",
        "input": list(images.shape),
        "stages": shapes,
        "loss": loss.item(),
        "finite_gradients": all(
            p.grad is not None and torch.isfinite(p.grad).all().item()
            for p in model.parameters()
        ),
    }


def main():
    """Expose the offline CPU smoke check."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.parse_args()
    print(json.dumps(smoke(), indent=2))


if __name__ == "__main__":
    main()