# examples/bert.py
"""Offline reduced BERT pretraining mechanics, with no tokenizer or checkpoint.

Conventions: google-research/bert revision
 eedf5716ce1268e56f0a50264a88cafad334ac61, original token-level masking.
Tiny ordered documents replace corpus sampling. One SGD step replaces the
original Adam/weight-decay schedule; dropout is disabled in smoke via eval().
"""

import argparse
import hashlib
import json
import math
from pathlib import Path
import random

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

REVISION = "eedf5716ce1268e56f0a50264a88cafad334ac61"
PAD, CLS, SEP, MASK = 0, 1, 2, 3
VOCAB_SIZE = 16
# Each tuple is a document of already-tokenized, ordered sentences.
CORPUS = (((4, 5, 6), (7, 8)), ((9, 10), (11, 12, 13)))


def check_ids(ids, upper, name):
    """Reject noninteger or out-of-range token/index tensors."""
    if ids.dtype != torch.long or (ids < 0).any() or (ids >= upper).any():
        raise ValueError(f"{name} must be int64 IDs in [0, {upper}).")


# region: sentence_pair
def sentence_pair(
    document=0,
    sentence=0,
    is_random_next=False,
    random_document=1,
    random_sentence=0,
    max_length=10,
):
    """Pack an explicit corpus pair; label 0 means IsNext, 1 means NotNext."""
    if not isinstance(is_random_next, bool):
        raise ValueError("is_random_next must be boolean.")
    if not (0 <= document < len(CORPUS)):
        raise ValueError("Invalid document.")
    if not (0 <= sentence < len(CORPUS[document]) - 1):
        raise ValueError("Sentence A must have a successor.")
    a = CORPUS[document][sentence]
    if is_random_next:
        if (
            not (0 <= random_document < len(CORPUS))
            or random_document == document
        ):
            raise ValueError("Random B must come from another document.")
        if not (0 <= random_sentence < len(CORPUS[random_document])):
            raise ValueError("Invalid random sentence.")
        b = CORPUS[random_document][random_sentence]
    else:
        b = CORPUS[document][sentence + 1]
    tokens = [CLS, *a, SEP, *b, SEP]
    if max_length < len(tokens):
        raise ValueError("Pair does not fit; this fixture does not truncate.")
    segments = [0] * (len(a) + 2) + [1] * (len(b) + 1)
    padding = max_length - len(tokens)
    return {
        "ids": torch.tensor(tokens + [PAD] * padding),
        "segments": torch.tensor(segments + [0] * padding),
        "valid": torch.tensor([True] * len(tokens) + [False] * padding),
        "nsp_label": int(is_random_next),
    }


# endregion: sentence_pair


# region: select_and_corrupt
def select_and_corrupt(ids, valid, rng, probability=0.15, max_predictions=3):
    """Use original capped selection and two RNG draws for 80/10/10."""
    check_ids(ids, VOCAB_SIZE, "ids")
    if ids.ndim != 1 or valid.shape != ids.shape or valid.dtype != torch.bool:
        raise ValueError("Expected one sequence and a boolean validity mask.")
    length = int(valid.sum())
    if length == 0 or not valid[:length].all() or valid[length:].any():
        raise ValueError("Valid tokens must form a nonempty prefix.")
    if not 0 < probability <= 1 or max_predictions < 1:
        raise ValueError(
            "Probability and maximum predictions must be positive."
        )
    candidates = [i for i in range(length) if int(ids[i]) >= 4]
    if not candidates:
        raise ValueError("At least one eligible content token is required.")
    rng.shuffle(candidates)
    count = min(max_predictions, max(1, int(round(length * probability))))
    selected = candidates[:count]
    output = ids.clone()
    branches = {}
    for index in selected:
        if rng.random() < 0.8:
            output[index], branches[index] = MASK, "mask"
        elif rng.random() < 0.5:
            branches[index] = "unchanged"
        else:
            # Original replacement samples the ENTIRE vocabulary, specials too.
            output[index] = rng.randint(0, VOCAB_SIZE - 1)
            branches[index] = "random"
    selected.sort()
    slots = max_predictions - len(selected)
    positions = torch.tensor(selected + [0] * slots)
    targets = torch.tensor([int(ids[i]) for i in selected] + [PAD] * slots)
    weights = torch.tensor([1.0] * len(selected) + [0.0] * slots)
    return output, positions, targets, weights, [branches[i] for i in selected]


