# examples/vitdet.py
"""Seeded offline ViTDet backbone/neck mechanics, stopping before a detector.

Conventions: Detectron2 1e3e13bbf607b54f62205c4c33922521822fb298,
detectron2/modeling/backbone/vit.py (Attention, Block, SimpleFeaturePyramid)
and utils.py (window helpers, get_abs_pos, get_rel_pos). Authored standalone
implementation; no Detectron2 dependency or upstream code download at runtime.
Paper: https://arxiv.org/html/2203.16527v2, section 3 and appendix A.2.
Padding is zero in BHWC after norm1, participates in attention without a mask,
and is cropped before residual addition. This is the pinned release behavior.
"""

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

REVISION = "1e3e13bbf607b54f62205c4c33922521822fb298"


def positive_int(value):
    """Return whether value is an integer size, excluding booleans."""
    return isinstance(value, int) and not isinstance(value, bool) and value > 0


def validate_grid(x):
    """Reject empty, nonfinite or non-floating four-axis spatial tensors."""
    if x.ndim != 4 or min(x.shape) < 1 or not x.is_floating_point():
        raise ValueError("Expected a nonempty floating four-axis tensor")
    if not torch.isfinite(x).all():
        raise ValueError("Spatial tensors must be finite")


# region: patch_grid
class PatchGrid(nn.Module):
    """Project nonoverlapping 16-pixel patches into a BHWC grid.

    Inputs need at least one full patch; incomplete right/bottom pixels are
    discarded by the unpadded convolution, as in the release. Detection inputs
    normally arrive padded to stride 32 upstream; this example does not invent
    image padding inside the patch projector.
    """

    def __init__(self, channels=8):
        super().__init__()
        if not positive_int(channels):
            raise ValueError("channels must be positive")
        self.proj = nn.Conv2d(3, channels, kernel_size=16, stride=16)

    def forward(self, image):
        """Return [B, floor(H/16), floor(W/16), C] patch embeddings."""
        validate_grid(image)
        if image.shape[1] != 3 or min(image.shape[-2:]) < 16:
            raise ValueError(
                "Expected RGB images with at least one 16x16 patch"
            )
        return self.proj(image).permute(0, 2, 3, 1)


def get_abs_pos(positions, has_cls_token, hw):
    """Remove the pretraining CLS position and bicubic-resize a square grid."""
    if len(hw) != 2 or not all(positive_int(size) for size in hw):
        raise ValueError("Target height and width must be positive integers")
    if positions.ndim != 3 or positions.shape[0] != 1:
        raise ValueError("Absolute positions must have shape [1, N, C]")
    spatial = positions[:, 1:] if has_cls_token else positions
    side = math.isqrt(spatial.shape[1])
    if side < 1 or side * side != spatial.shape[1] or spatial.shape[2] < 1:
        raise ValueError("Pretraining spatial positions must form a square")
    grid = spatial.reshape(1, side, side, -1)
    if (side, side) == tuple(hw):
        return grid
    resized = F.interpolate(
        grid.permute(0, 3, 1, 2),
        size=hw,
        mode="bicubic",
        align_corners=False,
    )
    return resized.permute(0, 2, 3, 1)


def feature_map(tokens):
    """Expose only spatial tokens as an NCHW detector-interface map."""
    validate_grid(tokens)
    return tokens.permute(0, 3, 1, 2)


# endregion: patch_grid


# region: window_partition
def window_partition(x, window_side=14):
    """Pad BHWC at bottom/right with zeros and return row-major windows.

    Returns:
        Windows [B*ceil(H/s)*ceil(W/s), s, s, C] and padded (Hp, Wp).
        No validity mask is returned: padded tokens enter ordinary attention.
    """
    validate_grid(x)
    if not positive_int(window_side):
        raise ValueError("window_side must be a positive integer")
    b, h, w, c = x.shape
    pad_h, pad_w = (-h) % window_side, (-w) % window_side
    padded = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
    hp, wp = h + pad_h, w + pad_w
    tiles = padded.reshape(
        b,
        hp // window_side,
        window_side,
        wp // window_side,
        window_side,
        c,
    )
    windows = tiles.permute(0, 1, 3, 2, 4, 5).reshape(
        -1, window_side, window_side, c
    )
    return windows, (hp, wp)


# endregion: window_partition


