data.py

"""Carga, preprocesado, split y augmentation del dataset Digit Recognizer."""
from __future__ import annotations

import math
from typing import Iterator

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F

import config as cfg


# --------------------------------------------------------------------------- #
# Carga
# --------------------------------------------------------------------------- #
def load_raw() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Lee los CSV y devuelve (X_train uint8 (N,28,28), y int64 (N,), X_test uint8)."""
    train_df = pd.read_csv(cfg.TRAIN_CSV)
    test_df = pd.read_csv(cfg.TEST_CSV)

    if train_df.isnull().to_numpy().any() or test_df.isnull().to_numpy().any():
        raise ValueError("Aparecieron nulos en los CSV; el EDA no encontro ninguno.")

    y = train_df["label"].to_numpy(dtype=np.int64)
    X = train_df.drop(columns="label").to_numpy(dtype=np.uint8)
    X_test = test_df.to_numpy(dtype=np.uint8)

    side = cfg.IMG_SIZE
    return (
        X.reshape(-1, side, side),
        y,
        X_test.reshape(-1, side, side),
    )


def normalize(X_uint8: np.ndarray) -> torch.Tensor:
    """uint8 (N,28,28) -> float32 (N,1,28,28) estandarizado con las constantes de train.

    float() explicito en las constantes: un escalar numpy promocionaria todo a
    float64, duplicando la RAM y rompiendo AMP mas adelante.
    """
    x = X_uint8.astype(np.float32) / 255.0
    x = (x - float(cfg.PIXEL_MEAN)) / float(cfg.PIXEL_STD)
    t = torch.from_numpy(x).unsqueeze(1)
    assert t.dtype == torch.float32, t.dtype
    return t


def stratified_split(
    X: torch.Tensor, y: torch.Tensor, val_split: float, seed: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Split estratificado por clase, sin depender de scikit-learn."""
    g = torch.Generator().manual_seed(seed)
    train_idx, val_idx = [], []
    for c in range(cfg.N_CLASSES):
        idx = torch.nonzero(y == c, as_tuple=True)[0]
        idx = idx[torch.randperm(len(idx), generator=g)]
        n_val = int(round(len(idx) * val_split))
        val_idx.append(idx[:n_val])
        train_idx.append(idx[n_val:])

    train_idx = torch.cat(train_idx)[torch.randperm(sum(len(i) for i in train_idx), generator=g)]
    val_idx = torch.cat(val_idx)
    return X[train_idx], y[train_idx], X[val_idx], y[val_idx]


def stratified_kfold(
    y: torch.Tensor, n_splits: int, seed: int
) -> list[tuple[torch.Tensor, torch.Tensor]]:
    """K folds estratificados. Devuelve [(train_idx, val_idx), ...].

    Cada muestra cae en el conjunto de validacion de exactamente un fold, asi que
    las predicciones out-of-fold cubren el dataset entero sin solapamiento.
    """
    g = torch.Generator().manual_seed(seed)
    fold_of = torch.empty(len(y), dtype=torch.long)

    for c in range(cfg.N_CLASSES):
        idx = torch.nonzero(y == c, as_tuple=True)[0]
        idx = idx[torch.randperm(len(idx), generator=g)]
        # reparte la clase en bloques contiguos: cada fold recibe la misma proporcion
        fold_of[idx] = torch.arange(len(idx)) % n_splits

    folds = []
    for k in range(n_splits):
        val_idx = torch.nonzero(fold_of == k, as_tuple=True)[0]
        train_idx = torch.nonzero(fold_of != k, as_tuple=True)[0]
        train_idx = train_idx[torch.randperm(len(train_idx), generator=g)]
        folds.append((train_idx, val_idx))
    return folds


