# examples/llava.py
"""Offline CPU reference for original LLaVA's learning contract.

Authored reduced mechanics, not pretrained CLIP/Vicuna or benchmark
reproduction.
Primary paper: https://arxiv.org/html/2304.08485v2#S4
Historical code convention (July 23, 2023, before LLaVA-1.5):
Primary code: https://github.com/haotian-liu/LLaVA
Revision: 5e89b7ed725f6bd026aa0245d5a0f948e85d613a
llava_arch.py: biased single Linear, IMAGE_TOKEN_INDEX expansion,
ignored images.
clip_encoder.py + train.py: patch-only selected hidden state, default layer -1.
llava_llama.py: one causal shift and mean valid-target cross entropy.
Right-padding/position reconstruction and role masks are authored explicitly;
we omit tokenizer templates, image wrappers, checkpoint loading, SGD recipe
fidelity, and special added-vocabulary training. Synthetic features stand in
for the selected vision hidden state; a one-block width-eight LM replaces LLaMA.
"""

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

IMAGE = -200
PAD = 0
IGNORE = -100
UPSTREAM_REVISION = "5e89b7ed725f6bd026aa0245d5a0f948e85d613a"


# region: vision_features
def vision_features(vision, selected_inputs):
    """Extract patch features from a frozen selected-state stand-in.

    Args:
        vision: Frozen linear module mapping three channels to d_v.
        selected_inputs: Float [B, N+1, 3], including a leading CLS slot.

    Returns:
        Detached patch features [B, N, d_v], with no pooled/CLS token.
    """
    if selected_inputs.ndim != 3 or selected_inputs.shape[1] < 2:
        raise ValueError("Expected [B, N+1, 3] with at least one patch")
    if selected_inputs.shape[-1] != vision.in_features:
        raise ValueError("Vision input channels do not match")
    if not torch.isfinite(selected_inputs).all():
        raise ValueError("Vision inputs must be finite")
    if any(p.requires_grad for p in vision.parameters()):
        raise ValueError("Vision must remain frozen")
    with torch.no_grad():
        selected_hidden_state = vision(selected_inputs)
        return selected_hidden_state[:, 1:, :]


# endregion: vision_features


# region: linear_projector
def linear_projector(features, projector):
    """Map row-major patches through one biased linear layer, not an MLP."""
    if features.ndim != 3 or features.shape[1] < 1:
        raise ValueError("Expected nonempty [B, N, d_v] patch features")
    if features.shape[-1] != projector.in_features:
        raise ValueError("Feature width must match projector")
    if not torch.isfinite(features).all():
        raise ValueError("Patch features must be finite")
    # W: [d_l, d_v]; b: [d_l]; output: [B, N, d_l].
    return features @ projector.weight.T + projector.bias


# endregion: linear_projector


# region: splice_image
def splice_image(ids, text_embeddings, visual_sequences, labels, attention):
    """Replace one IMAGE sentinel per visual sample and right-pad the batch.

    visual_sequences is a B-long list of [N_b, d_l] tensors or None for
    text-only samples. Labels are already masked by assistant_labels.
    Returns embeddings, attention, positions, labels with identical B/L axes.
    """
    if ids.ndim != 2 or ids.dtype != torch.long:
        raise ValueError("Token IDs must be long [B, T]")
    if text_embeddings.shape[:2] != ids.shape or text_embeddings.ndim != 3:
        raise ValueError("Text embeddings must be [B, T, d_l]")
    if labels.shape != ids.shape or labels.dtype != torch.long:
        raise ValueError("Labels must be long [B, T]")
    if attention.shape != ids.shape or attention.dtype != torch.bool:
        raise ValueError("Attention must be boolean [B, T]")
    if len(visual_sequences) != ids.shape[0] or ids.shape[0] < 1:
        raise ValueError("Supply one optional visual tensor per sample")
    if (attention[:, 1:] & ~attention[:, :-1]).any():
        raise ValueError(
            "Only contiguous valid prefix / right padding is supported"
        )
    if ((ids == PAD) & attention).any() or (ids[~attention] != PAD).any():
        raise ValueError("PAD IDs must coincide with masked suffix positions")
    sequences, targets = [], []
    for b in range(ids.shape[0]):
        count = int(attention[b].sum())
        if count < 1:
            raise ValueError("Every sample needs a valid context token")
        where = (ids[b, :count] == IMAGE).nonzero().flatten()
        visual = visual_sequences[b]
        if len(where) != (0 if visual is None else 1):
            raise ValueError(
                "Exactly one IMAGE is required for each supplied image"
            )
        x, y = text_embeddings[b, :count], labels[b, :count]
        if visual is not None:
            if (
                visual.ndim != 2
                or visual.shape[0] < 1
                or visual.shape[1] != x.shape[1]
            ):
                raise ValueError(
                    "Visual embeddings must be nonempty [N_b, d_l]"
                )
            if not torch.isfinite(visual).all():
                raise ValueError("Visual embeddings must be finite")
            k = int(where[0])
            x = torch.cat((x[:k], visual, x[k + 1 :]), dim=0)
            image_labels = torch.full((visual.shape[0],), IGNORE, dtype=y.dtype)
            y = torch.cat((y[:k], image_labels, y[k + 1 :]))
        sequences.append(x)
        targets.append(y)
    length = max(x.shape[0] for x in sequences)
    embeddings = torch.stack(
        [F.pad(x, (0, 0, 0, length - len(x))) for x in sequences]
    )
    new_labels = torch.stack(
        [F.pad(y, (0, length - len(y)), value=IGNORE) for y in targets]
    )
    lengths = torch.tensor([len(x) for x in sequences])
    new_attention = torch.arange(length)[None, :] < lengths[:, None]
    positions = (new_attention.long().cumsum(1) - 1).clamp_min(0)
    positions = positions.masked_fill(~new_attention, 0)
    return embeddings, new_attention, positions, new_labels


