Module 16. Multimodal LLMs: Vision, Audio, and Video

import importlib.util
import subprocess
import sys

for package in ("numpy", "torch", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])

import hashlib
import json
import math
from dataclasses import dataclass

import numpy as np

np.set_printoptions(precision=5, suppress=True)
SEED = 16
rng = np.random.default_rng(SEED)
print('Environment ready.')

def _positive_finite(name: str, value: float) -> float:
    """Validate and return one strictly positive finite floating-point value."""

    value = float(value)
    if not math.isfinite(value) or value <= 0.0:
        raise ValueError(f"{name} must be finite and positive")
    return value

REFERENCE_JSON = '{"infonce":{"batch_size":4,"negative_similarity":0.2,"positive_similarity":0.8,"tau_0_1":{"image_to_text":0.00740874388428169,"symmetric":0.00740874388428169,"text_to_image":0.00740874388428169},"tau_0_1_probability":0.9926186332070156,"tau_0_5":{"image_to_text":0.6437377084206329,"symmetric":0.6437377084206329,"text_to_image":0.6437377084206329},"uniform_floor":1.3862943611198906},"kv":{"device_budget_gib":80.0,"per_image":{"llava336":{"available_gib":65.0,"capacity_upper_bound":924,"kv_bytes":75497472,"kv_mib":72.0,"visual_tokens":576},"next672":{"available_gib":65.0,"capacity_upper_bound":184,"kv_bytes":377487360,"kv_mib":360.0,"visual_tokens":2880},"qwen1008_merged":{"available_gib":65.0,"capacity_upper_bound":410,"kv_bytes":169869312,"kv_mib":162.0,"visual_tokens":1296}},"spec":{"bytes_per_element":2,"bytes_per_token":131072,"head_dim":128,"kib_per_token":128.0,"n_kv_heads":8,"n_layers":32},"weights_gib":15.0},"patches":{"llava336":576,"next672":2880,"qwen1008":5184,"qwen1008_merged":1296,"rectangular_56x28":8,"rectangular_56x28_merged":2},"projector":{"hidden_dim":4096,"llm_dim":4096,"parameter_count":20979712,"vision_dim":1024},"schema_version":1,"training_floor":{"flops_multiplier":6.0,"mfu":0.4,"n_parameters":7000000000.0,"peak_flops_per_second":990000000000000.0,"stage1_hours":9.863636363636365,"stage1_tokens":334800000,"stage2_hours":19.59175084175084,"stage2_tokens":665000000,"total_hours":29.455387205387204},"video":{"source_frames":7200,"with_merger_2x2":{"full_prompt_kv_gib":31.25,"full_prompt_tokens":256000,"max_frames":1750,"sample_interval_seconds":4.114285714285714,"selected_frames":1750,"total_frames":7200,"used_visual_tokens":252000,"visual_budget_tokens":252000,"visual_kv_gib":30.76171875},"without_merger":{"full_prompt_kv_gib":31.21484375,"full_prompt_tokens":255712,"max_frames":437,"sample_interval_seconds":16.475972540045767,"selected_frames":437,"total_frames":7200,"used_visual_tokens":251712,"visual_budget_tokens":252000,"visual_kv_gib":30.7265625}}}'
EXPECTED_SHA256 = 'b22068921a7c496a72b525b7015d5c7450e3b0da8fb0b2d482de6e97a899449c'
actual_sha256 = hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest()
assert actual_sha256 == EXPECTED_SHA256
REFERENCE = json.loads(REFERENCE_JSON)
assert REFERENCE["schema_version"] == 1
print('The Module-16 payload was unpacked and verified', actual_sha256[:12])

def _positive_int(name: str, value: int) -> int:
    """Validate and return one strictly positive integer."""

    if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)



@dataclass(frozen=True)
class PatchLayout:
    """Spatial patch layout for one image before or after token merging."""

    image_height: int
    image_width: int
    patch_size: int
    merge_factor: int
    grid_height: int
    grid_width: int
    token_count: int



def patch_layout(
    image_height: int,
    image_width: int,
    *,
    patch_size: int = 14,
    merge_factor: int = 1,
) -> PatchLayout:
    """Return the exact patch grid and token count for a rectangular image.

    ``merge_factor=m`` groups each spatial ``m × m`` block into one visual token.
    Both the image and the resulting patch grid must be exactly divisible; silent
    truncation would under-count tokens and is therefore rejected.
    """

    image_height = _positive_int("image_height", image_height)
    image_width = _positive_int("image_width", image_width)
    patch_size = _positive_int("patch_size", patch_size)
    merge_factor = _positive_int("merge_factor", merge_factor)
    if image_height % patch_size or image_width % patch_size:
        raise ValueError("image dimensions must be divisible by patch_size")
    grid_height = image_height // patch_size
    grid_width = image_width // patch_size
    if grid_height % merge_factor or grid_width % merge_factor:
        raise ValueError("the patch grid must be divisible by merge_factor on both axes")
    merged_h = grid_height // merge_factor
    merged_w = grid_width // merge_factor
    return PatchLayout(
        image_height=image_height,
        image_width=image_width,
        patch_size=patch_size,
        merge_factor=merge_factor,
        grid_height=grid_height,
        grid_width=grid_width,
        token_count=merged_h * merged_w,
    )



P = REFERENCE["patches"]
assert patch_layout(336, 336).token_count == P["llava336"] == 576
assert 5 * patch_layout(336, 336).token_count == P["next672"] == 2880
assert patch_layout(1008, 1008).token_count == P["qwen1008"] == 5184
assert patch_layout(1008, 1008, merge_factor=2).token_count == P["qwen1008_merged"] == 1296
assert patch_layout(56, 28).token_count == P["rectangular_56x28"] == 8
assert patch_layout(56, 28, merge_factor=2).token_count == P["rectangular_56x28_merged"] == 2
print('Patch layout validated for square and rectangular images')

def patchify_rgb(image: np.ndarray, patch_size: int) -> tuple[np.ndarray, tuple[int, int]]:
    """Convert one ``[H, W, C]`` image into flattened row-major patch vectors.

    Returns ``patches`` with shape ``[N_patches, patch_size² × C]`` and the unmerged
    spatial grid ``(grid_height, grid_width)``.  The function accepts any numeric dtype
    but rejects non-finite floating-point values.
    """

    array = np.asarray(image)
    if array.ndim != 3:
        raise ValueError("image must have shape [H, W, C]")
    height, width, channels = array.shape
    if channels <= 0:
        raise ValueError("image must contain at least one channel")
    layout = patch_layout(height, width, patch_size=patch_size)
    if not np.issubdtype(array.dtype, np.number):
        raise TypeError("image dtype must be numeric")
    if np.issubdtype(array.dtype, np.floating) and not np.all(np.isfinite(array)):
        raise ValueError("image contains NaN or infinity")
    patches = (
        array.reshape(layout.grid_height, patch_size, layout.grid_width, patch_size, channels)
        .transpose(0, 2, 1, 3, 4)
        .reshape(layout.token_count, patch_size * patch_size * channels)
    )
    return patches, (layout.grid_height, layout.grid_width)



