# examples/qwen2vl.py
"""Offline reduced Qwen2-VL packing and rotary mechanics.

Conventions: Transformers v4.45.2, commit
53fad641cfdb5105e2470bcf3ef17ea8e25cc300, matching model and image processor.
Synthetic float pixels bypass RGB normalization; bicubic interpolation is a
PyTorch stand-in for the processor's PIL resize. No tokenizer, ViT attention,
pretrained weights, LLM, generation or native performance is reproduced.
Odd video clips duplicate the last frame as an authored preparation extension:
the pinned processor itself requires paired frames (or a single still image).
"""

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 = "53fad641cfdb5105e2470bcf3ef17ea8e25cc300"
VISION_START, VISION_END, IMAGE, VIDEO, PAD = 90, 91, 92, 93, 0


# region: smart_resize
def smart_resize(height, width, patch_size=14, merge_size=2,
                 min_pixels=56 * 56, max_pixels=14 * 14 * 4 * 1280):
    """Return the pinned processor's bounded, factor-aligned spatial size."""
    values = (height, width, patch_size, merge_size, min_pixels, max_pixels)
    if any(type(value) is not int or value <= 0 for value in values):
        raise ValueError("Sizes and bounds must be positive integers.")
    factor = patch_size * merge_size
    if min_pixels > max_pixels or min_pixels < factor ** 2:
        raise ValueError("Bounds must be ordered and allow a factor cell.")
    if min(height, width) < factor:
        raise ValueError("Each input dimension must be at least the factor.")
    if max(height, width) / min(height, width) > 200:
        raise ValueError("Absolute aspect ratio must not exceed 200.")
    h_bar = round(height / factor) * factor
    w_bar = round(width / factor) * factor
    if h_bar * w_bar > max_pixels:
        beta = math.sqrt(height * width / max_pixels)
        h_bar = math.floor(height / beta / factor) * factor
        w_bar = math.floor(width / beta / factor) * factor
    elif h_bar * w_bar < min_pixels:
        beta = math.sqrt(min_pixels / (height * width))
        h_bar = math.ceil(height * beta / factor) * factor
        w_bar = math.ceil(width * beta / factor) * factor
    # Extreme custom budgets can make the upstream rounding infeasible.
    if min(h_bar, w_bar) == 0 or not min_pixels <= h_bar * w_bar <= max_pixels:
        raise ValueError("Rounded dimensions cannot satisfy these bounds.")
    return h_bar, w_bar
# endregion: smart_resize


# region: tubelets
def prepare_tubelets(frames, size):
    """Resize float [F,C,H,W] pixels and pair frames, duplicating if odd."""
    if (frames.ndim != 4 or min(frames.shape) <= 0
            or not frames.is_floating_point()):
        raise ValueError("Expected nonempty floating [F,C,H,W] pixels.")
    if not torch.isfinite(frames).all():
        raise ValueError("Pixels must be finite.")
    if len(size) != 2 or any(type(x) is not int or x <= 0 for x in size):
        raise ValueError("Resize dimensions must be positive integers.")
    resized = F.interpolate(
        frames, size=size, mode="bicubic", align_corners=False
    )
    if resized.shape[0] % 2:
        # F=1 is the paper's still-image duplication. For odd videos this
        # is explicit authored preparation, not v4.45.2 processor behavior.
        resized = torch.cat((resized, resized[-1:]), dim=0)
    return resized
# endregion: tubelets


# region: patch_order
def flatten_patches(frames, patch_size=14, merge_size=2):
    """Match processor transpose: t, group-h/w, inner-h/w, C, T, pixels."""
    if frames.ndim != 4 or min(frames.shape) <= 0:
        raise ValueError("Expected nonempty [F,C,H,W].")
    if any(type(x) is not int or x <= 0 for x in (patch_size, merge_size)):
        raise ValueError("Patch and merge sizes must be positive integers.")
    count, channels, height, width = frames.shape
    factor = patch_size * merge_size
    if count % 2 or height % factor or width % factor:
        raise ValueError("Frames must pair; spatial sizes must divide factor.")
    grid_t = count // 2
    grid_h, grid_w = height // patch_size, width // patch_size
    patches = frames.reshape(
        grid_t, 2, channels, grid_h // merge_size, merge_size, patch_size,
        grid_w // merge_size, merge_size, patch_size
    ).permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
    patches = patches.reshape(
        grid_t * grid_h * grid_w, channels * 2 * patch_size ** 2
    )
    return patches, (grid_t, grid_h, grid_w)