# region: window_attention
def get_rel_pos(q_size, k_size, table):
    """Linearly resize a relative-position table and gather scaled offsets."""
    if not positive_int(q_size) or not positive_int(k_size):
        raise ValueError("Query/key sizes must be positive integers")
    if table.ndim != 2 or min(table.shape) < 1:
        raise ValueError("Relative positions must have shape [L, head_dim]")
    length = 2 * max(q_size, k_size) - 1
    if table.shape[0] != length:
        table = F.interpolate(
            table.T.unsqueeze(0),
            size=length,
            mode="linear",
            align_corners=False,
        )[0].T
    q = torch.arange(q_size, device=table.device)[:, None]
    k = torch.arange(k_size, device=table.device)[None, :]
    offsets = q * max(k_size / q_size, 1.0)
    offsets = offsets - k * max(q_size / k_size, 1.0)
    offsets = offsets + (k_size - 1) * max(q_size / k_size, 1.0)
    return table[offsets.long()]


def add_rel_pos(scores, q, rel_h, rel_w, hw):
    """Add query-dependent decomposed height/width terms to self-attention."""
    h, w = hw
    rh, rw = get_rel_pos(h, h, rel_h), get_rel_pos(w, w, rel_w)
    grid = q.reshape(q.shape[0], h, w, -1)
    height = torch.einsum("bhwc,hkc->bhwk", grid, rh)
    width = torch.einsum("bhwc,wkc->bhwk", grid, rw)
    expanded = scores.reshape(q.shape[0], h, w, h, w)
    return (expanded + height[..., None] + width[..., None, :]).reshape_as(
        scores
    )