def unpatchify_rgb(patches: np.ndarray, grid_shape: tuple[int, int], patch_size: int, channels: int) -> np.ndarray:
    """Invert patchify for a complete unmerged grid."""

    patches = np.asarray(patches)
    if len(grid_shape) != 2:
        raise ValueError("grid_shape must contain two dimensions")
    grid_height, grid_width = (_positive_int("grid dimension", v) for v in grid_shape)
    patch_size = _positive_int("patch_size", patch_size)
    channels = _positive_int("channels", channels)
    expected = (grid_height * grid_width, patch_size * patch_size * channels)
    if patches.shape != expected:
        raise ValueError(f"patches must have shape {expected}")
    return (
        patches.reshape(grid_height, grid_width, patch_size, patch_size, channels)
        .transpose(0, 2, 1, 3, 4)
        .reshape(grid_height * patch_size, grid_width * patch_size, channels)
    )


def merge_spatial_tokens(
    tokens: np.ndarray,
    grid_shape: tuple[int, int],
    *,
    merge_factor: int = 2,
) -> tuple[np.ndarray, tuple[int, int]]:
    """Group each spatial ``m × m`` block and concatenate its feature vectors.

    ``tokens`` may be ``[N, D]`` or batched ``[B, N, D]``.  The output keeps batch
    order and has shape ``[N/m², D·m²]`` or ``[B, N/m², D·m²]``.  This is the
    transparent pixel-unshuffle part of a merger; a learned projection can follow it.
    """

    array = np.asarray(tokens)
    if array.ndim not in {2, 3}:
        raise ValueError("tokens must have shape [N, D] or [B, N, D]")
    if len(grid_shape) != 2:
        raise ValueError("grid_shape must contain two integers")
    grid_height, grid_width = grid_shape
    grid_height = _positive_int("grid_height", grid_height)
    grid_width = _positive_int("grid_width", grid_width)
    merge_factor = _positive_int("merge_factor", merge_factor)
    n_tokens = array.shape[-2]
    feature_dim = array.shape[-1]
    if feature_dim == 0 or not np.issubdtype(array.dtype, np.number):
        raise ValueError("tokens must have nonempty numeric features")
    if n_tokens != grid_height * grid_width:
        raise ValueError("token count does not match grid_shape")
    if grid_height % merge_factor or grid_width % merge_factor:
        raise ValueError("grid_shape must be divisible by merge_factor")
    if np.issubdtype(array.dtype, np.floating) and not np.all(np.isfinite(array)):
        raise ValueError("tokens contain NaN or infinity")

    batched = array.ndim == 3
    work = array if batched else array[None, ...]
    batch = work.shape[0]
    merged_h = grid_height // merge_factor
    merged_w = grid_width // merge_factor
    merged = (
        work.reshape(batch, merged_h, merge_factor, merged_w, merge_factor, feature_dim)
        .transpose(0, 1, 3, 2, 4, 5)
        .reshape(batch, merged_h * merged_w, feature_dim * merge_factor * merge_factor)
    )
    return (merged if batched else merged[0]), (merged_h, merged_w)



image = np.arange(56 * 28 * 3, dtype=np.float32).reshape(56, 28, 3)
patches, grid = patchify_rgb(image, patch_size=14)
restored = unpatchify_rgb(patches, grid, patch_size=14, channels=3)
merged, merged_grid = merge_spatial_tokens(patches, grid, merge_factor=2)
assert patches.shape == (8, 14 * 14 * 3)
assert np.array_equal(restored, image)
# A round trip alone cannot establish the specified order: two inverse mistakes
# can cancel. Compare against an independently written coordinate fixture.
fixture = np.arange(16).reshape(4, 4, 1)
fixture_patches, _ = patchify_rgb(fixture, patch_size=2)
assert fixture_patches.tolist() == [[0,1,4,5], [2,3,6,7], [8,9,12,13], [10,11,14,15]]
fixture_merged, _ = merge_spatial_tokens(fixture_patches, (2,2), merge_factor=2)
assert fixture_merged.tolist() == [[0,1,4,5,2,3,6,7,8,9,12,13,10,11,14,15]]
assert merged.shape == (2, 4 * 14 * 14 * 3) and merged_grid == (2, 1)
print('Patchify round-trip and the 2×2 merger passed', "|", patches.shape, "->", merged.shape)

def gelu_exact(x: np.ndarray) -> np.ndarray:
    """Evaluate exact GELU in float64 without requiring SciPy."""

    array = np.asarray(x, dtype=np.float64)
    erf = np.vectorize(math.erf, otypes=[np.float64])
    return 0.5 * array * (1.0 + erf(array / math.sqrt(2.0)))



def count_projector_parameters(
    vision_dim: int,
    llm_dim: int,
    *,
    hidden_dim: int | None = None,
    bias: bool = True,
) -> int:
    """Count parameters of a two-layer MLP vision projector exactly."""

    vision_dim = _positive_int("vision_dim", vision_dim)
    llm_dim = _positive_int("llm_dim", llm_dim)
    hidden_dim = llm_dim if hidden_dim is None else _positive_int("hidden_dim", hidden_dim)
    total = vision_dim * hidden_dim + hidden_dim * llm_dim
    if bias:
        total += hidden_dim + llm_dim
    return total



@dataclass(frozen=True)
class ProjectorWeights:
    """Weights of a two-layer MLP mapping vision features into the LLM width."""

    w1: np.ndarray
    b1: np.ndarray
    w2: np.ndarray
    b2: np.ndarray

    def validate(self) -> tuple[int, int, int]:
        """Validate matrix compatibility and return ``(vision, hidden, llm)`` widths."""

        w1 = np.asarray(self.w1)
        b1 = np.asarray(self.b1)
        w2 = np.asarray(self.w2)
        b2 = np.asarray(self.b2)
        if w1.ndim != 2 or w2.ndim != 2 or b1.ndim != 1 or b2.ndim != 1:
            raise ValueError("projector weights must be rank-2 matrices and rank-1 biases")
        vision_dim, hidden_dim = w1.shape
        if b1.shape != (hidden_dim,) or w2.shape[0] != hidden_dim:
            raise ValueError("w1/b1/w2 hidden dimensions are incompatible")
        llm_dim = w2.shape[1]
        if b2.shape != (llm_dim,):
            raise ValueError("w2 and b2 output dimensions are incompatible")
        for name, value in {"w1": w1, "b1": b1, "w2": w2, "b2": b2}.items():
            if not np.issubdtype(value.dtype, np.number) or not np.all(np.isfinite(value)):
                raise ValueError(f"{name} must be finite and numeric")
        return vision_dim, hidden_dim, llm_dim



def init_toy_projector(
    vision_dim: int,
    hidden_dim: int,
    llm_dim: int,
    *,
    seed: int = 14,
) -> ProjectorWeights:
    """Create deterministic small weights for a CPU-safe teaching forward pass."""

    vision_dim = _positive_int("vision_dim", vision_dim)
    hidden_dim = _positive_int("hidden_dim", hidden_dim)
    llm_dim = _positive_int("llm_dim", llm_dim)
    rng = np.random.default_rng(seed)
    w1 = rng.normal(0.0, 1.0 / math.sqrt(vision_dim), size=(vision_dim, hidden_dim))
    w2 = rng.normal(0.0, 1.0 / math.sqrt(hidden_dim), size=(hidden_dim, llm_dim))
    return ProjectorWeights(
        w1=w1.astype(np.float32),
        b1=np.zeros(hidden_dim, dtype=np.float32),
        w2=w2.astype(np.float32),
        b2=np.zeros(llm_dim, dtype=np.float32),
    )



