"""Redes 3 y 4 de 4: una recurrente (LSTM) y una mini GPT (Transformer), entrenadas con tus textos el mismo tiempo.

Es la demo del vídeo «Cómo funcionan las redes neuronales por dentro: del perceptrón a GPT» (Frogames): allí
las dos se entrenaron 10 minutos con los 221 episodios del podcast (7,7 millones de caracteres) y empataron.
Uso:  uv run --with torch --with numpy python texto.py <fichero.txt o carpeta con .txt> [minutos_por_red]
Aprenden letra a letra a adivinar la siguiente. Cuanto más texto (mínimo, un par de millones de caracteres), mejor.
"""
import sys, time
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F

BLOCK, BATCH, LR = 256, 64, 1e-3      # letras de contexto, textos por lote, velocidad de aprendizaje
dev = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
torch.manual_seed(1337)

ruta, minutos = Path(sys.argv[1]), float(sys.argv[2]) if len(sys.argv) > 2 else 10
ficheros = sorted(ruta.glob("*.txt")) if ruta.is_dir() else [ruta]
texto = "\n".join(p.read_text() for p in ficheros)
chars = sorted(set(texto))
stoi = {c: i for i, c in enumerate(chars)}
V = len(chars)
data = torch.tensor([stoi[c] for c in texto], dtype=torch.long)
n = int(0.9 * len(data))
train, val = data[:n], data[n:]                                     # 90 % para aprender, 10 % para medir
i = texto.index("\n", n) + 1
PROMPT = texto[i:i + 35]                                            # inicio de las muestras: una línea que no ha visto
print(f"{len(texto):,} caracteres, {V} símbolos, en {dev}", flush=True)


def lote(d):
    ix = torch.randint(len(d) - BLOCK - 1, (BATCH,))
    x = torch.stack([d[i:i + BLOCK] for i in ix])
    y = torch.stack([d[i + 1:i + BLOCK + 1] for i in ix])          # la respuesta: la letra siguiente
    return x.to(dev), y.to(dev)


class Bloque(nn.Module):
    """Un bloque de Transformer: atención (cada letra mira a las anteriores) y una red densa."""

    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.drop = nn.Dropout(0.1)

    def atencion(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q, k, v = (z.view(B, T, self.h, C // self.h).transpose(1, 2) for z in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.1 if self.training else 0)
        return self.proj(y.transpose(1, 2).reshape(B, T, C))

    def forward(self, x):
        x = x + self.drop(self.atencion(self.ln1(x)))
        return x + self.drop(self.mlp(self.ln2(x)))


class MiniGPT(nn.Module):
    """6 bloques de Transformer apilados, 6 cabezas de atención, 384 dimensiones (~10,8 M de pesos)."""

    def __init__(self, d=384, capas=6, cabezas=6):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(BLOCK, d)
        self.bloques = nn.Sequential(*[Bloque(d, cabezas) for _ in range(capas)])
        self.ln, self.head = nn.LayerNorm(d), nn.Linear(d, V)

    def forward(self, x):
        h = self.tok(x) + self.pos(torch.arange(x.shape[1], device=x.device))
        return self.head(self.ln(self.bloques(h)))

    @torch.no_grad()
    def generar(self, ids, n, temp=0.8):
        for _ in range(n):
            p = F.softmax(self(ids[:, -BLOCK:])[:, -1] / temp, -1)
            ids = torch.cat([ids, torch.multinomial(p, 1)], 1)
        return ids


class Recurrente(nn.Module):
    """LSTM de 3 capas × 640: lee letra a letra arrastrando una memoria (~9,3 M de pesos)."""

    def __init__(self, d=384, oculto=640, capas=3):
        super().__init__()
        self.tok = nn.Embedding(V, d)
        self.lstm = nn.LSTM(d, oculto, capas, batch_first=True, dropout=0.1)
        self.head = nn.Linear(oculto, V)

    def forward(self, x):
        return self.head(self.lstm(self.tok(x))[0])

    @torch.no_grad()
    def generar(self, ids, n, temp=0.8):
        out, estado = self.lstm(self.tok(ids))
        for _ in range(n):
            p = F.softmax(self.head(out[:, -1]) / temp, -1)
            nxt = torch.multinomial(p, 1)
            ids = torch.cat([ids, nxt], 1)
            out, estado = self.lstm(self.tok(nxt), estado)
        return ids


@torch.no_grad()
def perdida_val(m):
    """Error medio (entropía cruzada) adivinando la letra siguiente en texto que no ha visto."""
    m.eval()
    g = torch.Generator().manual_seed(0)
    tot = 0
    for _ in range(20):
        ix = torch.randint(len(val) - BLOCK - 1, (BATCH,), generator=g)
        x = torch.stack([val[i:i + BLOCK] for i in ix]).to(dev)
        y = torch.stack([val[i + 1:i + BLOCK + 1] for i in ix]).to(dev)
        tot += F.cross_entropy(m(x).reshape(-1, V), y.reshape(-1)).item()
    m.train()
    return tot / 20


def entrenar(nombre, m):
    m.to(dev)
    print(f"\n== {nombre}: {sum(p.numel() for p in m.parameters()):,} pesos", flush=True)
    opt = torch.optim.AdamW(m.parameters(), lr=LR)
    paso, t0 = 0, time.time()
    while time.time() - t0 < minutos * 60:
        x, y = lote(train)
        loss = F.cross_entropy(m(x).reshape(-1, V), y.reshape(-1))
        opt.zero_grad(set_to_none=True)
        loss.backward()                                             # retropropagación
        opt.step()
        paso += 1
    m.eval()
    torch.manual_seed(7)
    ids = torch.tensor([[stoi[c] for c in PROMPT]], device=dev)
    muestra = "".join(chars[i] for i in m.generar(ids, 300)[0].tolist())
    print(f"{paso:,} pasos en {minutos:g} min · error de validación {perdida_val(m):.3f}\n{muestra}\n", flush=True)


entrenar("Recurrente (LSTM)", Recurrente())
entrenar("Mini GPT (Transformer)", MiniGPT())