class PatchEmbed(nn.Module):
    """Pinned depth-two Conv3d over each processor-ordered patch row."""

    def __init__(self, channels=3, patch_size=2, width=6):
        super().__init__()
        self.channels, self.patch_size = channels, patch_size
        kernel = (2, patch_size, patch_size)
        self.proj = nn.Conv3d(
            channels, width, kernel, stride=kernel, bias=False
        )

    def forward(self, patches):
        pixels = patches.reshape(
            -1, self.channels, 2, self.patch_size, self.patch_size
        )
        return self.proj(pixels).flatten(1)
# endregion: patch_order


# region: merge
class PatchMerger(nn.Module):
    """Normalize each patch, concatenate 2x2 neighbors, then apply an MLP."""

    def __init__(self, context_width=6, output_width=12):
        super().__init__()
        self.context_width = context_width
        self.ln_q = nn.LayerNorm(context_width, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(4 * context_width, 4 * context_width), nn.GELU(),
            nn.Linear(4 * context_width, output_width)
        )

    def forward(self, features):
        if (features.ndim != 2 or features.shape[1] != self.context_width
                or features.shape[0] == 0 or features.shape[0] % 4):
            raise ValueError("Expected four adjacent patch rows per token.")
        neighbors = self.ln_q(features).reshape(-1, 4 * self.context_width)
        return self.mlp(neighbors)
# endregion: merge