def project_vision(features: np.ndarray, weights: ProjectorWeights) -> np.ndarray:
    """Project ``[..., N_visual, D_vision]`` features into ``D_llm`` with GELU."""

    array = np.asarray(features)
    if array.ndim < 2:
        raise ValueError("features must have shape [..., N_visual, D_vision]")
    vision_dim, _, _ = weights.validate()
    if array.shape[-1] != vision_dim:
        raise ValueError("features last dimension does not match projector vision_dim")
    if not np.issubdtype(array.dtype, np.number) or not np.all(np.isfinite(array)):
        raise ValueError("features must be finite and numeric")
    work = array.astype(np.float32, copy=False)
    hidden = gelu_exact(work @ weights.w1 + weights.b1).astype(np.float32)
    output = hidden @ weights.w2 + weights.b2
    if not np.all(np.isfinite(output)):
        raise FloatingPointError("projector produced NaN or infinity")
    return output.astype(np.float32, copy=False)



full = REFERENCE["projector"]
assert count_projector_parameters(full["vision_dim"], full["llm_dim"], hidden_dim=full["hidden_dim"]) == full["parameter_count"] == 20_979_712
vision_features = merged.reshape(2, 8, -1).mean(axis=-1, dtype=np.float64).astype(np.float32)[None, ...]
weights = init_toy_projector(vision_dim=8, hidden_dim=12, llm_dim=16)
visual_embeddings = project_vision(vision_features, weights)
assert visual_embeddings.shape == (1, 2, 16)
print('The projector preserved batch/token axes and changed only feature width', "|", vision_features.shape, "->", visual_embeddings.shape, "| full params:", full["parameter_count"])

@dataclass(frozen=True)
class MultimodalSequence:
    """Decoder input after replacing one image placeholder by continuous features."""

    embeddings: np.ndarray
    token_labels: np.ndarray
    modality_ids: np.ndarray
    position_ids: np.ndarray
    attention_mask: np.ndarray



def assemble_multimodal_sequence(
    token_ids: np.ndarray,
    token_embeddings: np.ndarray,
    *,
    image_placeholder_id: int,
    image_embeddings: np.ndarray,
    image_start_embedding: np.ndarray,
    image_end_embedding: np.ndarray,
    ignore_index: int = -100,
) -> MultimodalSequence:
    """Replace exactly one placeholder token by start/vision/end embeddings.

    Text positions retain their token IDs as labels.  Continuous visual tokens and
    image-boundary embeddings receive ``ignore_index``.  The returned causal mask is
    ``[sequence, sequence]`` and positions are contiguous across modalities.
    """

    token_ids = np.asarray(token_ids)
    token_embeddings = np.asarray(token_embeddings)
    image_embeddings = np.asarray(image_embeddings)
    image_start_embedding = np.asarray(image_start_embedding)
    image_end_embedding = np.asarray(image_end_embedding)
    if token_ids.ndim != 1 or not np.issubdtype(token_ids.dtype, np.integer):
        raise ValueError("token_ids must be a rank-1 integer array")
    if token_embeddings.ndim != 2 or token_embeddings.shape[0] != token_ids.size:
        raise ValueError("token_embeddings must have shape [len(token_ids), D]")
    if image_embeddings.ndim != 2 or image_embeddings.shape[0] == 0:
        raise ValueError("image_embeddings must have shape [N_visual>0, D]")
    width = token_embeddings.shape[1]
    if image_embeddings.shape[1] != width:
        raise ValueError("text and image embeddings must share the LLM width")
    if image_start_embedding.shape != (width,) or image_end_embedding.shape != (width,):
        raise ValueError("image boundary embeddings must have shape [D]")
    for name, value in {
        "token_embeddings": token_embeddings,
        "image_embeddings": image_embeddings,
        "image_start_embedding": image_start_embedding,
        "image_end_embedding": image_end_embedding,
    }.items():
        if not np.all(np.isfinite(value)):
            raise ValueError(f"{name} contains NaN or infinity")
    if (isinstance(image_placeholder_id, bool)
            or not isinstance(image_placeholder_id, (int, np.integer))):
        raise ValueError("image_placeholder_id must be an integer")
    if np.any(token_ids < 0) or width == 0:
        raise ValueError("token IDs must be nonnegative and embedding width positive")
    placeholder_positions = np.flatnonzero(token_ids == image_placeholder_id)
    if placeholder_positions.size != 1:
        raise ValueError("token_ids must contain exactly one image placeholder")
    placeholder = int(placeholder_positions[0])

    before_embeddings = token_embeddings[:placeholder]
    after_embeddings = token_embeddings[placeholder + 1 :]
    visual_block = np.concatenate(
        [image_start_embedding[None, :], image_embeddings, image_end_embedding[None, :]],
        axis=0,
    )
    embeddings = np.concatenate([before_embeddings, visual_block, after_embeddings], axis=0)

    labels = np.concatenate(
        [
            token_ids[:placeholder],
            np.full(visual_block.shape[0], ignore_index, dtype=np.int64),
            token_ids[placeholder + 1 :],
        ]
    ).astype(np.int64)
    modality_ids = np.concatenate(
        [
            np.zeros(before_embeddings.shape[0], dtype=np.int8),
            np.full(visual_block.shape[0], 1, dtype=np.int8),
            np.zeros(after_embeddings.shape[0], dtype=np.int8),
        ]
    )
    sequence = embeddings.shape[0]
    position_ids = np.arange(sequence, dtype=np.int64)
    attention_mask = np.tril(np.ones((sequence, sequence), dtype=bool))
    return MultimodalSequence(
        embeddings=embeddings.astype(np.float32, copy=False),
        token_labels=labels,
        modality_ids=modality_ids,
        position_ids=position_ids,
        attention_mask=attention_mask,
    )



IMAGE_ID = 999
token_ids = np.array([101, 102, IMAGE_ID, 103, 104], dtype=np.int64)
text_embeddings = rng.normal(size=(token_ids.size, 16)).astype(np.float32)
sequence = assemble_multimodal_sequence(
    token_ids,
    text_embeddings,
    image_placeholder_id=IMAGE_ID,
    image_embeddings=visual_embeddings[0],
    image_start_embedding=np.zeros(16, dtype=np.float32),
    image_end_embedding=np.ones(16, dtype=np.float32),
)
assert sequence.embeddings.shape == (8, 16)
assert sequence.token_labels.tolist() == [101, 102, -100, -100, -100, -100, 103, 104]
assert sequence.modality_ids.tolist() == [0, 0, 1, 1, 1, 1, 0, 0]
assert np.array_equal(sequence.position_ids, np.arange(8))
assert np.array_equal(sequence.attention_mask, np.tril(np.ones((8, 8), dtype=bool)))
assert not sequence.attention_mask[2, 3] and sequence.attention_mask[7, 0]
print('The placeholder became a visual block; labels, modality IDs, and the causal mask agree', "| sequence:", sequence.embeddings.shape)

def _logsumexp(values: np.ndarray, axis: int) -> np.ndarray:
    """Stable log-sum-exp for finite arrays."""

    maximum = np.max(values, axis=axis, keepdims=True)
    return np.squeeze(maximum + np.log(np.sum(np.exp(values - maximum), axis=axis, keepdims=True)), axis=axis)



