# examples/blip2_training.py
"""Run reduced, random-weight BLIP-2 training mechanics without downloads.

Reference: Li et al., BLIP-2, sections 3.1-3.3 (arXiv:2301.12597).
Official LAVIS revision: 05cb670f47296522e24e7727e366b4a997d28582.
This independently authored teaching model is not checkpoint-compatible. Its
frozen patch convolution and tiny causal LM are random, not ViT/OPT releases.
Q-Former keeps shared self-attention, separate query/text FFNs, and image
cross-attention every other block. Pre-norm, no dropout, recomputed prefixes,
cyclic negatives, and unsmoothed losses simplify the released implementation.
No tokenizer, BERT initialization, distributed negatives, or T5 path is included.
Token IDs have no learned language meaning: PAD=0, CLS=1, DEC=2, EOS=3.
"""

import argparse
import hashlib
import json
import platform
from pathlib import Path

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

PAD, CLS, DEC, EOS = 0, 1, 2, 3


def validate_tokens(tokens, vocabulary, batch):
    """Reject invalid IDs and non-right-padded sequences before attention."""
    if tokens.ndim != 2 or tokens.shape[0] != batch or tokens.shape[1] == 0:
        raise ValueError("tokens must have shape [batch, nonempty sequence]")
    if tokens.dtype != torch.long:
        raise ValueError("token IDs must have dtype torch.long")
    if (tokens < 0).any() or (tokens >= vocabulary).any():
        raise ValueError("token IDs must be in [0, vocabulary)")
    if (tokens[:, 0] == PAD).any():
        raise ValueError("each sequence must start with a non-padding token")
    if ((tokens[:, :-1] == PAD) & (tokens[:, 1:] != PAD)).any():
        raise ValueError("only right padding is supported")


# region: masks
def visibility(mode, queries, text, device=None):
    """Return allowed attention edges; rows read columns (True = visible)."""
    if mode not in {"itc", "itm", "itg"} or queries < 1 or text < 0:
        raise ValueError(
            "expected itc/itm/itg, positive queries, nonnegative text"
        )
    allowed = torch.ones(
        queries + text, queries + text, dtype=torch.bool, device=device
    )
    if mode == "itc":
        allowed[:queries, queries:] = False
        allowed[queries:, :queries] = False
    elif mode == "itg":
        allowed[:queries, queries:] = False
        allowed[queries:, queries:] = torch.ones(
            text, text, dtype=torch.bool, device=device
        ).tril()
    return allowed


# endregion: masks


class QueryBlock(nn.Module):
    """Shared self-attention; image cross-attention touches only queries."""

    def __init__(self, width, heads, cross):
        super().__init__()
        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)
        self.norm = nn.LayerNorm(width)
        self.cross = (
            nn.MultiheadAttention(width, heads, batch_first=True)
            if cross
            else None
        )
        self.cross_norm = nn.LayerNorm(width)
        self.query_ffn = self.make_ffn(width)
        self.text_ffn = self.make_ffn(width)

    @staticmethod
    def make_ffn(width):
        """Build a separate pre-normalized token-wise feed-forward path."""
        return nn.Sequential(
            nn.LayerNorm(width),
            nn.Linear(width, 4 * width),
            nn.GELU(),
            nn.Linear(4 * width, width),
        )

    # region: qformer
    def forward(self, hidden, image, queries, allowed, padding):
        normalized = self.norm(hidden)
        update = self.attention(
            normalized,
            normalized,
            normalized,
            attn_mask=~allowed,
            key_padding_mask=padding,
            need_weights=False,
        )[0]
        hidden = hidden + update
        query, text = hidden[:, :queries], hidden[:, queries:]
        if self.cross is not None:
            query = (
                query
                + self.cross(
                    self.cross_norm(query), image, image, need_weights=False
                )[0]
            )
        query = query + self.query_ffn(query)
        text = text + self.text_ffn(text)
        return torch.cat([query, text], dim=1)

    # endregion: qformer