# endregion: splice_image


# region: assistant_labels
def assistant_labels(ids, assistant, attention, vocabulary):
    """Keep assistant response/stop-token targets, ignore all context/padding.

    Role masks are supplied explicitly rather than inferred from a tokenizer.
    Ignored targets have zero direct loss, but context embeddings can get
    gradients through causal attention to supervised assistant predictions.
    """
    if ids.ndim != 2 or ids.dtype != torch.long:
        raise ValueError("Token IDs must be long [B, T]")
    if assistant.shape != ids.shape or assistant.dtype != torch.bool:
        raise ValueError("Assistant role mask must be boolean [B, T]")
    if attention.shape != ids.shape or attention.dtype != torch.bool:
        raise ValueError("Attention mask must be boolean [B, T]")
    if (
        vocabulary < 2
        or ((ids != IMAGE) & ((ids < 0) | (ids >= vocabulary))).any()
    ):
        raise ValueError("IDs must be IMAGE or valid vocabulary tokens")
    valid = assistant & attention & (ids != IMAGE) & (ids != PAD)
    return ids.masked_fill(~valid, IGNORE)


# endregion: assistant_labels


# region: causal_loss
def causal_loss(logits, labels):
    """Shift exactly once, averaging NLL over valid assistant targets."""
    if (
        logits.ndim != 3
        or labels.shape != logits.shape[:2]
        or labels.dtype != torch.long
    ):
        raise ValueError("Expected logits [B, L, V] and long labels [B, L]")
    if not torch.isfinite(logits).all():
        raise ValueError("Logits must be finite")
    if (
        (labels != IGNORE) & ((labels < 0) | (labels >= logits.shape[-1]))
    ).any():
        raise ValueError("Target is outside vocabulary")
    shifted_logits = logits[:, :-1, :]
    shifted_labels = labels[:, 1:]
    valid = shifted_labels != IGNORE
    if not valid.any():
        # Defined finite empty-target extension, unlike mean CE's NaN.
        return logits.sum() * 0.0
    return F.cross_entropy(shifted_logits[valid], shifted_labels[valid])


# endregion: causal_loss


class TinyLM(nn.Module):
    """Random causal attention + FFN stand-in; no meaningful language output."""

    def __init__(self, vocabulary=24, width=8):
        super().__init__()
        self.embedding = nn.Embedding(vocabulary, width)
        self.position = nn.Embedding(64, width)
        self.qkv = nn.Linear(width, 3 * width)
        self.out = nn.Linear(width, width)
        self.ffn = nn.Sequential(
            nn.Linear(width, 16), nn.GELU(), nn.Linear(16, width)
        )
        self.head = nn.Linear(width, vocabulary, bias=False)

    def forward(self, embeddings, attention, positions):
        """Attend only to valid keys at the current or earlier position."""
        h = embeddings + self.position(positions)
        q, k, v = self.qkv(h).chunk(3, dim=-1)
        scores = q @ k.transpose(-1, -2) / math.sqrt(h.shape[-1])
        length = h.shape[1]
        causal = torch.ones(length, length, dtype=torch.bool).tril()
        allowed = causal[None] & attention[:, None, :]
        weights = scores.masked_fill(~allowed, -torch.inf).softmax(-1)
        h = h + self.out(weights @ v)
        h = h + self.ffn(h)
        return self.head(h)