def l2_normalize_rows(embeddings: np.ndarray, *, eps: float = 1e-12) -> np.ndarray:
    """L2-normalize a finite ``[batch, width]`` matrix and reject zero rows."""

    array = np.asarray(embeddings, dtype=np.float64)
    if array.ndim != 2 or array.shape[0] < 2 or array.shape[1] == 0:
        raise ValueError("embeddings must have shape [batch>=2, width>0]")
    if not np.all(np.isfinite(array)):
        raise ValueError("embeddings contain NaN or infinity")
    eps = _positive_finite("eps", eps)
    scale = np.max(np.abs(array), axis=1, keepdims=True)
    if np.any(scale == 0):
        raise ValueError("embedding rows must have non-zero norm")
    scaled = array / scale
    norms = np.linalg.norm(scaled, axis=1, keepdims=True)
    if np.any(scale <= eps / norms):
        raise ValueError("embedding row norm must exceed eps")
    return scaled / norms



def symmetric_infonce(similarity: np.ndarray, temperature: float) -> dict[str, float]:
    """Return stable image→text, text→image, and symmetric InfoNCE losses.

    ``similarity`` must be a finite square matrix whose diagonal contains the aligned
    image-text pairs.  Log-sum-exp arithmetic avoids overflow at small temperatures.
    """

    similarity = np.asarray(similarity, dtype=np.float64)
    if similarity.ndim != 2 or similarity.shape[0] != similarity.shape[1] or similarity.shape[0] < 2:
        raise ValueError("similarity must be a square [batch>=2, batch] matrix")
    if not np.all(np.isfinite(similarity)):
        raise ValueError("similarity contains NaN or infinity")
    temperature = _positive_finite("temperature", temperature)
    # Center BEFORE dividing; equal huge finite scores need not overflow.
    def directional(values):
        with np.errstate(over='ignore', invalid='ignore'):
            logits = (values - np.max(values, axis=1, keepdims=True)) / temperature
            loss = _logsumexp(logits, axis=1) - np.diag(logits)
        if not np.all(np.isfinite(loss)):
            raise FloatingPointError("InfoNCE exceeds the float64 range")
        return float(np.mean(loss))
    image_to_text = directional(similarity)
    text_to_image = directional(similarity.T)
    return {
        "image_to_text": image_to_text,
        "text_to_image": text_to_image,
        "symmetric": 0.5 * (image_to_text + text_to_image),
    }



I = REFERENCE["infonce"]
similarity = np.full((I["batch_size"], I["batch_size"]), I["negative_similarity"], dtype=np.float64)
np.fill_diagonal(similarity, I["positive_similarity"])
loss01 = symmetric_infonce(similarity, 0.1)
loss05 = symmetric_infonce(similarity, 0.5)
for key in ("image_to_text", "text_to_image", "symmetric"):
    assert math.isclose(loss01[key], I["tau_0_1"][key], rel_tol=0, abs_tol=1e-12)
    assert math.isclose(loss05[key], I["tau_0_5"][key], rel_tol=0, abs_tol=1e-12)
assert loss01["symmetric"] < loss05["symmetric"] < I["uniform_floor"]
normalized = l2_normalize_rows(rng.normal(size=(4, 8)))
assert np.allclose(np.linalg.norm(normalized, axis=1), 1.0)
assert np.max(np.abs(normalized @ normalized.T)) <= 1.0 + 1e-12
print('Image→text, text→image, and symmetric InfoNCE were reproduced stably', "| tau=.1:", loss01["symmetric"], "| tau=.5:", loss05["symmetric"])

@dataclass(frozen=True)
class KVCacheSpec:
    """K+V geometry for one decoder token across all layers."""

    n_layers: int
    n_kv_heads: int
    head_dim: int
    bytes_per_element: int

    def bytes_per_token(self) -> int:
        """Return exact K+V bytes per token across all layers."""

        n_layers = _positive_int("n_layers", self.n_layers)
        n_kv_heads = _positive_int("n_kv_heads", self.n_kv_heads)
        head_dim = _positive_int("head_dim", self.head_dim)
        bytes_per_element = _positive_int("bytes_per_element", self.bytes_per_element)
        return 2 * n_layers * n_kv_heads * head_dim * bytes_per_element