def validate_grid(grid):
    """Validate this reference's fixed 2x2 spatial merge contract."""
    if (len(grid) != 3 or any(type(x) is not int or x <= 0 for x in grid)
            or grid[1] % 2 or grid[2] % 2):
        raise ValueError("Grid needs positive t and even h/w patch counts.")
    return grid[0] * (grid[1] // 2) * (grid[2] // 2)


# region: position_ids
def position_ids(input_ids, attention_mask, grids):
    """Assign [3,B,L] IDs; grids follow visual order, ignoring padding."""
    if input_ids.ndim != 2 or attention_mask.shape != input_ids.shape:
        raise ValueError("Token IDs and mask must have the same [B,L] shape.")
    if input_ids.dtype != torch.long or len(grids) != input_ids.shape[0]:
        raise ValueError("Use integer tokens and one grid list per sample.")
    if not ((attention_mask == 0) | (attention_mask == 1)).all():
        raise ValueError("Attention mask must be binary.")
    positions = torch.ones(3, *input_ids.shape, dtype=torch.long)
    deltas = []
    pure_text = not any(grids)
    for batch, row in enumerate(input_ids):
        valid = attention_mask[batch].bool()
        tokens = row[valid].tolist()
        if not tokens:
            raise ValueError("Each sample needs at least one valid token.")
        chunks, cursor, next_id, visual_index = [], 0, 0, 0
        while cursor < len(tokens):
            token = tokens[cursor]
            if token == VISION_START:
                chunks.append(torch.full((3, 1), next_id))
                next_id += 1
                if visual_index >= len(grids[batch]):
                    raise ValueError("Missing visual grid.")
                kind, grid = grids[batch][visual_index]
                if kind not in ("image", "video"):
                    raise ValueError("Unknown modality.")
                if kind == "image" and grid[0] != 1:
                    raise ValueError("A still image has one temporal tubelet.")
                count = validate_grid(grid)
                marker = IMAGE if kind == "image" else VIDEO
                end = cursor + 1 + count
                if (tokens[cursor + 1:end] != [marker] * count
                        or end >= len(tokens) or tokens[end] != VISION_END):
                    raise ValueError("Visual count or sentinels disagree.")
                t, h, w = grid[0], grid[1] // 2, grid[2] // 2
                axes = torch.meshgrid(
                    torch.arange(t), torch.arange(h), torch.arange(w),
                    indexing="ij"
                )
                visual = torch.stack([axis.flatten() for axis in axes])
                visual = visual + next_id
                chunks.append(visual)
                next_id = int(visual.max()) + 1
                # The end sentinel is a text position, handled next.
                cursor = end
                visual_index += 1
            else:
                if token in (IMAGE, VIDEO, PAD):
                    raise ValueError("Unexpected visual or padding token.")
                unmatched = (
                    cursor == 0 or tokens[cursor - 1] not in (IMAGE, VIDEO)
                )
                if token == VISION_END and unmatched:
                    raise ValueError("Unmatched end sentinel.")
                chunks.append(torch.full((3, 1), next_id))
                next_id += 1
                cursor += 1
        if visual_index != len(grids[batch]):
            raise ValueError("Unused visual grids.")
        positions[:, batch, valid] = torch.cat(chunks, dim=1)
        # The pure-text release branch includes padding IDs in its maximum.
        # Mixed visual streams retain the valid-position maximum convention.
        maximum = int(positions[:, batch].max()) if pure_text else next_id - 1
        deltas.append(maximum + 1 - input_ids.shape[1])
    return positions, torch.tensor(deltas).unsqueeze(1)
# endregion: position_ids


# region: mrope
def apply_mrope(query, key, positions, sections=(1, 2, 3), base=10000.0):
    """Select t/h/w channel rotations in each half; rotate [B,H,L,D] q/k."""
    if (query.ndim != 4 or key.ndim != 4
            or query.shape[0] != key.shape[0]
            or query.shape[2:] != key.shape[2:]
            or positions.shape != (3, query.shape[0], query.shape[2])):
        raise ValueError("Expected compatible q/k and [3,B,L] positions.")
    width = query.shape[-1]
    invalid_sections = (
        len(sections) != 3
        or any(type(x) is not int or x <= 0 for x in sections)
    )
    if invalid_sections or 2 * sum(sections) != width or base <= 1:
        raise ValueError("Three sections must sum to half the rotary width.")
    frequency = base ** (
        -torch.arange(0, width, 2, dtype=query.dtype) / width
    )
    angles = positions.to(query.dtype).unsqueeze(-1) * frequency
    angles = torch.cat((angles, angles), dim=-1)
    # Python list repetition gives [t,h,w,t,h,w], not doubled section widths.
    sizes = list(sections) * 2
    cosine = torch.cat(
        [block[i % 3]
         for i, block in enumerate(angles.cos().split(sizes, -1))],
        dim=-1
    ).unsqueeze(1)
    sine = torch.cat(
        [block[i % 3]
         for i, block in enumerate(angles.sin().split(sizes, -1))],
        dim=-1
    ).unsqueeze(1)

    def rotate_half(value):
        first, second = value.chunk(2, dim=-1)
        return torch.cat((-second, first), dim=-1)

    return (query * cosine + rotate_half(query) * sine,
            key * cosine + rotate_half(key) * sine)
# endregion: mrope


# region: packing
def pack_stream(parts, embed, merger, padding=0, pad_left=False):
    """Pack synthetic text IDs and actual reduced image/video tensors."""
    if type(padding) is not int or padding < 0:
        raise ValueError("Padding must be a nonnegative integer.")
    ids, features, grids, records = [], [], [], []
    for kind, payload in parts:
        if kind == "text":
            if not payload or any(type(x) is not int or not 0 < x < 90
                                  for x in payload):
                raise ValueError("Synthetic text IDs must be in [1,89].")
            ids.extend(payload)
            continue
        if kind not in ("image", "video"):
            raise ValueError("Unknown modality.")
        if (not isinstance(payload, torch.Tensor) or payload.ndim != 4
                or payload.shape[1] != embed.channels):
            raise ValueError("Visual tensors must match embedding channels.")
        if kind == "image" and payload.shape[0] != 1:
            raise ValueError("A still image must supply one frame.")
        height, width = payload.shape[-2:]
        size = smart_resize(height, width, 2, 2, 16, 256)
        paired = prepare_tubelets(payload, size)
        patches, grid = flatten_patches(paired, 2, 2)
        merged = merger(embed(patches))
        count = validate_grid(grid)
        if patches.shape[0] != 4 * count or merged.shape[0] != count:
            raise ValueError("Patch, feature and placeholder counts differ.")
        marker = IMAGE if kind == "image" else VIDEO
        ids.extend([VISION_START] + [marker] * count + [VISION_END])
        features.append(merged)
        grids.append((kind, grid))
        records.append({
            "kind": kind, "input_shape": list(payload.shape),
            "paired_shape": list(paired.shape), "grid_thw": list(grid),
            "patch_shape": list(patches.shape),
            "merged_shape": list(merged.shape), "sentinels": 2,
        })
    padded = [PAD] * padding + ids if pad_left else ids + [PAD] * padding
    mask = ([0] * padding + [1] * len(ids) if pad_left
            else [1] * len(ids) + [0] * padding)
    tokens = torch.tensor([padded], dtype=torch.long)
    attention = torch.tensor([mask], dtype=torch.long)
    positions, deltas = position_ids(tokens, attention, [grids])
    return {
        "input_ids": tokens, "attention_mask": attention,
        "positions": positions, "deltas": deltas,
        "features": features, "records": records,
    }
# endregion: packing


def smoke(seed=7):
    """Execute seeded CPU examples and emit finite measured observations."""
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    embed, merger = PatchEmbed(), PatchMerger()
    image = torch.randn(1, 3, 9, 17)
    image_two = torch.randn(1, 3, 8, 8)
    video = torch.randn(3, 3, 8, 12)
    cases = {
        "rectangular_image": [("text", [1, 2]), ("image", image)],
        "two_images": [("image", image), ("text", [3]),
                       ("image", image_two)],
        "mixed_video": [("text", [4, 5]), ("image", image),
                        ("text", [6]), ("video", video), ("text", [7])],
    }
    parameters = list(embed.parameters()) + list(merger.parameters())
    before = [parameter.detach().clone() for parameter in parameters]
    results, loss_terms, rotary_leaves = {}, [], {}
    norm_error = 0.0
    for name, parts in cases.items():
        packed = pack_stream(parts, embed, merger)
        left = pack_stream(parts, embed, merger, padding=3, pad_left=True)
        length = packed["input_ids"].shape[1]
        query = torch.randn(1, 2, length, 12, requires_grad=True)
        key = torch.randn(1, 1, length, 12, requires_grad=True)
        rotated_q, rotated_k = apply_mrope(query, key, packed["positions"])
        rotary_leaves[name] = (query, key)
        error = (rotated_q.norm(dim=-1) - query.norm(dim=-1)).abs().max()
        norm_error = max(norm_error, float(error.detach()))
        loss_terms.extend([feature.square().mean()
                           for feature in packed["features"]])
        loss_terms.append(
            rotated_q.square().mean() + rotated_k.square().mean()
        )
        stable = torch.equal(packed["positions"], left["positions"][:, :, 3:])
        results[name] = {
            "visuals": packed["records"], "sequence_length": length,
            "positions": packed["positions"].tolist(),
            "padding_positions_stable": stable,
            "position_delta": packed["deltas"].tolist(),
            "query_shape": list(query.shape), "key_shape": list(key.shape),
        }
    diagnostic = torch.stack(loss_terms).mean()
    diagnostic.backward()
    for name, (query, key) in rotary_leaves.items():
        results[name]["rotary_gradient_norms"] = {
            "query": float(query.grad.detach().norm()),
            "key": float(key.grad.detach().norm()),
        }
    parameter_change = max(float((p.detach() - old).abs().max())
                           for p, old in zip(parameters, before))
    named_parameters = list(embed.named_parameters()) + [
        ("merger." + name, parameter)
        for name, parameter in merger.named_parameters()
    ]
    gradients = {name: float(parameter.grad.detach().norm())
                 for name, parameter in named_parameters}
    report = {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "revision": REVISION, "seed": seed, "device": "cpu",
        "reductions": {
            "patch_size": 2, "temporal_depth": 2, "merge_size": 2,
            "vision_width": 6, "rotary_width": 12, "sections": [1, 2, 3],
            "synthetic_pixels": True, "pretrained_weights": False,
            "vision_attention": False, "optimizer_steps": 0,
        },
        "cases": results, "diagnostic_loss": float(diagnostic.detach()),
        "rotation_norm_max_error": norm_error, "gradient_norms": gradients,
        "parameters_unchanged_after_backward": parameter_change == 0.0,
        "parameter_max_change": parameter_change,
        "all_gradients_finite": all(
            torch.isfinite(p.grad).all().item()
            for p in parameters + [leaf for pair in rotary_leaves.values()
                                   for leaf in pair]
        ),
        "limitations": ["No pretrained semantics or native benchmark",
                        "PyTorch resize stand-in; no RGB normalization",
                        "Odd video padding is authored preparation"],
    }
    json.dumps(report, allow_nan=False)
    return report


def main():
    """Run the explicit offline smoke entry point."""
    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("Use --smoke to execute the reduced CPU reference.")
    print(json.dumps(smoke(args.seed), allow_nan=False, sort_keys=True))


if __name__ == "__main__":
    main()