# endregion: select_and_corrupt


class ReducedBERT(nn.Module):
    """One original-style post-norm block and both pretraining heads."""

    def __init__(self, width=8, heads=2, max_length=10, dropout=0.1, seed=7):
        super().__init__()
        if width < 1 or heads < 1 or width % heads or max_length < 1:
            raise ValueError("Positive width divisible by heads is required.")
        if not 0 <= dropout < 1:
            raise ValueError("Dropout must lie in [0, 1).")
        self.width, self.heads, self.max_length = width, heads, max_length
        # Construction and initialization do not change the caller's RNG state.
        with torch.random.fork_rng(devices=[]):
            torch.manual_seed(seed)
            self.token = nn.Embedding(VOCAB_SIZE, width)
            self.position = nn.Embedding(max_length, width)
            self.segment = nn.Embedding(2, width)
            self.embedding_norm = nn.LayerNorm(width, eps=1e-12)
            self.q, self.k, self.v = [nn.Linear(width, width) for _ in range(3)]
            self.attention_output = nn.Linear(width, width)
            self.attention_norm = nn.LayerNorm(width, eps=1e-12)
            self.intermediate = nn.Linear(width, 4 * width)
            self.output = nn.Linear(4 * width, width)
            self.output_norm = nn.LayerNorm(width, eps=1e-12)
            self.mlm_dense = nn.Linear(width, width)
            self.mlm_norm = nn.LayerNorm(width, eps=1e-12)
            self.vocab_bias = nn.Parameter(torch.zeros(VOCAB_SIZE))
            self.pooler = nn.Linear(width, width)
            self.nsp = nn.Linear(width, 2)
            self.drop = nn.Dropout(dropout)
            for module in self.modules():
                if isinstance(module, (nn.Linear, nn.Embedding)):
                    nn.init.trunc_normal_(
                        module.weight, std=0.02, a=-0.04, b=0.04
                    )
                    if isinstance(module, nn.Linear):
                        nn.init.zeros_(module.bias)

    # region: input_embeddings
    def input_embeddings(self, ids, segments):
        """Sum learned token, position and segment embeddings, then LN/dropout."""
        if (
            ids.ndim != 2
            or min(ids.shape) < 1
            or ids.shape[1] > self.max_length
        ):
            raise ValueError(
                "Expected nonempty [B, N] IDs within position capacity."
            )
        if segments.shape != ids.shape:
            raise ValueError("Segment shape must match tokens.")
        check_ids(ids, VOCAB_SIZE, "ids")
        check_ids(segments, 2, "segments")
        positions = torch.arange(ids.shape[1], device=ids.device)
        x = self.token(ids) + self.position(positions) + self.segment(segments)
        # CLS uses its learned token row, just like other vocabulary entries.
        return self.drop(self.embedding_norm(x))

    # endregion: input_embeddings

    # region: bidirectional_encoder
    def encode(self, ids, segments, valid):
        """Attend to every valid key in both directions, then post-norm FFN."""
        x = self.input_embeddings(ids, segments)
        if valid.shape != ids.shape or valid.dtype != torch.bool:
            raise ValueError("Expected boolean validity [B, N].")
        if not valid.any(dim=1).all():
            raise ValueError("Every sequence needs at least one valid key.")
        batch, length, width = x.shape
        head_width = width // self.heads
        q, k, v = [
            projection(x)
            .reshape(batch, length, self.heads, head_width)
            .transpose(1, 2)
            for projection in (self.q, self.k, self.v)
        ]
        scores = q @ k.transpose(-1, -2) / math.sqrt(head_width)
        scores = scores + (~valid[:, None, None, :]).to(x.dtype) * -10000.0
        attention = torch.softmax(scores, dim=-1)  # No causal triangle.
        context = (self.drop(attention) @ v).transpose(1, 2).reshape(x.shape)
        z = self.attention_norm(x + self.drop(self.attention_output(context)))
        # Original modeling.gelu uses this tanh approximation, not exact erf.
        ffn = self.output(F.gelu(self.intermediate(z), approximate="tanh"))
        return self.output_norm(z + self.drop(ffn)), attention

    # endregion: bidirectional_encoder

    # region: mlm_head
    def mlm_logits(self, states):
        """Dense/GELU/LN, then tied token matrix and output-only vocab bias."""
        h = self.mlm_norm(F.gelu(self.mlm_dense(states), approximate="tanh"))
        return F.linear(h, self.token.weight, self.vocab_bias)

    # endregion: mlm_head

    # region: nsp_head
    def nsp_logits(self, encoded):
        """Pool final CLS through dense+tanh, then predict IsNext/NotNext."""
        pooled_cls = torch.tanh(self.pooler(encoded[:, 0]))
        return self.nsp(pooled_cls)

    # endregion: nsp_head

    def forward(self, ids, segments, valid, positions):
        """Return encoded states, selected MLM logits, NSP logits and attention."""
        encoded, attention = self.encode(ids, segments, valid)
        selected = gather_selected_states(encoded, positions)
        return (
            encoded,
            self.mlm_logits(selected),
            self.nsp_logits(encoded),
            attention,
        )