def image_kv_metrics(
    visual_tokens: int,
    spec: KVCacheSpec,
    *,
    device_budget_gib: float,
    weights_gib: float,
    runtime_reserve_gib: float = 0.0,
) -> dict[str, float | int]:
    """Compute one image's KV memory and a KV-only capacity upper bound.

    The capacity ignores activations, allocator fragmentation, temporary buffers, and
    scheduling overhead.  It is therefore explicitly an upper bound, not a deployment
    promise.
    """

    visual_tokens = _positive_int("visual_tokens", visual_tokens)
    device_budget_gib = _positive_finite("device_budget_gib", device_budget_gib)
    weights_gib = _positive_finite("weights_gib", weights_gib)
    runtime_reserve_gib = float(runtime_reserve_gib)
    if not math.isfinite(runtime_reserve_gib) or runtime_reserve_gib < 0.0:
        raise ValueError("runtime_reserve_gib must be finite and non-negative")
    available_gib = device_budget_gib - weights_gib - runtime_reserve_gib
    if available_gib <= 0.0:
        raise ValueError("weights plus reserve leave no positive KV budget")
    kv_bytes = visual_tokens * spec.bytes_per_token()
    kv_mib = kv_bytes / 2**20
    capacity_upper_bound = int((available_gib * 2**30) // kv_bytes)
    return {
        "visual_tokens": visual_tokens,
        "kv_bytes": kv_bytes,
        "kv_mib": kv_mib,
        "available_gib": available_gib,
        "capacity_upper_bound": capacity_upper_bound,
    }



K = REFERENCE["kv"]
spec = KVCacheSpec(**{key: K["spec"][key] for key in ("n_layers", "n_kv_heads", "head_dim", "bytes_per_element")})
assert spec.bytes_per_token() == K["spec"]["bytes_per_token"] == 128 * 2**10
rows = []
for name in ("llava336", "next672", "qwen1008_merged"):
    expected = K["per_image"][name]
    actual = image_kv_metrics(expected["visual_tokens"], spec, device_budget_gib=K["device_budget_gib"], weights_gib=K["weights_gib"])
    assert math.isclose(actual["kv_mib"], expected["kv_mib"], abs_tol=1e-12)
    assert actual["capacity_upper_bound"] == expected["capacity_upper_bound"]
    reserved = image_kv_metrics(expected["visual_tokens"], spec, device_budget_gib=80.0, weights_gib=15.0, runtime_reserve_gib=8.0)
    assert reserved["capacity_upper_bound"] <= actual["capacity_upper_bound"]
    rows.append((name, expected["visual_tokens"], actual["kv_mib"], actual["capacity_upper_bound"], reserved["capacity_upper_bound"]))
print('KV bytes/token, MiB/image, and capacity upper bounds passed')
for row in rows:
    print(f"{row[0]:>18}: {row[1]:4d} tok | {row[2]:6.1f} MiB | KV-only cap {row[3]:3d} | with 8 GiB reserve {row[4]:3d}")

def estimate_dense_training_hours(
    n_parameters: float,
    training_tokens: float,
    *,
    peak_flops_per_second: float,
    mfu: float,
    flops_multiplier: float = 6.0,
) -> float:
    """Return a transparent dense-equivalent training scenario in hours."""

    n_parameters = _positive_finite("n_parameters", n_parameters)
    training_tokens = _positive_finite("training_tokens", training_tokens)
    peak_flops_per_second = _positive_finite("peak_flops_per_second", peak_flops_per_second)
    mfu = float(mfu)
    if not math.isfinite(mfu) or not 0.0 < mfu <= 1.0:
        raise ValueError("mfu must lie in (0, 1]")
    flops_multiplier = _positive_finite("flops_multiplier", flops_multiplier)
    return flops_multiplier * n_parameters * training_tokens / (peak_flops_per_second * mfu) / 3600.0



C = REFERENCE["training_floor"]
kwargs = dict(
    n_parameters=C["n_parameters"],
    peak_flops_per_second=C["peak_flops_per_second"],
    mfu=C["mfu"],
    flops_multiplier=C["flops_multiplier"],
)
s1 = estimate_dense_training_hours(training_tokens=C["stage1_tokens"], **kwargs)
s2 = estimate_dense_training_hours(training_tokens=C["stage2_tokens"], **kwargs)
assert math.isclose(s1, C["stage1_hours"], abs_tol=1e-12)
assert math.isclose(s2, C["stage2_hours"], abs_tol=1e-12)
assert math.isclose(s1 + s2, C["total_hours"], abs_tol=1e-12)
print('The exact projector count and dense-equivalent scenario passed', "| projector params:", REFERENCE["projector"]["parameter_count"], "| scenario hours:", round(s1 + s2, 2))

def uniform_sample_indices(total_frames: int, max_frames: int) -> np.ndarray:
    """Choose deterministic, approximately uniform frame indices including endpoints."""

    total_frames = _positive_int("total_frames", total_frames)
    max_frames = _positive_int("max_frames", max_frames)
    if max_frames >= total_frames:
        return np.arange(total_frames, dtype=np.int64)
    if max_frames == 1:
        return np.array([0], dtype=np.int64)
    indices = np.rint(np.linspace(0, total_frames - 1, max_frames)).astype(np.int64)
    if np.unique(indices).size != max_frames:
        raise AssertionError("uniform sampler produced duplicate indices")
    if indices[0] != 0 or indices[-1] != total_frames - 1:
        raise AssertionError("uniform sampler must include first and last frame")
    return indices



def video_budget(
    *,
    duration_seconds: float,
    source_fps: float,
    context_tokens: int,
    text_reserve_tokens: int,
    visual_tokens_per_frame: int,
    kv_spec: KVCacheSpec,
) -> dict[str, float | int]:
    """Compute frame capacity, sampling interval, and KV for one video prompt."""

    duration_seconds = _positive_finite("duration_seconds", duration_seconds)
    source_fps = _positive_finite("source_fps", source_fps)
    context_tokens = _positive_int("context_tokens", context_tokens)
    if not isinstance(text_reserve_tokens, (int, np.integer)) or isinstance(text_reserve_tokens, bool):
        raise ValueError("text_reserve_tokens must be an integer")
    text_reserve_tokens = int(text_reserve_tokens)
    if text_reserve_tokens < 0 or text_reserve_tokens >= context_tokens:
        raise ValueError("text_reserve_tokens must lie in [0, context_tokens)")
    visual_tokens_per_frame = _positive_int("visual_tokens_per_frame", visual_tokens_per_frame)
    total_frames = int(math.floor(duration_seconds * source_fps + 1e-12))
    if total_frames < 1:
        raise ValueError("the floor-based source-frame convention yields zero frames")
    visual_budget = context_tokens - text_reserve_tokens
    max_frames = visual_budget // visual_tokens_per_frame
    if max_frames <= 0:
        raise ValueError("the context budget cannot fit one frame")
    selected_frames = min(total_frames, max_frames)
    sample_interval_seconds = duration_seconds / selected_frames
    used_visual_tokens = selected_frames * visual_tokens_per_frame
    full_prompt_tokens = text_reserve_tokens + used_visual_tokens
    kv_bytes = full_prompt_tokens * kv_spec.bytes_per_token()
    return {
        "total_frames": total_frames,
        "visual_budget_tokens": visual_budget,
        "max_frames": max_frames,
        "selected_frames": selected_frames,
        "sample_interval_seconds": sample_interval_seconds,
        "used_visual_tokens": used_visual_tokens,
        "visual_kv_gib": used_visual_tokens * kv_spec.bytes_per_token() / 2**30,
        "full_prompt_tokens": full_prompt_tokens,
        "full_prompt_kv_gib": kv_bytes / 2**30,
    }



V = REFERENCE["video"]
plain = video_budget(duration_seconds=7200, source_fps=1.0, context_tokens=256_000, text_reserve_tokens=4_000, visual_tokens_per_frame=576, kv_spec=spec)
merged_video = video_budget(duration_seconds=7200, source_fps=1.0, context_tokens=256_000, text_reserve_tokens=4_000, visual_tokens_per_frame=144, kv_spec=spec)
for key, value in V["without_merger"].items():
    assert math.isclose(plain[key], value, rel_tol=0, abs_tol=1e-12)
for key, value in V["with_merger_2x2"].items():
    assert math.isclose(merged_video[key], value, rel_tol=0, abs_tol=1e-12)
indices = uniform_sample_indices(plain["total_frames"], plain["selected_frames"])
indices_merged = uniform_sample_indices(merged_video["total_frames"], merged_video["selected_frames"])
assert len(indices) == 437 and len(indices_merged) == 1750
assert indices[0] == indices_merged[0] == 0 and indices[-1] == indices_merged[-1] == 7199
print('Video capacity, interval, frame indices, and full-prompt KV passed')
print(f"plain: {len(indices)} frames | 1/{plain['sample_interval_seconds']:.1f}s | visual {plain['visual_kv_gib']:.2f} GiB | full {plain['full_prompt_kv_gib']:.2f} GiB")
print(f"merge: {len(indices_merged)} frames | 1/{merged_video['sample_interval_seconds']:.1f}s | visual {merged_video['visual_kv_gib']:.2f} GiB | full {merged_video['full_prompt_kv_gib']:.2f} GiB")
print("first/last sampled IDs:", indices[:4].tolist(), "...", indices[-4:].tolist())
print("inverse density / mean timestamp gap (s):",
      plain['sample_interval_seconds'], float(np.diff(indices).mean()),
      merged_video['sample_interval_seconds'], float(np.diff(indices_merged).mean()))

def expect_error(label, fn, error_type=ValueError):
    """Run one negative test and require the declared exception type."""

    try:
        fn()
    except error_type:
        print("expected:", label)
        return
    raise AssertionError(f"{label}: expected {error_type.__name__}")


expect_error("resolution not divisible by patch", lambda: patch_layout(335, 336))
expect_error("grid not divisible by merger", lambda: patch_layout(42, 28, merge_factor=2))
expect_error("token count/grid mismatch", lambda: merge_spatial_tokens(np.zeros((7, 3)), (4, 2)))
expect_error("projector feature width mismatch", lambda: project_vision(np.zeros((1, 2, 7)), weights))
expect_error("zero image placeholders", lambda: assemble_multimodal_sequence(np.array([1, 2]), np.zeros((2, 16)), image_placeholder_id=999, image_embeddings=visual_embeddings[0], image_start_embedding=np.zeros(16), image_end_embedding=np.zeros(16)))
expect_error("two image placeholders", lambda: assemble_multimodal_sequence(np.array([999, 999]), np.zeros((2, 16)), image_placeholder_id=999, image_embeddings=visual_embeddings[0], image_start_embedding=np.zeros(16), image_end_embedding=np.zeros(16)))
expect_error("non-positive InfoNCE temperature", lambda: symmetric_infonce(np.eye(2), 0.0))
expect_error("zero embedding row", lambda: l2_normalize_rows(np.array([[0.0, 0.0], [1.0, 0.0]])))
expect_error("no KV budget after weights", lambda: image_kv_metrics(576, spec, device_budget_gib=10.0, weights_gib=10.0))
expect_error("context cannot fit one frame", lambda: video_budget(duration_seconds=10, source_fps=1, context_tokens=100, text_reserve_tokens=90, visual_tokens_per_frame=20, kv_spec=spec))
print('Expected errors were caught; silent truncation and hidden NaNs do not pass')

"""A small, fully trained pixels-to-token model; no pretrained VLM or external data.

Images are generated by this module. Held-out images change rendering brightness,
not the underlying bar-height combinations. This tests rendering transfer only.
"""
import io
import math
import torch
from torch import nn
from torch.nn import functional as F


def bar_data():
    """Return disjoint 8x8 grayscale renders of two bars with heights 1..4."""
    images, labels, split = [], [], []
    for brightness in (0.6, 1.0, 0.8):
        for left in range(1, 5):
            for right in range(1, 5):
                image = torch.zeros(1, 8, 8)
                image[0, 7-left:7, 1:3] = brightness
                image[0, 7-right:7, 5:7] = brightness
                images.append(image)
                # IDs: prompt=0, left=1, equal=2, right=3, EOS=4.
                labels.append([1 if left > right else 3 if right > left else 2, 4])
                split.append(brightness == 0.8)
    x, y, test = torch.stack(images), torch.tensor(labels), torch.tensor(split)
    assert len({row.numpy().tobytes() for row in x}) == 48
    return x[~test], y[~test], x[test], y[test]


class TinyVLM(nn.Module):
    """Patch encoder -> MLP connector -> one causal decoder -> answer and EOS.

    merge=2 uses lossless pixel unshuffle BEFORE the learned encoder: it changes
    the number of positions, not the number of observed pixels. This is not a ViT.
    """
    def __init__(self, merge=1):
        super().__init__()
        if type(merge) is not int or merge not in (1, 2):
            raise ValueError("merge must be 1 or 2")
        self.merge = merge
        self.n_visual = 16 // merge**2
        self.encoder = nn.Linear(4 * merge**2, 16)
        self.projector = nn.Sequential(nn.Linear(16, 32), nn.GELU(), nn.Linear(32, 32))
        self.embedding = nn.Embedding(5, 32)
        self.position = nn.Parameter(torch.randn(1, self.n_visual + 3, 32) * 0.02)
        self.norm1 = nn.LayerNorm(32)
        self.qkv = nn.Linear(32, 96)
        self.out = nn.Linear(32, 32)
        self.norm2 = nn.LayerNorm(32)
        self.ff = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32))
        self.head = nn.Linear(32, 5)

    def forward(self, images, text):
        # unfold: [B,C,H,W] -> [B,16,4], row-major 2x2 patches.
        b = len(images)
        patches = F.unfold(images, kernel_size=2, stride=2).transpose(1, 2)
        if self.merge == 2:
            patches = patches.reshape(b, 2, 2, 2, 2, 4).permute(0, 1, 3, 2, 4, 5).reshape(b, 4, 16)
        visual = self.projector(torch.tanh(self.encoder(patches)))
        x = torch.cat([visual, self.embedding(text)], dim=1)
        x = x + self.position[:, :x.shape[1]]
        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)
        scores = q @ k.transpose(-1, -2) / math.sqrt(32)
        allowed = torch.ones(x.shape[1], x.shape[1], dtype=torch.bool, device=x.device).tril()
        attention = scores.masked_fill(~allowed, float('-inf')).softmax(-1)
        x = x + self.out(attention @ v)
        x = x + self.ff(self.norm2(x))
        return self.head(x)

    def loss(self, images, answers):
        # Prompt predicts answer; answer predicts EOS. Exactly one causal shift.
        text = torch.cat([torch.zeros(len(images), 1, dtype=torch.long), answers[:, :1]], dim=1)
        logits = self(images, text)[:, self.n_visual:]
        return F.cross_entropy(logits.reshape(-1, 5), answers.reshape(-1))

    @torch.no_grad()
    def generate(self, images):
        text = torch.zeros(len(images), 1, dtype=torch.long)
        for _ in range(2):
            token = self(images, text)[:, -1].argmax(-1, keepdim=True)
            text = torch.cat([text, token], dim=1)
        return text[:, 1:]


