"""Train a tiny character language model with causal attention. Requires torch."""
from __future__ import annotations
import math
import torch
from torch import nn
from torch.nn import functional as F


TEXT = "the cat sat.\nthe dog ran.\n"*80


class TinyLM(nn.Module):
    """Pre: token IDs (B,T), T<=16. Post: causal logits (B,T,V), with learned positions."""
    def __init__(self, vocabulary, width=32, heads=2, context=16):
        super().__init__()
        if min(vocabulary, width, heads, context) <= 0 or width % heads:
            raise ValueError("positive dimensions and width divisible by heads required")
        self.width, self.heads, self.context = width, heads, context
        self.token = nn.Embedding(vocabulary, width)
        self.position = nn.Embedding(context, width)
        self.norm1, self.norm2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.qkv, self.output = nn.Linear(width, 3*width), nn.Linear(width, width)
        self.ffn = nn.Sequential(nn.Linear(width, 2*width), nn.GELU(), nn.Linear(2*width, width))
        self.head = nn.Linear(width, vocabulary)

    def forward(self, tokens):
        b, t = tokens.shape
        if t > self.context:
            raise ValueError("sequence exceeds allocated position embeddings")
        x = self.token(tokens) + self.position(torch.arange(t, device=tokens.device))
        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)
        q, k, v = [value.reshape(b, t, self.heads, self.width//self.heads).transpose(1, 2)
                   for value in (q, k, v)]
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.width//self.heads)
        mask = torch.ones(t, t, dtype=torch.bool, device=tokens.device).triu(1)
        attention = scores.masked_fill(mask, -torch.inf).softmax(-1)
        mixed = (attention @ v).transpose(1, 2).contiguous().reshape(b, t, self.width)
        x = x + self.output(mixed)
        x = x + self.ffn(self.norm2(x))
        return self.head(x)


def train(*, steps=120):
    """Pre: steps>0. Post: model fitted to repeated demonstration text, not general English."""
    if steps <= 0:
        raise ValueError("steps must be positive")
    torch.set_num_threads(1)
    torch.manual_seed(7)
    alphabet = sorted(set(TEXT))
    ids = torch.tensor([alphabet.index(c) for c in TEXT], dtype=torch.long)
    model = TinyLM(len(alphabet))
    optimizer = torch.optim.AdamW(model.parameters(), lr=.01)
    history = []
    # Invariant: targets are inputs shifted by one; previous losses record completed updates.
    # Bound: steps-step.
    for _ in range(steps):
        starts = torch.randint(0, len(ids)-17, (16,))
        inputs = torch.stack([ids[start:start+16] for start in starts])
        targets = torch.stack([ids[start+1:start+17] for start in starts])
        optimizer.zero_grad(set_to_none=True)
        loss = F.cross_entropy(model(inputs).reshape(-1, len(alphabet)), targets.reshape(-1))
        loss.backward()
        optimizer.step()
        history.append(loss.item())
    return model, alphabet, history


def generate(model, alphabet, prefix="the ", count=40):
    """Pre: prefix characters belong to vocabulary. Post: greedy autoregressive continuation."""
    tokens = torch.tensor([[alphabet.index(c) for c in prefix]], dtype=torch.long)
    model.eval()
    with torch.no_grad():
        # Invariant: tokens contains prefix and all generated IDs. Bound: remaining count.
        for _ in range(count):
            next_id = model(tokens[:, -model.context:])[:, -1].argmax(-1, keepdim=True)
            tokens = torch.cat((tokens, next_id), dim=1)
    return "".join(alphabet[index] for index in tokens[0].tolist())


def demo():
    model, alphabet, history = train()
    print("Vocabulary:", repr("".join(alphabet)))
    print("Initial/final next-character loss:", history[0], history[-1])
    print("Generated:", repr(generate(model, alphabet)))
    print("The small corpus demonstrates teacher forcing and causality; it does not establish general language ability.")


def self_test():
    model, alphabet, history = train()
    check(history[-1] < history[0], True)
    check(history[-1] < 1, True)
    model.eval()
    first = torch.zeros(1, 8, dtype=torch.long)
    second = first.clone()
    second[:, 4:] = 1
    with torch.no_grad():
        a, b = model(first), model(second)
    check(torch.allclose(a[:, :4], b[:, :4], atol=1e-6), True)
    check(a.shape, torch.Size([1, 8, len(alphabet)]))
    check(len(generate(model, alphabet, count=10)), 14)
    return 5

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()