import math, torch
import torch.nn as nn, torch.nn.functional as F

T       = 1024
N_LAYER = 24
N_HEAD  = 12
N_EMBD  = 768
VOCAB   = 50304   # gpt2 vocab padded to a multiple of 64

class Block(nn.Module):
    def __init__(s):
        super().__init__()
        s.ln1 = nn.LayerNorm(N_EMBD); s.ln2 = nn.LayerNorm(N_EMBD)
        s.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False)
        s.proj = nn.Linear(N_EMBD, N_EMBD, bias=False)
        s.fc = nn.Linear(N_EMBD, 4 * N_EMBD, bias=False)
        s.out = nn.Linear(4 * N_EMBD, N_EMBD, bias=False)
    def forward(s, x):
        B, L, C = x.shape
        q, k, v = s.qkv(s.ln1(x)).split(C, dim=2)
        q, k, v = [t.view(B, L, N_HEAD, C // N_HEAD).transpose(1, 2) for t in (q, k, v)]
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + s.proj(a.transpose(1, 2).reshape(B, L, C))
        return x + s.out(F.gelu(s.fc(s.ln2(x))))

class GPT(nn.Module):
    def __init__(s):
        super().__init__()
        s.wte = nn.Embedding(VOCAB, N_EMBD)
        s.wpe = nn.Embedding(T, N_EMBD)
        s.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
        s.lnf = nn.LayerNorm(N_EMBD)
        s.head = nn.Linear(N_EMBD, VOCAB, bias=False)
        s.head.weight = s.wte.weight  # tied
        s.apply(s._init)
        for n, p in s.named_parameters():
            if n.endswith('proj.weight') or n.endswith('out.weight'):
                nn.init.normal_(p, std=0.02 / math.sqrt(2 * N_LAYER))
    def _init(s, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)
    def forward(s, idx):
        x = s.wte(idx) + s.wpe(torch.arange(idx.size(1), device=idx.device))
        for b in s.blocks: x = b(x)
        return s.head(s.lnf(x))

def load_iau1(path='model.safetensors', device='cpu'):
    from safetensors.torch import load_file
    m = GPT()
    # head.weight is tied to wte.weight, so it is not stored in the file
    m.load_state_dict(load_file(path), strict=False)
    return m.to(device).eval()

@torch.no_grad()
def generate(model, prompt, n=80, temp=0.8, top_k=50, rep_penalty=1.2, no_repeat_ngram=3):
    import tiktoken
    enc = tiktoken.get_encoding('gpt2')
    dev = next(model.parameters()).device
    idx = torch.tensor([enc.encode_ordinary(prompt)], device=dev)
    for _ in range(n):
        logits = model(idx[:, -T:])[:, -1, :].float()
        if rep_penalty != 1.0:   # push down tokens that already appeared
            seen = torch.unique(idx[0])
            l = logits[0, seen]
            logits[0, seen] = torch.where(l > 0, l / rep_penalty, l * rep_penalty)
        if no_repeat_ngram and no_repeat_ngram >= 2 and idx.size(1) >= no_repeat_ngram:
            k = no_repeat_ngram
            toks = idx[0].tolist()
            prefix = toks[-(k - 1):]
            banned = [toks[i + k - 1] for i in range(len(toks) - k + 1) if toks[i:i + k - 1] == prefix]
            if banned: logits[0, banned] = -float('inf')   # never repeat an exact n-gram
        logits = logits / temp
        v, _ = torch.topk(logits, top_k)
        logits[logits < v[:, [-1]]] = -float('inf')
        idx = torch.cat([idx, torch.multinomial(F.softmax(logits, -1), 1)], dim=1)
    return enc.decode(idx[0].tolist())