def run_vision_lab(steps=240):
    """Train two declared models, save/reload, and return stored predictions."""
    if type(steps) is not int or steps <= 0:
        raise ValueError("steps must be a positive integer")
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    train_x, train_y, test_x, test_y = bar_data()
    metrics, records = {}, {}
    for merge in (1, 2):
        torch.manual_seed(1600 + merge)
        model = TinyVLM(merge)
        initial = float(model.loss(test_x, test_y).detach())
        optimizer = torch.optim.AdamW(model.parameters(), lr=0.005, weight_decay=0.01)
        losses = []
        for _ in range(steps):
            loss = model.loss(train_x, train_y)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
            optimizer.step()
            losses.append(float(loss.detach()))
        model.eval()
        predictions = model.generate(test_x)
        blank_predictions = model.generate(torch.zeros_like(test_x))
        train_predictions = model.generate(train_x)
        # The first answer logit cannot see the teacher-forced answer token.
        a = torch.zeros(len(test_x), 2, dtype=torch.long)
        b = a.clone(); b[:, 1] = 3
        with torch.no_grad():
            torch.testing.assert_close(model(test_x, a)[:, model.n_visual], model(test_x, b)[:, model.n_visual])
        # Serialization must preserve the model's actual outputs.
        buffer = io.BytesIO()
        torch.save(model.state_dict(), buffer)
        buffer.seek(0)
        restored = TinyVLM(merge)
        restored.load_state_dict(torch.load(buffer, weights_only=True))
        restored.eval()
        assert torch.equal(restored.generate(test_x), predictions)
        # Freezing the decoder must NOT cut the input gradient to the projector.
        for name, p in model.named_parameters():
            p.requires_grad_(name.startswith('projector.'))
            p.grad = None
        model.loss(train_x, train_y).backward()
        assert sum(float(p.grad.abs().sum()) for p in model.projector.parameters()) > 0
        assert all(p.grad is None for name, p in model.named_parameters() if not name.startswith('projector.'))
        key = str(merge)
        metrics[key] = dict(visual_tokens=model.n_visual, parameters=sum(p.numel() for p in model.parameters()),
            train_images=len(train_x), test_images=len(test_x), steps=steps,
            initial_test_ce=initial, final_test_ce=float(model.loss(test_x, test_y).detach()),
            train_correct=int((train_predictions == train_y).all(-1).sum()),
            exact_correct=int((predictions == test_y).all(-1).sum()),
            blank_correct=int((blank_predictions == test_y).all(-1).sum()))
        records[key] = dict(predictions=predictions.tolist(), targets=test_y.tolist(),
                            blank_predictions=blank_predictions.tolist(), loss=losses)
    return metrics, records