def get_full_train() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Las 42.000 imagenes de train sin dividir, mas el test. Para validacion cruzada."""
    X_raw, y_raw, X_test_raw = load_raw()
    return normalize(X_raw), torch.from_numpy(y_raw.copy()), normalize(X_test_raw)


def get_datasets(
    val_split: float = cfg.VAL_SPLIT, seed: int = cfg.SEED, verbose: bool = True
):
    """Devuelve (X_tr, y_tr, X_va, y_va, X_test) listos para entrenar."""
    X_raw, y_raw, X_test_raw = load_raw()
    X = normalize(X_raw)
    y = torch.from_numpy(y_raw.copy())  # copy(): to_numpy() puede devolver un array de solo lectura
    X_test = normalize(X_test_raw)

    X_tr, y_tr, X_va, y_va = stratified_split(X, y, val_split, seed)

    if verbose:
        print(f"train    : {tuple(X_tr.shape)}")
        print(f"val      : {tuple(X_va.shape)}")
        print(f"test     : {tuple(X_test.shape)}")
        print(f"rango    : [{X_tr.min():.3f}, {X_tr.max():.3f}]  "
              f"mean={X_tr.mean():.4f} std={X_tr.std():.4f}")
    return X_tr, y_tr, X_va, y_va, X_test


# --------------------------------------------------------------------------- #
# Batching
# --------------------------------------------------------------------------- #
def iterate_batches(
    X: torch.Tensor, y: torch.Tensor, batch_size: int, shuffle: bool,
    generator: torch.Generator | None = None, drop_last: bool = False,
) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
    """Itera en batches indexando los tensores directamente.

    Sin DataLoader a proposito: el dataset entero cabe en RAM (~130 MB en float32),
    asi que los workers de Windows solo anadirian sobrecarga de 'spawn' por epoca.
    """
    n = len(X)
    order = torch.randperm(n, generator=generator) if shuffle else torch.arange(n)
    last = n - (n % batch_size) if drop_last else n
    for start in range(0, last, batch_size):
        idx = order[start:start + batch_size]
        yield X[idx], y[idx]


def n_batches(n_samples: int, batch_size: int, drop_last: bool = False) -> int:
    return n_samples // batch_size if drop_last else math.ceil(n_samples / batch_size)


# --------------------------------------------------------------------------- #
# Augmentation
# --------------------------------------------------------------------------- #
def random_affine(
    x: torch.Tensor,
    degrees: float = cfg.AUG_DEGREES,
    translate: float = cfg.AUG_TRANSLATE,
    scale: tuple[float, float] = cfg.AUG_SCALE,
    shear: float = cfg.AUG_SHEAR,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Afin aleatoria por muestra, vectorizada sobre el batch entero.

    Equivale a transforms.RandomAffine pero sin convertir a PIL imagen por imagen:
    construye una matriz 2x3 distinta para cada muestra y usa affine_grid +
    grid_sample. Cada imagen recibe su propia transformacion (no una comun al
    batch) y corre igual en CPU que en GPU.
    """
    n = x.size(0)
    dev, dtype = x.device, x.dtype

    def rand(lo: float, hi: float) -> torch.Tensor:
        r = torch.rand(n, generator=generator, dtype=dtype)
        return (r * (hi - lo) + lo).to(dev)

    ang = rand(-degrees, degrees) * math.pi / 180.0
    shr = rand(-shear, shear) * math.pi / 180.0
    sc = rand(scale[0], scale[1])
    tx = rand(-translate, translate) * 2.0  # coords normalizadas: [-1, 1]
    ty = rand(-translate, translate) * 2.0

    cos, sin, tan = torch.cos(ang), torch.sin(ang), torch.tan(shr)

    # A = (rotacion @ cizalla) / escala
    a11 = cos / sc
    a12 = (cos * tan - sin) / sc
    a21 = sin / sc
    a22 = (sin * tan + cos) / sc

    theta = torch.stack(
        [torch.stack([a11, a12, tx], dim=1),
         torch.stack([a21, a22, ty], dim=1)], dim=1
    )  # (N, 2, 3)

    grid = F.affine_grid(theta, list(x.shape), align_corners=False)

    # Ojo con el relleno: los datos vienen normalizados, asi que el fondo negro
    # (pixel 0) vale -MEAN/STD = -0.425, no 0. Con padding_mode="zeros" a secas,
    # las esquinas que quedan vacias al rotar se rellenarian con 0, que es un gris
    # medio (pixel ~33) y le dejaria al modelo un borde falso que aprender.
    # Se corrige desplazando el fondo a 0, muestreando, y deshaciendo el desplazamiento.
    bg = -float(cfg.PIXEL_MEAN) / float(cfg.PIXEL_STD)
    out = F.grid_sample(x - bg, grid, mode="bilinear", padding_mode="zeros",
                        align_corners=False)
    return out + bg