# region: gather_selected_states
def gather_selected_states(encoded, positions):
    """Gather [B, M, D]; padded slots remain present but have zero loss weight."""
    if encoded.ndim != 3 or positions.ndim != 2 or min(positions.shape) < 1:
        raise ValueError(
            "Expected encoded [B, N, D] and nonempty positions [B, M]."
        )
    if positions.shape[0] != encoded.shape[0]:
        raise ValueError("Position batch must match encoded batch.")
    check_ids(positions, encoded.shape[1], "positions")
    return encoded.gather(
        1, positions[..., None].expand(-1, -1, encoded.shape[-1])
    )


# endregion: gather_selected_states


def validate_objectives(logits, targets, weights, nsp_logits, labels):
    """Check objective shapes, labels, and binary selected-slot weights."""
    if logits.ndim != 3 or min(logits.shape) < 1:
        raise ValueError("MLM logits must be nonempty [B, M, V].")
    if targets.shape != logits.shape[:2] or weights.shape != targets.shape:
        raise ValueError("MLM targets and weights must match [B, M].")
    check_ids(targets, logits.shape[-1], "targets")
    if (
        not torch.isfinite(weights).all()
        or ((weights != 0) & (weights != 1)).any()
    ):
        raise ValueError("Weights must be finite zero/one values.")
    if nsp_logits.shape != (logits.shape[0], 2) or labels.shape != (
        logits.shape[0],
    ):
        raise ValueError("Expected NSP logits [B, 2] and labels [B].")
    check_ids(labels, 2, "NSP labels")
    if not torch.isfinite(logits).all() or not torch.isfinite(nsp_logits).all():
        raise ValueError("Logits must be finite.")


# region: joint_loss_and_update
def joint_loss(logits, targets, weights, nsp_logits, labels):
    """Normalize selected MLM slots and batch NSP separately, then add."""
    validate_objectives(logits, targets, weights, nsp_logits, labels)
    per_slot = (
        -F.log_softmax(logits, dim=-1)
        .gather(-1, targets[..., None])
        .squeeze(-1)
    )
    mlm_loss = (per_slot * weights).sum() / (weights.sum() + 1e-5)
    nsp_loss = F.cross_entropy(nsp_logits, labels, reduction="mean")
    return mlm_loss + nsp_loss, mlm_loss, nsp_loss