vision_metrics, vision_records = run_vision_lab()
print(json.dumps(vision_metrics, indent=2))

"""Audio framing, temporal pooling and a tiny residual quantizer, all on CPU.

The spectra and hand-written codebooks are teaching examples, not a speech model
or a trained audio codec. Codec-rate accounting is a separate declared scenario.
"""
from numbers import Integral
import math
import numpy as np


def count(value, name):
    if isinstance(value, bool) or not isinstance(value, Integral) or value <= 0:
        raise ValueError(name + ' must be a positive integer')
    return int(value)


def frame_signal(waveform, window=400, hop=160):
    """Start a window at 0, hop, ... < length; right-pad the last windows."""
    signal = np.asarray(waveform)
    window, hop = count(window, 'window'), count(hop, 'hop')
    if signal.ndim != 1 or signal.size == 0 or signal.dtype.kind not in 'fiu':
        raise ValueError('waveform must be a nonempty real numeric vector')
    if not np.isfinite(signal).all():
        raise ValueError('waveform must be finite')
    starts = np.arange(0, len(signal), hop)
    indices = starts[:, None] + np.arange(window)[None, :]
    valid = indices < len(signal)
    frames = np.zeros(indices.shape, dtype=np.float64)
    frames[valid] = signal[indices[valid]]
    return frames, valid


def pool_frames(features, factor=4):
    """Average consecutive frame vectors; the final short group uses its own size."""
    x = np.asarray(features, dtype=np.float64)
    factor = count(factor, 'factor')
    if x.ndim != 2 or 0 in x.shape or not np.isfinite(x).all():
        raise ValueError('features must be a finite nonempty matrix')
    return np.stack([x[i:i+factor].mean(axis=0) for i in range(0, len(x), factor)])


def residual_quantize(features, codebooks):
    """Greedily encode residuals; books have shape [books, entries, width]."""
    x = np.asarray(features, dtype=np.float64)
    books = np.asarray(codebooks, dtype=np.float64)
    if x.ndim != 2 or 0 in x.shape or books.ndim != 3 or 0 in books.shape:
        raise ValueError('nonempty feature matrix and codebook tensor required')
    if x.shape[1] != books.shape[2] or not np.isfinite(x).all() or not np.isfinite(books).all():
        raise ValueError('widths must agree and values must be finite')
    residual = x.copy()
    reconstructed = np.zeros_like(x)
    ids, squared_errors = [], [np.square(residual).sum(axis=1)]
    for book in books:
        distances = np.square(residual[:, None, :] - book[None, :, :]).sum(axis=-1)
        chosen = distances.argmin(axis=1)
        reconstructed += book[chosen]
        residual -= book[chosen]
        ids.append(chosen)
        squared_errors.append(np.square(residual).sum(axis=1))
    return np.stack(ids, axis=1), reconstructed, np.stack(squared_errors, axis=1)


def audio_budget(n_samples=96000, sample_rate=16000, hop=160, pooling=4,
                 codec_rate=75, books=8, entries=1024):
    """Two separate paths: acoustic features and illustrative discrete codec codes."""
    n_samples, sample_rate, hop, pooling, codec_rate, books, entries = (
        count(v,n) for v,n in zip((n_samples,sample_rate,hop,pooling,codec_rate,books,entries),
                                ('samples','sample_rate','hop','pooling','codec_rate','books','entries')))
    if entries < 2 or entries & (entries-1):
        raise ValueError('entries must be a power of two >= 2 for fixed-width packing')
    frames = (n_samples+hop-1)//hop
    codec_frames = (n_samples*codec_rate+sample_rate-1)//sample_rate
    ids = codec_frames*books
    return dict(samples=n_samples, seconds=n_samples/sample_rate,
                acoustic_frames=frames, pooled_positions=(frames+pooling-1)//pooling,
                codec_frames=codec_frames, codec_ids=ids,
                codec_payload_bits=ids*(entries.bit_length()-1),
                codec_bits_per_second=codec_rate*books*(entries.bit_length()-1))


def run_audio_example():
    t = np.arange(96000)/16000
    wave = .6*np.sin(2*np.pi*220*t) + .2*np.sin(2*np.pi*660*t)
    frames, valid = frame_signal(wave)
    # Simple magnitude spectra: not log-mel features or a trained audio encoder.
    features = np.abs(np.fft.rfft(frames*np.hanning(400), axis=1))
    pooled = pool_frames(features)
    assert frames.shape == (600,400) and features.shape == (600,201)
    assert pooled.shape == (150,201) and valid[-1].sum() == 160
    budget = audio_budget()
    assert budget['acoustic_frames'] == len(frames)
    assert budget['pooled_positions'] == len(pooled)
    books = np.array([[[0,0],[1,0],[0,-1]], [[0,0],[.5,0],[.5,-.5]]], dtype=float)
    ids, reconstructed, errors = residual_quantize([[1.6,-.6]], books)
    assert ids.tolist() == [[1,2]]
    np.testing.assert_allclose(reconstructed, [[1.5,-.5]], atol=1e-15)
    np.testing.assert_allclose(errors, [[2.92,.72,.02]], atol=1e-15)
    # Zero is available in each book, so this toy can never increase residual norm.
    assert np.all(np.diff(errors, axis=1) <= 0)
    # Illustrative timings; no speech decoder is measured by this script.
    rtf = 2.4/6
    assert math.isclose(rtf,.4)
    result = dict(budget=budget, spectral_shape=list(features.shape),
                  pooled_shape=list(pooled.shape), rvq_ids=ids.tolist(),
                  rvq_squared_errors=errors.tolist(), illustrative_rtf=rtf,
                  illustrative_first_audio_seconds=.18)
    return result



audio_result = run_audio_example()
print(json.dumps(audio_result, indent=2))

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

import importlib.util, subprocess, sys, json
if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
"""A small, fully trained pixels-to-token model; no pretrained VLM or external data.

Images are generated by this module. Held-out images change rendering brightness,
not the underlying bar-height combinations. This tests rendering transfer only.
"""
import io
import math
import torch
from torch import nn
from torch.nn import functional as F


def bar_data():
    """Return disjoint 8x8 grayscale renders of two bars with heights 1..4."""
    images, labels, split = [], [], []
    for brightness in (0.6, 1.0, 0.8):
        for left in range(1, 5):
            for right in range(1, 5):
                image = torch.zeros(1, 8, 8)
                image[0, 7-left:7, 1:3] = brightness
                image[0, 7-right:7, 5:7] = brightness
                images.append(image)
                # IDs: prompt=0, left=1, equal=2, right=3, EOS=4.
                labels.append([1 if left > right else 3 if right > left else 2, 4])
                split.append(brightness == 0.8)
    x, y, test = torch.stack(images), torch.tensor(labels), torch.tensor(split)
    assert len({row.numpy().tobytes() for row in x}) == 48
    return x[~test], y[~test], x[test], y[test]


