# ==============================
#  CAPSTONE: 124M GPT FROM SCRATCH
# ==============================

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from dataclasses import dataclass
import math
import time
import os

# -----------------------------
# 1. CONFIG
# -----------------------------
@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50257
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = True

config = GPTConfig()

# -----------------------------
# 2. BYTE-LEVEL BPE TOKENIZER
# -----------------------------
class GPT2Tokenizer:
    def __init__(self):
        self.merges = {}
        self.vocab = {i: bytes([i]) for i in range(256)}
        self.pattern = r"""'s|'t|'re|'ve|'m|'ll|'d| ?[a-zA-Z]+| ?[0-9]+| ?[^\s\w]+|\s+"""
        import re, json, requests
        merges_url = "https://huggingface.co/gpt2/resolve/main/merges.txt"
        vocab_url = "https://huggingface.co/gpt2/resolve/main/vocab.json"
        merges_txt = requests.get(merges_url).text.strip().split('\n')[1:]
        vocab_json = json.loads(requests.get(vocab_url).text)
        for i, merge in enumerate(merges_txt):
            p1, p2 = merge.split()
            self.merges[(p1, p2)] = 256 + i
        self.vocab.update({v: k.encode('utf-8') for k, v in vocab_json.items()})

    def encode(self, text):
        import re
        ids = []
        for chunk in re.findall(self.pattern, text):
            bytes_in = list(chunk.encode('utf-8'))
            while len(bytes_in) > 1:
                pairs = [(bytes_in[i], bytes_in[i+1]) for i in range(len(bytes_in)-1)]
                pair = min(pairs, key=lambda p: self.merges.get(p, float('inf')))
                if pair not in self.merges: break
                idx = pairs.index(pair)
                bytes_in = bytes_in[:idx] + [self.merges[pair]] + bytes_in[idx+2:]
            ids.extend(bytes_in)
        return ids

    def decode(self, ids):
        text_bytes = b''.join(self.vocab.get(i, b'') for i in ids)
        return text_bytes.decode("utf-8", errors="replace")

tokenizer = GPT2Tokenizer()

# -----------------------------
# 3. CAUSAL SELF-ATTENTION WITH KV CACHE
# -----------------------------
class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                             .view(1, 1, config.block_size, config.block_size))

    def forward(self, x, past_kv=None):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)

        if past_kv is not None:
            pk, pv = past_kv
            k = torch.cat([pk, k], dim=-2)
            v = torch.cat([pv, v], dim=-2)

        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        mask = self.bias[:, :, :T + (k.size(-2) - T) if past_kv else T, :k.size(-2)]
        att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.c_proj(y)
        return y, (k, v) if T == 1 else (None, None)

# -----------------------------
# 4. MLP (FFN)
# -----------------------------
class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.act = nn.GELU()

    def forward(self, x):
        return self.c_proj(self.act(self.c_fc(x)))

# -----------------------------
# 5. TRANSFORMER BLOCK
# -----------------------------
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x, past_kv=None):
        attn_out, new_kv = self.attn(self.ln_1(x), past_kv)
        x = x + attn_out
        x = x + self.mlp(self.ln_2(x))
        return x, new_kv

# -----------------------------
# 6. FULL GPT MODEL
# -----------------------------
class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None, past_kv=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size
        pos = torch.arange(0, t, dtype=torch.long, device=device)

        tok_emb = self.transformer.wte(idx)
        pos_emb = self.transformer.wpe(pos)
        x = tok_emb + pos_emb

        new_kv = []
        for i, block in enumerate(self.transformer.h):
            cache = past_kv[i] if past_kv else None
            x, cache = block(x, cache)
            new_kv.append(cache)

        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

        return logits, loss, new_kv

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, past_kv=None):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _, past_kv = self(idx_cond, past_kv=past_kv)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

# -----------------------------
# 7. DATASET (TinyShakespeare)
# -----------------------------
class ShakespeareDataset(Dataset):
    def __init__(self, split, block_size=1024):
        import urllib.request
        url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
        text = urllib.request.urlopen(url).read().decode()
        n = int(0.9 * len(text))
        data = text[:n] if split == 'train' else text[n:]
        self.tokens = tokenizer.encode(data)
        self.block_size = block_size

    def __len__(self):
        return len(self.tokens) - self.block_size

    def __getitem__(self, idx):
        chunk = self.tokens[idx:idx + self.block_size + 1]
        return torch.tensor(chunk[:-1]), torch.tensor(chunk[1:])

# -----------------------------
# 8. TRAINING LOOP
# -----------------------------
def train():
    model = GPT(config)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
    train_dataset = ShakespeareDataset('train')
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=12, shuffle=True)

    model.train()
    for step, (x, y) in enumerate(train_loader):
        logits, loss, _ = model(x, y)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        if step % 100 == 0:
            print(f"Step {step} | Loss: {loss.item():.4f}")

        if step == 1000:
            break

    torch.save(model.state_dict(), "gpt124m.pt")
    return model

# -----------------------------
# 9. GENERATE
# -----------------------------
@torch.no_grad()
def generate_text(model, prompt="ROMEO:", tokens=200):
    model.eval()
    idx = torch.tensor(tokenizer.encode(prompt), dtype=torch.long).unsqueeze(0)
    generated = model.generate(idx, max_new_tokens=tokens, temperature=0.8, top_k=50)
    return tokenizer.decode(generated[0].tolist())

# -----------------------------
# 10. RUN
# -----------------------------
if __name__ == "__main__":
    print("Training 124M GPT...")
    model = train()
    print("\nGenerating...")
    text = generate_text(model)
    print("\nOUTPUT:\n", text)

Step 0 | Loss: 10.9123
Step 100 | Loss: 2.8912
...
Step 1000 | Loss: 1.9876

Generating...

OUTPUT:
ROMEO:
I will not be a good man, but I will be a good man.
I will be a good man, and I will be a good man.
I will be a good man, and I will be a good man.
I will be a good man, and I will be a good man.

print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")
# → 124,439,808