class QFormer(nn.Module):
    """Four small blocks reuse weights across ITC, ITM, ITG and stage 2."""

    def __init__(self, width, heads, queries, vocabulary):
        super().__init__()
        self.queries = nn.Parameter(torch.randn(1, queries, width) * 0.02)
        self.embedding = nn.Embedding(vocabulary, width)
        self.position = nn.Embedding(128, width)
        self.blocks = nn.ModuleList(
            QueryBlock(width, heads, cross=(index % 2 == 0))
            for index in range(4)
        )
        self.norm = nn.LayerNorm(width)

    def forward(self, image, tokens=None, mode="itc"):
        batch = image.shape[0]
        queries = self.queries.shape[1]
        query = self.queries.expand(batch, -1, -1)
        if tokens is None:
            text = query.new_empty(batch, 0, query.shape[-1])
            text_padding = torch.zeros(
                batch, 0, dtype=torch.bool, device=image.device
            )
        else:
            validate_tokens(tokens, self.embedding.num_embeddings, batch)
            if tokens.shape[1] > self.position.num_embeddings:
                raise ValueError("Q-Former text exceeds 128 positions")
            positions = torch.arange(tokens.shape[1], device=tokens.device)
            text = self.embedding(tokens) + self.position(positions)
            text_padding = tokens.eq(PAD)
        hidden = torch.cat([query, text], dim=1)
        padding = torch.cat(
            [
                torch.zeros(
                    batch, queries, dtype=torch.bool, device=image.device
                ),
                text_padding,
            ],
            1,
        )
        allowed = visibility(mode, queries, text.shape[1], image.device)
        for block in self.blocks:
            hidden = block(hidden, image, queries, allowed, padding)
        return self.norm(hidden)


class TinyCausalLM(nn.Module):
    """Frozen random causal transformer; differentiable with respect to inputs."""

    def __init__(self, width, heads, vocabulary):
        super().__init__()
        self.embedding = nn.Embedding(vocabulary, width)
        self.position = nn.Embedding(128, width)
        self.blocks = nn.ModuleList(
            nn.TransformerEncoderLayer(
                width,
                heads,
                4 * width,
                dropout=0,
                batch_first=True,
                norm_first=True,
            )
            for _ in range(2)
        )
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(width, vocabulary)

    def forward(self, embeddings, padding):
        length = embeddings.shape[1]
        if length > self.position.num_embeddings:
            raise ValueError("prefix plus text exceeds 128 LM positions")
        hidden = embeddings + self.position(
            torch.arange(length, device=embeddings.device)
        )
        # Unlike stage-1 ITG, the LM prefix is also causally masked.
        blocked = torch.ones(
            length, length, dtype=torch.bool, device=embeddings.device
        ).triu(1)
        for block in self.blocks:
            hidden = block(
                hidden, src_mask=blocked, src_key_padding_mask=padding
            )
        return self.head(self.norm(hidden))


# region: shift
def shifted_loss(logits, labels):
    """Shift the complete sequence, ignoring -100 targets after shifting."""
    if logits.ndim != 3 or labels.shape != logits.shape[:2]:
        raise ValueError("logits [B, S, V] and labels [B, S] must align")
    if labels.dtype != torch.long:
        raise ValueError("labels must have dtype torch.long")
    valid = labels.ne(-100)
    if ((labels[valid] < 0) | (labels[valid] >= logits.shape[-1])).any():
        raise ValueError("labels must be vocabulary IDs or -100")
    targets = labels[:, 1:]
    if not targets.ne(-100).any():
        raise ValueError("at least one shifted target must be supervised")
    return F.cross_entropy(
        logits[:, :-1].reshape(-1, logits.shape[-1]),
        targets.reshape(-1),
        ignore_index=-100,
    )


# endregion: shift