class Attention(nn.Module):
    """Ordinary multi-head spatial attention, including padded tokens."""

    def __init__(
        self, channels=8, heads=2, use_rel_pos=False, relative_grid=(14, 14)
    ):
        super().__init__()
        if (
            not positive_int(channels)
            or not positive_int(heads)
            or channels % heads
        ):
            raise ValueError("Positive channels must be divisible by heads")
        if len(relative_grid) != 2 or not all(
            positive_int(s) for s in relative_grid
        ):
            raise ValueError("relative_grid must contain two positive sizes")
        self.channels, self.heads = channels, heads
        self.scale = (channels // heads) ** -0.5
        self.qkv = nn.Linear(channels, 3 * channels, bias=True)
        self.proj = nn.Linear(channels, channels, bias=True)
        self.use_rel_pos = use_rel_pos
        if use_rel_pos:
            self.rel_h = nn.Parameter(
                torch.zeros(2 * relative_grid[0] - 1, channels // heads)
            )
            self.rel_w = nn.Parameter(
                torch.zeros(2 * relative_grid[1] - 1, channels // heads)
            )

    def forward(self, x):
        """Return BHWC with scaled QK scores and optional relative positions."""
        validate_grid(x)
        b, h, w, c = x.shape
        if c != self.channels:
            raise ValueError(
                "Input channels disagree with attention projection"
            )
        qkv = self.qkv(x).reshape(b, h * w, 3, self.heads, c // self.heads)
        q, k, v = (
            qkv.permute(2, 0, 3, 1, 4)
            .reshape(3, b * self.heads, h * w, c // self.heads)
            .unbind(0)
        )
        scores = (q * self.scale) @ k.transpose(-2, -1)
        if self.use_rel_pos:
            scores = add_rel_pos(scores, q, self.rel_h, self.rel_w, (h, w))
        values = scores.softmax(-1) @ v
        merged = values.reshape(b, self.heads, h, w, -1)
        merged = merged.permute(0, 2, 3, 1, 4).reshape(b, h, w, c)
        return self.proj(merged)


# endregion: window_attention


# region: window_unpartition
def window_unpartition(windows, window_side, padded_hw, original_hw):
    """Restore row-major BHWC windows and crop bottom/right padding exactly."""
    validate_grid(windows)
    if not positive_int(window_side):
        raise ValueError("window_side must be a positive integer")
    if len(padded_hw) != 2 or len(original_hw) != 2:
        raise ValueError("Geometry must contain height and width")
    hp, wp = padded_hw
    h, w = original_hw
    if not all(positive_int(s) for s in (hp, wp, h, w)):
        raise ValueError("Geometry sizes must be positive integers")
    if hp != h + (-h) % window_side or wp != w + (-w) % window_side:
        raise ValueError("Padded geometry must match minimal window padding")
    if tuple(windows.shape[1:3]) != (window_side, window_side):
        raise ValueError("Window tensor has the wrong tile size")
    per_image = (hp // window_side) * (wp // window_side)
    if windows.shape[0] % per_image:
        raise ValueError("Window count does not represent complete images")
    b, c = windows.shape[0] // per_image, windows.shape[-1]
    tiles = windows.reshape(
        b,
        hp // window_side,
        wp // window_side,
        window_side,
        window_side,
        c,
    )
    restored = tiles.permute(0, 1, 3, 2, 4, 5).reshape(b, hp, wp, c)
    return restored[:, :h, :w].contiguous()


# endregion: window_unpartition


# region: global_blocks
def global_slots(depth, count=4):
    """Choose zero-based last blocks of equally sized depth subsets."""
    if not positive_int(depth) or not positive_int(count) or depth % count:
        raise ValueError("depth must be positive and divisible by global count")
    return tuple(range(depth // count - 1, depth, depth // count))


class Block(nn.Module):
    """Pre-LN attention then pre-LN MLP; window_side=0 selects global mixing."""

    def __init__(self, channels=8, heads=2, window_side=14, use_rel_pos=False):
        super().__init__()
        if (
            not isinstance(window_side, int)
            or isinstance(window_side, bool)
            or window_side < 0
        ):
            raise ValueError("window_side must be a nonnegative integer")
        self.window_side = window_side
        self.norm1 = nn.LayerNorm(channels, eps=1e-6)
        self.attn = Attention(channels, heads, use_rel_pos)
        self.norm2 = nn.LayerNorm(channels, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(channels, 4 * channels),
            nn.GELU(),
            nn.Linear(4 * channels, channels),
        )

    def forward(self, x):
        """Pad only normalized spatial tokens; crop before shortcut addition."""
        normalized = self.norm1(x)
        if self.window_side:
            windows, padded = window_partition(normalized, self.window_side)
            attended = window_unpartition(
                self.attn(windows),
                self.window_side,
                padded,
                tuple(x.shape[1:3]),
            )
        else:
            attended = self.attn(normalized)
        x = x + attended
        return x + self.mlp(self.norm2(x))


class Backbone(nn.Module):
    """Reduced plain ViT with one final map, no depth taps or classifier token."""

    def __init__(self, channels=8, heads=2, depth=4, globals_count=2):
        super().__init__()
        self.global_indexes = global_slots(depth, globals_count)
        self.patch = PatchGrid(channels)
        self.positions = nn.Parameter(torch.empty(1, 197, channels))
        nn.init.trunc_normal_(self.positions, std=0.02)
        self.blocks = nn.ModuleList(
            [
                Block(
                    channels,
                    heads,
                    0 if index in self.global_indexes else 14,
                    use_rel_pos=True,
                )
                for index in range(depth)
            ]
        )
        for layer in self.modules():
            if isinstance(layer, nn.Linear):
                nn.init.trunc_normal_(layer.weight, std=0.02)
                nn.init.zeros_(layer.bias)

    def forward(self, image):
        """Return final [B, C, H/16, W/16] when image sides divide by 16."""
        x = self.patch(image)
        x = x + get_abs_pos(self.positions, True, tuple(x.shape[1:3]))
        for block in self.blocks:
            x = block(x)
        return feature_map(x)


# endregion: global_blocks


class ChannelNorm(nn.Module):
    """LayerNorm over channels at each spatial position, with epsilon 1e-6."""

    def __init__(self, channels):
        super().__init__()
        self.norm = nn.LayerNorm(channels, eps=1e-6)

    def forward(self, x):
        """Normalize NCHW channels independently at each pixel."""
        return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)


# region: simple_pyramid
class SimplePyramid(nn.Module):
    """Four independent release-style rescaling branches from ONE final map.

    Patch stride 16 yields p2/p3/p4/p5 at nominal strides 4/8/16/32.
    Output channels are reduced to 4, rather than the paper's 256.
    No lateral/top-down connections and no optional fifth top-block level.
    """

    strides = {"p2": 4, "p3": 8, "p4": 16, "p5": 32}

    def __init__(self, channels=8, output_channels=4):
        super().__init__()
        if (
            not positive_int(channels)
            or channels % 4
            or not positive_int(output_channels)
        ):
            raise ValueError(
                "Input width must divide by 4; output width must be positive"
            )
        self.channels = channels
        self.stages = nn.ModuleList()
        rescalers = [
            (
                [
                    nn.ConvTranspose2d(channels, channels // 2, 2, 2),
                    ChannelNorm(channels // 2),
                    nn.GELU(),
                    nn.ConvTranspose2d(channels // 2, channels // 4, 2, 2),
                ],
                channels // 4,
            ),
            (
                [nn.ConvTranspose2d(channels, channels // 2, 2, 2)],
                channels // 2,
            ),
            ([], channels),
            ([nn.MaxPool2d(2, 2)], channels),
        ]
        for rescale, width in rescalers:
            self.stages.append(
                nn.Sequential(
                    *rescale,
                    nn.Conv2d(width, output_channels, 1, bias=False),
                    ChannelNorm(output_channels),
                    nn.Conv2d(
                        output_channels,
                        output_channels,
                        3,
                        padding=1,
                        bias=False,
                    ),
                    ChannelNorm(output_channels),
                )
            )

    def forward(self, final_map):
        """Return the detector interface; odd sides use floor max-pool sizes."""
        validate_grid(final_map)
        if final_map.shape[1] != self.channels or min(final_map.shape[-2:]) < 2:
            raise ValueError("Expected correct channels and spatial sides >= 2")
        return {
            name: stage(final_map)
            for name, stage in zip(self.strides, self.stages)
        }


# endregion: simple_pyramid


def smoke(seed=7):
    """Execute synthetic CPU mechanics and report actual finite measurements."""
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    backbone, neck = Backbone(), SimplePyramid()
    image = torch.randn(1, 3, 256, 288)
    grid = backbone.patch(image)
    windows, padded = window_partition(grid)
    restored = window_unpartition(windows, 14, padded, (16, 18))
    final = backbone(image)
    pyramid = neck(final)
    # Arbitrary spatial targets avoid a nearly constant LayerNorm energy loss.
    targets = {name: torch.randn_like(value) for name, value in pyramid.items()}
    loss = sum(
        (value - targets[name]).square().mean()
        for name, value in pyramid.items()
    )
    parameters = list(backbone.parameters()) + list(neck.parameters())
    before = [parameter.detach().clone() for parameter in parameters]
    optimizer = torch.optim.SGD(parameters, lr=0.01)
    loss.backward()
    gradient_norm = math.sqrt(
        sum(
            parameter.grad.square().sum().item()
            for parameter in parameters
            if parameter.grad is not None
        )
    )
    optimizer.step()
    backbone.eval()
    neck.eval()
    update_norm = math.sqrt(
        sum(
            (parameter - old).square().sum().item()
            for parameter, old in zip(parameters, before)
        )
    )
    # Identical blocks contrast locality using the same weights and token change.
    local, global_block = Block(8, 2, 14), Block(8, 2, 0)
    global_block.load_state_dict(local.state_dict())
    changed = grid.detach().clone()
    changed[0, 15, 17, 0] += 3
    with torch.no_grad():
        local_delta = (
            (local(grid)[0, 0, 0] - local(changed)[0, 0, 0]).abs().max().item()
        )
        global_delta = (
            (global_block(grid)[0, 0, 0] - global_block(changed)[0, 0, 0])
            .abs()
            .max()
            .item()
        )
        abs_positions = get_abs_pos(backbone.positions, True, (16, 18))
        rel_positions = get_rel_pos(16, 18, backbone.blocks[0].attn.rel_h)
        repeat_error = (backbone(image) - backbone(image)).abs().max().item()
    receipt = {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "upstream_revision": REVISION,
        "seed": seed,
        "device": "cpu",
        "torch_version": torch.__version__,
        "reductions": [
            "synthetic 256x288 RGB",
            "width 8, 2 heads, 4 blocks",
            "2 global blocks instead of original 4",
            "pyramid width 4 instead of 256",
            "zero stochastic depth",
            "relative tables use a synthetic 14x14 initial grid",
            "no detector head, checkpoint, dataset or AP",
        ],
        "patch_size": 16,
        "window_side": 14,
        "tokens_per_window": 196,
        "image_shape": list(image.shape),
        "grid_shape": list(grid.shape),
        "window_shape": list(windows.shape),
        "padded_hw": list(padded),
        "final_shape": list(final.shape),
        "pyramid_shapes": {
            name: list(value.shape) for name, value in pyramid.items()
        },
        "strides": neck.strides,
        "global_indexes": list(backbone.global_indexes),
        "original_vit_l_global_indexes": list(global_slots(24)),
        "absolute_position_shape": list(abs_positions.shape),
        "relative_position_shape": list(rel_positions.shape),
        "round_trip_max_error": (grid - restored).abs().max().item(),
        "local_distant_delta": local_delta,
        "global_distant_delta": global_delta,
        "loss": loss.item(),
        "gradient_norm": gradient_norm,
        "update_norm": update_norm,
        "repeat_eval_error": repeat_error,
    }
    return receipt


def main():
    """Print JSON only for an explicitly requested smoke run."""
    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("Choose --smoke to run the offline reference")
    print(json.dumps(smoke(args.seed), sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()