def joint_update(model, batch, learning_rate=0.05):
    """Take one illustrative SGD step; both objectives reach shared parameters."""
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("Learning rate must be finite and positive.")
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    optimizer.zero_grad(set_to_none=True)
    encoded, logits, nsp_logits, attention = model(
        batch["ids"], batch["segments"], batch["valid"], batch["positions"]
    )
    losses = joint_loss(
        logits, batch["targets"], batch["weights"], nsp_logits, batch["labels"]
    )
    losses[0].backward()
    norms = {name: float(p.grad.norm()) for name, p in model.named_parameters()}
    if not all(math.isfinite(value) for value in norms.values()):
        raise ValueError("Nonfinite gradients.")
    optimizer.step()
    return encoded, logits, nsp_logits, attention, losses, norms


# endregion: joint_loss_and_update


def fixture_batch(seed=7):
    """Create one consecutive and one declared cross-document negative pair."""
    rng = random.Random(seed)
    pairs = [sentence_pair(), sentence_pair(is_random_next=True)]
    corruptions = [
        select_and_corrupt(pair["ids"], pair["valid"], rng) for pair in pairs
    ]
    return {
        "ids": torch.stack([item[0] for item in corruptions]),
        "segments": torch.stack([pair["segments"] for pair in pairs]),
        "valid": torch.stack([pair["valid"] for pair in pairs]),
        "positions": torch.stack([item[1] for item in corruptions]),
        "targets": torch.stack([item[2] for item in corruptions]),
        "weights": torch.stack([item[3] for item in corruptions]),
        "labels": torch.tensor([pair["nsp_label"] for pair in pairs]),
        "branches": [item[4] for item in corruptions],
    }


def smoke(seed=7):
    """Measure finite CPU tensors, objective gradients and masking invariants."""
    model = ReducedBERT(seed=seed).eval()
    batch = fixture_batch(seed)
    with torch.no_grad():
        original = model.encode(
            batch["ids"], batch["segments"], batch["valid"]
        )[0]
        padded_ids = batch["ids"].clone()
        padded_ids[~batch["valid"]] = 15
        padded = model.encode(padded_ids, batch["segments"], batch["valid"])[0]
        right_ids = batch["ids"].clone()
        right_ids[:, 6] = 14
        right = model.encode(right_ids, batch["segments"], batch["valid"])[0]
        padding_delta = float((padded - original)[batch["valid"]].abs().max())
        right_delta = float((right[:, 0] - original[:, 0]).abs().max())
    before = model.token.weight.detach().clone()
    encoded, logits, nsp_logits, attention, losses, norms = joint_update(
        model, batch
    )
    result = {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "primary_revision": REVISION,
        "seed": seed,
        "device": "cpu",
        "boundary": "random reduced mechanics, no pretrained semantics or benchmarks",
        "reductions": {
            "layers": 1,
            "width": 8,
            "heads": 2,
            "ffn_width": 32,
            "vocabulary": VOCAB_SIZE,
            "optimizer": "one SGD step",
            "dropout": "eval disabled",
        },
        "shapes": {
            "input": list(batch["ids"].shape),
            "encoded": list(encoded.shape),
            "mlm": list(logits.shape),
            "nsp": list(nsp_logits.shape),
            "attention": list(attention.shape),
        },
        "nsp_labels": batch["labels"].tolist(),
        "selected_positions": batch["positions"].tolist(),
        "target_ids": batch["targets"].tolist(),
        "loss_weights": batch["weights"].tolist(),
        "branches": batch["branches"],
        "loss": float(losses[0].detach()),
        "mlm_loss": float(losses[1].detach()),
        "nsp_loss": float(losses[2].detach()),
        "gradient_norms": norms,
        "tied_token_update_norm": float(
            (model.token.weight - before).detach().norm()
        ),
        "valid_padding_delta": padding_delta,
        "cls_right_context_delta": right_delta,
        "padded_key_probability": float(
            attention.detach()
            .masked_select(
                (~batch["valid"])[:, None, None, :].expand_as(attention)
            )
            .abs()
            .max()
        ),
        "finite": bool(
            all(
                torch.isfinite(t).all()
                for t in (encoded, logits, nsp_logits, attention)
            )
        ),
    }
    return result


def main():
    """Print measurements only when explicitly run with --smoke."""
    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("Specify --smoke for the offline reference.")
    torch.set_num_threads(1)
    print(json.dumps(smoke(args.seed), sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()