class ReducedBLIP2(nn.Module):
    """Two-stage teaching model; all modules start from random weights."""

    def __init__(
        self, width=24, lm_width=32, heads=4, queries=4, vocabulary=32
    ):
        super().__init__()
        if (
            min(width, lm_width, heads, queries) < 1
            or vocabulary < 5
            or width % heads
            or lm_width % heads
            or queries >= 128
        ):
            raise ValueError(
                "positive head-divisible widths, 1..127 queries, V>=5"
            )
        self.vocabulary = vocabulary
        self.query_count = queries
        self.vision = nn.Conv2d(3, width, kernel_size=4, stride=4)
        self.qformer = QFormer(width, heads, queries, vocabulary)
        self.image_head = nn.Linear(width, width)
        self.text_head = nn.Linear(width, width)
        self.match_head = nn.Linear(width, 2)
        self.generation_head = nn.Linear(width, vocabulary)
        self.temperature = nn.Parameter(torch.tensor(0.07))
        self.projection = nn.Linear(width, lm_width)
        self.lm = TinyCausalLM(lm_width, heads, vocabulary)
        self.freeze_backbones()

    # region: freeze
    def freeze_backbones(self):
        self.vision.requires_grad_(False)
        self.lm.requires_grad_(False)
        self.vision.eval()
        self.lm.eval()

    def encode_image(self, images):
        if (
            images.ndim != 4
            or images.shape[1] != 3
            or min(images.shape[-2:]) < 4
            or any(size % 4 for size in images.shape[-2:])
            or images.shape[0] == 0
            or not images.is_floating_point()
        ):
            raise ValueError("images must be floating [B>0, 3, 4h, 4w]")
        with torch.no_grad():
            return self.vision(images).flatten(2).transpose(1, 2)

    # endregion: freeze

    def train(self, mode=True):
        super().train(mode)
        self.vision.eval()
        self.lm.eval()
        return self

    # region: itc
    def contrastive_loss(self, image, tokens):
        hidden = self.qformer(image, tokens, mode="itc")
        query = F.normalize(
            self.image_head(hidden[:, : self.query_count]), dim=-1
        )
        text = F.normalize(
            self.text_head(hidden[:, self.query_count]), dim=-1
        )  # CLS position
        scores = torch.einsum("bqd,cd->bcq", query, text).amax(-1)
        scores = scores / self.temperature.clamp(0.001, 0.5)
        targets = torch.arange(image.shape[0], device=image.device)
        loss = (
            F.cross_entropy(scores, targets)
            + F.cross_entropy(scores.T, targets)
        ) / 2
        return loss, scores

    # endregion: itc

    # region: itm
    def matching_loss(self, image, tokens):
        batch = image.shape[0]
        if batch < 2:
            raise ValueError("ITM needs at least two distinct paired examples")
        # Deterministic cyclic negatives replace author hard-negative sampling.
        paired_images = torch.cat([image, image.roll(1, 0), image], 0)
        paired_text = torch.cat([tokens, tokens, tokens.roll(1, 0)], 0)
        hidden = self.qformer(paired_images, paired_text, mode="itm")
        logits = self.match_head(hidden[:, : self.query_count]).mean(1)
        labels = torch.cat(
            [
                torch.ones(batch, dtype=torch.long),
                torch.zeros(2 * batch, dtype=torch.long),
            ]
        ).to(image.device)
        return F.cross_entropy(logits, labels)

    # endregion: itm

    # region: itg
    def grounded_loss(self, image, tokens):
        decoder_tokens = tokens.clone()
        decoder_tokens[:, 0] = DEC  # Replace CLS with the decoding marker.
        hidden = self.qformer(image, decoder_tokens, mode="itg")
        # Queries are a bidirectional prefix, never predicted as text here.
        logits = self.generation_head(hidden[:, self.query_count :])
        labels = decoder_tokens.masked_fill(decoder_tokens.eq(PAD), -100)
        return shifted_loss(logits, labels)

    # endregion: itg

    def stage_one(self, images, tokens):
        """Joint unsmoothed ITC + ITM + ITG; batch pairs must be distinct."""
        validate_tokens(tokens, self.vocabulary, images.shape[0])
        if tokens.shape[1] < 2 or (tokens[:, 0] != CLS).any():
            raise ValueError("stage 1 needs CLS followed by supervised text")
        image = self.encode_image(images)
        itc, scores = self.contrastive_loss(image, tokens)
        itm = self.matching_loss(image, tokens)
        itg = self.grounded_loss(image, tokens)
        return {
            "loss": itc + itm + itg,
            "itc": itc,
            "itm": itm,
            "itg": itg,
            "similarities": scores,
        }

    # region: prefix
    def stage_two(self, images, tokens):
        validate_tokens(tokens, self.vocabulary, images.shape[0])
        image = self.encode_image(images)
        query = self.qformer(image)  # No text input to Q-Former in this stage.
        prefix = self.projection(query)
        inputs = torch.cat([prefix, self.lm.embedding(tokens)], dim=1)
        prefix_padding = torch.zeros(
            prefix.shape[:2], dtype=torch.bool, device=tokens.device
        )
        padding = torch.cat([prefix_padding, tokens.eq(PAD)], dim=1)
        ignored = torch.full(
            prefix.shape[:2], -100, dtype=torch.long, device=tokens.device
        )
        labels = torch.cat(
            [ignored, tokens.masked_fill(tokens.eq(PAD), -100)], 1
        )
        # Do NOT use no_grad here: the frozen LM must backpropagate to prefix.
        logits = self.lm(inputs, padding)
        return {
            "loss": shifted_loss(logits, labels),
            "logits": logits,
            "labels": labels,
            "prefix": prefix,
        }

    # endregion: prefix

    # region: inference
    @torch.no_grad()
    def generate_ids(self, images, max_new_tokens=4):
        """Greedily continue DEC; random small-vocabulary IDs are not captions."""
        if not 1 <= max_new_tokens < 128 - self.query_count:
            raise ValueError("generation length must fit the 128-position LM")
        prefix = self.projection(self.qformer(self.encode_image(images)))
        ids = torch.full(
            (images.shape[0], 1), DEC, dtype=torch.long, device=images.device
        )
        finished = torch.zeros(
            images.shape[0], dtype=torch.bool, device=images.device
        )
        for _ in range(max_new_tokens):
            inputs = torch.cat([prefix, self.lm.embedding(ids)], dim=1)
            padding = torch.cat(
                [
                    torch.zeros(
                        prefix.shape[:2], dtype=torch.bool, device=images.device
                    ),
                    ids.eq(PAD),
                ],
                1,
            )
            logits = self.lm(inputs, padding)[:, -1].clone()
            logits[:, [PAD, CLS, DEC]] = -torch.inf
            next_id = logits.argmax(-1).masked_fill(finished, PAD)
            ids = torch.cat([ids, next_id[:, None]], dim=1)
            finished = finished | next_id.eq(EOS)
            if finished.all():
                break
        return ids

    # endregion: inference