class ReducedLLaVA(nn.Module):
    """Reduced dimensions: three input channels, d_v=6, d_l=8, V=24."""

    def __init__(self):
        super().__init__()
        self.vision = nn.Linear(3, 6).requires_grad_(False)
        self.projector = nn.Linear(6, 8)
        self.lm = TinyLM()

    def assemble(self, ids, assistant, attention, selected_inputs):
        """Execute extraction, projection, role masking and token insertion."""
        patches = vision_features(self.vision, selected_inputs)
        projected = linear_projector(patches, self.projector)
        labels = assistant_labels(ids, assistant, attention, 24)
        # IMAGE is never passed to the vocabulary embedding lookup.
        text = self.lm.embedding(ids.masked_fill(ids == IMAGE, PAD))
        return splice_image(ids, text, list(projected), labels, attention)


# region: training_stage
def configure_stage(model, stage):
    """Stage 1 trains only projector; stage 2 trains projector plus LM."""
    if stage not in (1, 2):
        raise ValueError("Stage must be 1 or 2")
    model.zero_grad(set_to_none=True)
    model.vision.requires_grad_(False)
    model.projector.requires_grad_(True)
    model.lm.requires_grad_(stage == 2)


def train_step(model, batch, stage, learning_rate=0.05):
    """Run one illustrative SGD update and expose input/parameter gradients."""
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("Learning rate must be finite and positive")
    configure_stage(model, stage)
    optimizer = torch.optim.SGD(
        (p for p in model.parameters() if p.requires_grad), lr=learning_rate
    )
    embeddings, attention, positions, labels = model.assemble(*batch)
    embeddings.retain_grad()
    # Do not put the frozen LM in no_grad: stage 1 needs dL/d(image embeddings).
    logits = model.lm(embeddings, attention, positions)
    loss = causal_loss(logits, labels)
    loss.backward()
    observations = {
        "loss": float(loss.detach()),
        "projector_gradient_norm": float(
            model.projector.weight.grad.norm().detach()
        ),
        "input_gradient_norm": float(embeddings.grad.norm().detach()),
        "vision_gradients_absent": all(
            p.grad is None for p in model.vision.parameters()
        ),
        "lm_gradients_absent": all(
            p.grad is None for p in model.lm.parameters()
        ),
    }
    optimizer.step()
    return observations


# endregion: training_stage


def smoke(seed=7):
    """Return measured finite mechanics, never semantic visual answers."""
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    model = ReducedLLaVA()
    ids = torch.tensor([[1, IMAGE, 4, 5, 6, 2], [IMAGE, 1, 7, 8, 2, PAD]])
    assistant = torch.tensor(
        [
            [False, False, False, True, True, True],
            [False, False, False, True, True, False],
        ]
    )
    attention = ids != PAD
    inputs = torch.randn(2, 4, 3)
    batch = (ids, assistant, attention, inputs)
    embedded, mask, positions, labels = model.assemble(*batch)
    projected = linear_projector(
        vision_features(model.vision, inputs), model.projector
    )
    stages = {}
    for stage in (1, 2):
        before = {
            name: p.detach().clone() for name, p in model.named_parameters()
        }
        observation = train_step(model, batch, stage)
        for prefix in ("vision", "projector", "lm"):
            observation[prefix + "_parameter_delta"] = sum(
                float((p.detach() - before[name]).abs().sum())
                for name, p in model.named_parameters()
                if name.startswith(prefix + ".")
            )
        stages[str(stage)] = observation
    probe = torch.randn(2, embedded.shape[1], 24, requires_grad=True)
    probe_loss = causal_loss(probe, labels)
    probe_loss.backward()
    ignored = labels[:, 1:] == IGNORE
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "upstream_revision": UPSTREAM_REVISION,
        "seed": seed,
        "boundary": (
            "offline synthetic CPU; random vision/LM; "
            "no checkpoints or semantic answers"
        ),
        "reductions": {
            "input_channels": 3,
            "vision_width": 6,
            "lm_width": 8,
            "patches": 3,
            "vocabulary": 24,
            "lm_blocks": 1,
            "optimizer": "one SGD step per stage",
        },
        "shapes": {
            "selected_inputs": list(inputs.shape),
            "projected": list(projected.shape),
            "spliced_embeddings": list(embedded.shape),
            "logits": list(probe.shape),
        },
        "valid_lengths": mask.sum(1).tolist(),
        "positions": positions.tolist(),
        "labels": labels.tolist(),
        "valid_shifted_assistant_targets": int((labels[:, 1:] != IGNORE).sum()),
        "ignored_target_logit_gradient_max": float(
            probe.grad[:, :-1][ignored].abs().max()
        ),
        "final_unshifted_logit_gradient_max": float(
            probe.grad[:, -1].abs().max()
        ),
        "probe_loss": float(probe_loss.detach()),
        "stages": stages,
    }


def main():
    """Print actual smoke measurements with strict finite JSON encoding."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--smoke", action="store_true", help="Run offline reduced CPU mechanics"
    )
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if args.smoke:
        print(json.dumps(smoke(args.seed), indent=2, allow_nan=False))
    else:
        parser.print_help()


if __name__ == "__main__":
    main()