"""Red 2 de 4: dígitos escritos a mano (MNIST), red clásica (capas densas) contra red convolucional.

Es la demo del vídeo «Cómo funcionan las redes neuronales por dentro: del perceptrón a GPT» (Frogames).
Entrena las dos 3 épocas y las prueba con los dígitos tal cual y movidos 4 píxeles a la derecha.
Uso:  uv run --with torch --with torchvision python digitos.py      (descarga MNIST a ./mnist)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets

dev = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
torch.manual_seed(0)
tr = datasets.MNIST("mnist", train=True, download=True)
te = datasets.MNIST("mnist", train=False, download=True)
Xtr, ytr = tr.data.float().div(255).unsqueeze(1), tr.targets      # 60.000 imágenes de 28×28
Xte, yte = te.data.float().div(255).unsqueeze(1), te.targets      # 10.000 de prueba


class Densa(nn.Module):
    """Cada neurona mira todos los píxeles: 784 → 128 → 10."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))

    def forward(self, x):
        return self.net(x)


class Convolucional(nn.Module):
    """8 filtros de 5×5 que recorren la imagen, luego 16 más, y una capa final."""

    def __init__(self):
        super().__init__()
        self.c1, self.c2 = nn.Conv2d(1, 8, 5, padding=2), nn.Conv2d(8, 16, 5, padding=2)
        self.fc = nn.Linear(16 * 7 * 7, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.c1(x)), 2)                     # 28×28 → 14×14
        x = F.max_pool2d(F.relu(self.c2(x)), 2)                     # 14×14 → 7×7
        return self.fc(x.flatten(1))


@torch.no_grad()
def aciertos(m, X):
    m.eval()
    pred = torch.cat([m(X[i:i + 2000].to(dev)).argmax(1).cpu() for i in range(0, len(X), 2000)])
    m.train()
    return 100 * (pred == yte).float().mean().item()


def entrenar(m, epocas=3):
    m.to(dev)
    opt = torch.optim.Adam(m.parameters(), 1e-3)
    for _ in range(epocas):
        for idx in torch.randperm(len(Xtr)).split(128):
            loss = F.cross_entropy(m(Xtr[idx].to(dev)), ytr[idx].to(dev))
            opt.zero_grad(); loss.backward(); opt.step()          # retropropagación y ajuste de pesos


desplazado = torch.roll(Xte, shifts=(0, 4), dims=(2, 3))           # el mismo test, 4 píxeles a la derecha
for nombre, m in [("Densa", Densa()), ("Convolucional", Convolucional())]:
    entrenar(m)
    pesos = sum(p.numel() for p in m.parameters())
    print(f"{nombre}: {pesos:,} pesos · aciertos {aciertos(m, Xte):.2f} % · "
          f"con el dígito movido 4 píxeles {aciertos(m, desplazado):.2f} %")