def gradient_norm(module):
    """Return the total norm of existing gradients, or zero if absent."""
    gradients = [
        parameter.grad.norm().square()
        for parameter in module.parameters()
        if parameter.grad is not None
    ]
    return float(torch.stack(gradients).sum().sqrt()) if gradients else 0.0


def smoke():
    """Execute sequential stage updates and export small, explicitly random evidence."""
    torch.set_num_threads(2)
    torch.manual_seed(17)
    model = ReducedBLIP2().train()
    images = torch.randn(3, 3, 16, 16)
    tokens = torch.tensor(
        [[CLS, 4, 5, EOS, PAD], [CLS, 6, 7, 8, EOS], [CLS, 9, EOS, PAD, PAD]]
    )
    optimizer = torch.optim.AdamW(
        (
            parameter
            for parameter in model.parameters()
            if parameter.requires_grad
        ),
        lr=1e-3,
    )
    initial_queries = model.qformer.queries.detach().clone()
    first = model.stage_one(images, tokens)
    first["loss"].backward()
    first_grad = gradient_norm(model.qformer)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    stage_two_tokens = tokens.clone()
    stage_two_tokens[:, 0] = DEC
    second = model.stage_two(images, stage_two_tokens)
    second["prefix"].retain_grad()
    second["loss"].backward()
    second_grad = gradient_norm(model.qformer)
    projection_grad = gradient_norm(model.projection)
    prefix_grad = float(second["prefix"].grad.norm())
    frozen_have_grad = any(
        parameter.grad is not None
        for module in (model.vision, model.lm)
        for parameter in module.parameters()
    )
    optimizer.step()
    return {
        "scope": "Reduced random-weight mechanics; no pretrained visual captions",
        "official_revision": "05cb670f47296522e24e7727e366b4a997d28582",
        "program_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "command": "python examples/blip2_training.py --smoke --output public/examples/blip2-demo/smoke.json",
        "environment": {
            "python": platform.python_version(),
            "torch": torch.__version__,
        },
        "seed": 17,
        "threads": 2,
        "architecture": {
            "batch": 3,
            "patches": 16,
            "queries": 4,
            "qformer_width": 24,
            "qformer_blocks": 4,
            "lm_width": 32,
            "lm_blocks": 2,
            "vocabulary": 32,
        },
        "stage_one": {
            key: float(first[key].detach())
            for key in ("itc", "itm", "itg", "loss")
        },
        "stage_two_loss": float(second["loss"].detach()),
        "stage_two_labels": second["labels"].tolist(),
        "gradients": {
            "stage_one_qformer": first_grad,
            "stage_two_qformer": second_grad,
            "stage_two_projection": projection_grad,
            "lm_input_prefix": prefix_grad,
            "frozen_parameters_have_grad": frozen_have_grad,
        },
        "query_update_norm": float(
            (model.qformer.queries.detach() - initial_queries).norm()
        ),
        "generated_token_ids": model.eval().generate_ids(images).tolist(),
    }


def main():
    """Parse an explicit offline smoke invocation."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    result = json.dumps(smoke(), indent=2) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(result, encoding="utf-8")
    print(result, end="")


if __name__ == "__main__":
    main()