class TinyVLM(nn.Module):
    """Patch encoder -> MLP connector -> one causal decoder -> answer and EOS.

    merge=2 uses lossless pixel unshuffle BEFORE the learned encoder: it changes
    the number of positions, not the number of observed pixels. This is not a ViT.
    """
    def __init__(self, merge=1):
        super().__init__()
        if type(merge) is not int or merge not in (1, 2):
            raise ValueError("merge must be 1 or 2")
        self.merge = merge
        self.n_visual = 16 // merge**2
        self.encoder = nn.Linear(4 * merge**2, 16)
        self.projector = nn.Sequential(nn.Linear(16, 32), nn.GELU(), nn.Linear(32, 32))
        self.embedding = nn.Embedding(5, 32)
        self.position = nn.Parameter(torch.randn(1, self.n_visual + 3, 32) * 0.02)
        self.norm1 = nn.LayerNorm(32)
        self.qkv = nn.Linear(32, 96)
        self.out = nn.Linear(32, 32)
        self.norm2 = nn.LayerNorm(32)
        self.ff = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32))
        self.head = nn.Linear(32, 5)

    def forward(self, images, text):
        # unfold: [B,C,H,W] -> [B,16,4], row-major 2x2 patches.
        b = len(images)
        patches = F.unfold(images, kernel_size=2, stride=2).transpose(1, 2)
        if self.merge == 2:
            patches = patches.reshape(b, 2, 2, 2, 2, 4).permute(0, 1, 3, 2, 4, 5).reshape(b, 4, 16)
        visual = self.projector(torch.tanh(self.encoder(patches)))
        x = torch.cat([visual, self.embedding(text)], dim=1)
        x = x + self.position[:, :x.shape[1]]
        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)
        scores = q @ k.transpose(-1, -2) / math.sqrt(32)
        allowed = torch.ones(x.shape[1], x.shape[1], dtype=torch.bool, device=x.device).tril()
        attention = scores.masked_fill(~allowed, float('-inf')).softmax(-1)
        x = x + self.out(attention @ v)
        x = x + self.ff(self.norm2(x))
        return self.head(x)

    def loss(self, images, answers):
        # Prompt predicts answer; answer predicts EOS. Exactly one causal shift.
        text = torch.cat([torch.zeros(len(images), 1, dtype=torch.long), answers[:, :1]], dim=1)
        logits = self(images, text)[:, self.n_visual:]
        return F.cross_entropy(logits.reshape(-1, 5), answers.reshape(-1))

    @torch.no_grad()
    def generate(self, images):
        text = torch.zeros(len(images), 1, dtype=torch.long)
        for _ in range(2):
            token = self(images, text)[:, -1].argmax(-1, keepdim=True)
            text = torch.cat([text, token], dim=1)
        return text[:, 1:]


def run_vision_lab(steps=240):
    """Train two declared models, save/reload, and return stored predictions."""
    if type(steps) is not int or steps <= 0:
        raise ValueError("steps must be a positive integer")
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    train_x, train_y, test_x, test_y = bar_data()
    metrics, records = {}, {}
    for merge in (1, 2):
        torch.manual_seed(1600 + merge)
        model = TinyVLM(merge)
        initial = float(model.loss(test_x, test_y).detach())
        optimizer = torch.optim.AdamW(model.parameters(), lr=0.005, weight_decay=0.01)
        losses = []
        for _ in range(steps):
            loss = model.loss(train_x, train_y)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
            optimizer.step()
            losses.append(float(loss.detach()))
        model.eval()
        predictions = model.generate(test_x)
        blank_predictions = model.generate(torch.zeros_like(test_x))
        train_predictions = model.generate(train_x)
        # The first answer logit cannot see the teacher-forced answer token.
        a = torch.zeros(len(test_x), 2, dtype=torch.long)
        b = a.clone(); b[:, 1] = 3
        with torch.no_grad():
            torch.testing.assert_close(model(test_x, a)[:, model.n_visual], model(test_x, b)[:, model.n_visual])
        # Serialization must preserve the model's actual outputs.
        buffer = io.BytesIO()
        torch.save(model.state_dict(), buffer)
        buffer.seek(0)
        restored = TinyVLM(merge)
        restored.load_state_dict(torch.load(buffer, weights_only=True))
        restored.eval()
        assert torch.equal(restored.generate(test_x), predictions)
        # Freezing the decoder must NOT cut the input gradient to the projector.
        for name, p in model.named_parameters():
            p.requires_grad_(name.startswith('projector.'))
            p.grad = None
        model.loss(train_x, train_y).backward()
        assert sum(float(p.grad.abs().sum()) for p in model.projector.parameters()) > 0
        assert all(p.grad is None for name, p in model.named_parameters() if not name.startswith('projector.'))
        key = str(merge)
        metrics[key] = dict(visual_tokens=model.n_visual, parameters=sum(p.numel() for p in model.parameters()),
            train_images=len(train_x), test_images=len(test_x), steps=steps,
            initial_test_ce=initial, final_test_ce=float(model.loss(test_x, test_y).detach()),
            train_correct=int((train_predictions == train_y).all(-1).sum()),
            exact_correct=int((predictions == test_y).all(-1).sum()),
            blank_correct=int((blank_predictions == test_y).all(-1).sum()))
        records[key] = dict(predictions=predictions.tolist(), targets=test_y.tolist(),
                            blank_predictions=blank_predictions.tolist(), loss=losses)
    return metrics, records



plt.rcParams.update({"font.family": "DejaVu Sans", "figure.figsize": (11, 6),
    "figure.dpi": 150, "font.size": 12, "axes.titlesize": 17, "axes.labelsize": 13,
    "xtick.labelsize": 11, "ytick.labelsize": 11, "legend.fontsize": 11})
metrics, records = run_vision_lab()
print(json.dumps(metrics, indent=2))
for key, record in records.items():
    plt.semilogy(record["loss"], label=str(metrics[key]["visual_tokens"]) +
                 (" визуальных позиций" if DEMO_LANGUAGE == "ru" else " visual positions"))
plt.xlabel("Шаг оптимизации" if DEMO_LANGUAGE == "ru" else "Optimization step")
plt.ylabel("Обучающая кросс-энтропия" if DEMO_LANGUAGE == "ru" else "Training cross-entropy")
plt.legend(); plt.tight_layout(); plt.show()

# Consume saved predictions; routing and scoring never invoke the VLM again.
ALLOWED = {1: "mark_left", 2: "mark_equal", 3: "mark_right"}

def route_saved(tokens):
    if (not isinstance(tokens, list) or len(tokens) != 2
            or any(type(t) is not int for t in tokens)
            or tokens[1] != 4 or tokens[0] not in ALLOWED):
        return "clarify"
    return ALLOWED[tokens[0]]

traces = {}
for key, record in records.items():
    log = []
    for request_id, (prediction, target) in enumerate(zip(record["predictions"], record["targets"])):
        action = route_saved(prediction)
        # The mock tool only records an annotation; it has no external side effect.
        assert action in set(ALLOWED.values()) | {"clarify"}
        log.append(dict(request_id=request_id, prediction=prediction, action=action,
                        correct=action == route_saved(target)))
    assert len(log) == metrics[key]["test_images"] == 16
    assert sum(row["correct"] for row in log) == metrics[key]["exact_correct"]
    traces[key] = log
assert route_saved(["delete_files", 4]) == "clarify"
assert route_saved([True, 4]) == "clarify"
assert route_saved([1, 0]) == "clarify"
assert json.loads(json.dumps(traces)) == traces
print({key: {"recorded_actions": len(rows), "correct": sum(r["correct"] for r in rows)}
       for key, rows in traces.items()})