"""Funciones de apoyo del capítulo «Transferencia de aprendizaje y ajuste fino».

El notebook del capítulo importa este archivo con ``import apoyo_transferencia as apoyo``.
Aquí está el código largo o repetitivo (el bucle de entrenamiento, la evaluación, las
gráficas, el diagnóstico de errores y el guardado), para que el notebook muestre solo las
decisiones del experimento: qué red se usa, qué se congela, cuántas épocas y cómo se elige.

Funciones:

- ``parametros``: cuenta los parámetros de un modelo.
- ``inicializar``: da valores iniciales a los pesos de una red nueva.
- ``conjunto``: convierte arreglos de NumPy en un conjunto de datos de PyTorch.
- ``cargador``: entrega un conjunto de datos por lotes.
- ``evaluar``: calcula pérdida, exactitud y probabilidades sobre un conjunto.
- ``entrenar``: entrena un modelo y conserva la versión con menor pérdida de validación.
- ``curvas``: dibuja la pérdida y la exactitud por época.
- ``diagnostico``: informe por clase, matriz de confusión y galería de errores.
- ``guardar``: guarda el modelo con sus metadatos y comprueba que se recarga igual.

Las tres variables globales (``SEMILLA``, ``DEVICE`` y ``ARTEFACTOS``) tienen un valor por
defecto, pero el notebook las reemplaza al comenzar para que todo el capítulo use la misma
semilla, el mismo dispositivo y la misma carpeta de resultados.

Si modificas este archivo con el notebook abierto, reinicia el kernel antes de volver a
ejecutar: Python conserva en memoria la versión que importó la primera vez.
"""

from pathlib import Path
import json, time, copy, hashlib, platform, gzip, struct, urllib.request
import numpy as np
import pandas as pd
import sklearn
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix, f1_score
import matplotlib.pyplot as plt
from PIL import Image
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset, Dataset

# Valores por defecto; el notebook los reemplaza en su primera celda.
SEMILLA = 42                                                           # semilla para el orden de los lotes
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # GPU si existe; si no, CPU
ARTEFACTOS = Path('artefactos_pytorch')                                # carpeta donde se guardan resultados


def parametros(modelo):
    """Cuenta todos los parámetros de un modelo (entrenables y congelados).

    Ejemplo: una capa ``nn.Linear(1280, 5)`` tiene 1280 × 5 pesos + 5 sesgos = 6405.
    """
    return sum(p.numel() for p in modelo.parameters())


def inicializar(modelo):
    """Da valores iniciales a los pesos de una red nueva y la mueve al dispositivo.

    Se usa solo en la estrategia A, la red que aprende desde cero.

    - Capas convolucionales y lineales: inicialización de He (``kaiming_normal_``),
      pensada para capas seguidas de ReLU. Los sesgos empiezan en cero.
    - Última capa lineal (la que produce las puntuaciones de las clases): se
      reinicializa con Xavier (``xavier_uniform_``), porque no la sigue una ReLU.

    Devuelve el mismo modelo, ya en ``DEVICE``.
    """
    for m in modelo.modules():
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
            if m.bias is not None:
                nn.init.zeros_(m.bias)
    lineales = [m for m in modelo.modules() if isinstance(m, nn.Linear)]
    if lineales:
        nn.init.xavier_uniform_(lineales[-1].weight)
    return modelo.to(DEVICE)


def conjunto(x, y):
    """Convierte imágenes ``x`` y etiquetas ``y`` de NumPy en un ``TensorDataset``.

    Las imágenes quedan como ``float32`` y las etiquetas como enteros ``long``, que es
    lo que espera la entropía cruzada de PyTorch.
    """
    return TensorDataset(torch.as_tensor(np.ascontiguousarray(x), dtype=torch.float32),
                         torch.as_tensor(np.asarray(y), dtype=torch.long))


def cargador(ds, lote=128, mezclar=False):
    """Entrega el conjunto ``ds`` en lotes de ``lote`` filas.

    Con ``mezclar=True`` el orden cambia en cada época (se usa al entrenar). El generador
    con ``SEMILLA`` hace que ese orden sea el mismo en cada ejecución.
    """
    return DataLoader(ds, batch_size=lote, shuffle=mezclar, num_workers=0,
                      generator=torch.Generator().manual_seed(SEMILLA))


@torch.inference_mode()
def evaluar(modelo, ds, lote=128):
    """Mide un modelo sobre un conjunto sin modificar sus pesos.

    Pone el modelo en modo de evaluación (``eval()``) y no construye el cálculo de
    gradientes (``inference_mode``). Devuelve un diccionario con:

    - ``ce``: entropía cruzada media (menor es mejor).
    - ``exactitud``: proporción de imágenes clasificadas correctamente.
    - ``logits``: puntuaciones de salida, una fila por imagen y una columna por clase.
    - ``y``: etiquetas reales.
    - ``probabilidades``: los logits convertidos con softmax (cada fila suma 1).
    """
    modelo.eval()
    perdidas = 0.
    aciertos = 0
    logits = []
    etiquetas = []
    for x, y in cargador(ds, lote):
        x = x.to(DEVICE)
        y = y.to(DEVICE)
        z = modelo(x)
        # Suma (no media) para poder dividir al final por el total de imágenes.
        perdidas += F.cross_entropy(z, y, reduction='sum').item()
        aciertos += (z.argmax(1) == y).sum().item()
        logits.append(z.cpu())
        etiquetas.append(y.cpu())
    z = torch.cat(logits)
    y = torch.cat(etiquetas)
    assert torch.isfinite(z).all(), 'Logits no finitos'
    return {'ce': perdidas / len(ds), 'exactitud': aciertos / len(ds),
            'logits': z.numpy(), 'y': y.numpy(), 'probabilidades': z.softmax(1).numpy()}


def entrenar(modelo, train, val, epocas, lote=128, lr=1e-3, paciencia=None, nombre='modelo',
             modo_train=None, incluir_inicial=False):
    """Entrena ``modelo`` y le deja cargada la versión con menor pérdida de validación.

    Parámetros:

    - ``train``, ``val``: conjuntos de entrenamiento y de validación. La prueba nunca
      entra aquí.
    - ``epocas``: número máximo de épocas.
    - ``lote``: imágenes por lote.
    - ``lr``: tasa de aprendizaje del optimizador Adam.
    - ``paciencia``: si la pérdida de validación no mejora durante este número de épocas
      seguidas, el entrenamiento se detiene antes del máximo. ``None`` la desactiva.
    - ``nombre``: nombre de la estrategia; se usa en los mensajes y en los archivos.
    - ``modo_train``: función opcional que se aplica al modelo después de ``train()``.
      En la estrategia C mantiene BatchNorm en modo de evaluación.
    - ``incluir_inicial``: si es ``True``, el modelo de partida cuenta como «época 0» y
      puede resultar elegido si ninguna época lo mejora (se usa en C, que parte de B).

    En cada época ejecuta el bucle de cinco pasos del Apéndice A (borrar gradientes,
    predecir, medir la pérdida, calcular gradientes y corregir) sobre todos los lotes.
    Al cerrar la época mide entrenamiento y validación en modo de evaluación, para que
    las dos curvas sean comparables.

    Guarda el historial en ``<nombre>-historial.csv`` y un resumen en
    ``<nombre>-entrenamiento.json``. Devuelve ``(historial, resumen)``: el historial es una
    tabla con una fila por época; el resumen incluye la mejor época, su pérdida de
    validación, el tiempo y los parámetros totales y entrenables.
    """
    # El optimizador se crea después de decidir qué parámetros se entrenan
    # (requires_grad): solo recibe esos.
    optimizador = torch.optim.Adam([p for p in modelo.parameters() if p.requires_grad], lr=lr)
    dl = cargador(train, lote, True)
    historial = []
    mejor = float('inf')    # menor pérdida de validación vista hasta ahora
    pesos = None            # copia del estado del modelo en esa mejor época
    espera = 0              # épocas seguidas sin mejorar (para la paciencia)
    inicio = time.perf_counter()
    epoca_mejor = None
    if incluir_inicial:
        inicial = evaluar(modelo, val, lote)
        mejor = inicial['ce']
        pesos = copy.deepcopy(modelo.state_dict())
        epoca_mejor = 0
    for epoca in range(1, epocas + 1):
        modelo.train()
        if modo_train is not None:
            modo_train(modelo)
        suma = 0.
        aciertos = 0
        for x, y in dl:
            x = x.to(DEVICE)
            y = y.to(DEVICE)
            optimizador.zero_grad(set_to_none=True)   # 1. borrar gradientes anteriores
            z = modelo(x)                             # 2. predecir (logits)
            perdida = F.cross_entropy(z, y)           # 3. medir el error
            if not torch.isfinite(perdida):
                raise FloatingPointError(nombre + ' pérdida no finita')
            perdida.backward()                        # 4. calcular gradientes
            optimizador.step()                        # 5. corregir los parámetros
            suma += perdida.item() * len(y)
            aciertos += (z.argmax(1) == y).sum().item()
        # Medir entrenamiento y validación en el MISMO modo (eval) al cierre de la época.
        tr = evaluar(modelo, train, lote)
        va = evaluar(modelo, val, lote)
        historial.append({'epoca': epoca, 'ce_online': suma / len(train),
                          'acc_online': aciertos / len(train), 'ce_train': tr['ce'], 'ce_val': va['ce'],
                          'acc_train': tr['exactitud'], 'acc_val': va['exactitud']})
        if va['ce'] < mejor:
            # Nueva mejor versión: se guarda una copia y se reinicia la cuenta de paciencia.
            mejor = va['ce']
            pesos = copy.deepcopy(modelo.state_dict())
            epoca_mejor = epoca
            espera = 0
        else:
            espera += 1
        print(f"{nombre} {epoca}/{epocas}: CE val={va['ce']:.4f}, acc val={va['exactitud']:.4f}", flush=True)
        if paciencia is not None and espera >= paciencia:
            break
    # Al terminar, el modelo vuelve a la mejor versión, no a la última época.
    modelo.load_state_dict(pesos)
    modelo.eval()
    h = pd.DataFrame(historial)
    h.to_csv(ARTEFACTOS / f'{nombre}-historial.csv', index=False)
    resumen = {'modelo': nombre, 'mejor_epoca': epoca_mejor, 'epocas': len(h), 'ce_val': mejor,
               'segundos': time.perf_counter() - inicio, 'parametros': parametros(modelo),
               'entrenables': sum(p.numel() for p in modelo.parameters() if p.requires_grad)}
    (ARTEFACTOS / f'{nombre}-entrenamiento.json').write_text(json.dumps(resumen, indent=2), encoding='utf-8')
    return h, resumen


def curvas(h, titulo, mejor=None):
    """Dibuja dos gráficas por época: pérdida (entropía cruzada) y exactitud.

    ``h`` es el historial que devuelve ``entrenar``. Entrenamiento se dibuja con línea
    continua y círculos; validación, con línea discontinua y cuadrados, para que se
    distingan también sin color. Si se indica ``mejor``, una línea vertical marca la
    época elegida.
    """
    fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))
    for ax, a, b, t in [(axes[0], 'ce_train', 'ce_val', 'Entropía cruzada'),
                        (axes[1], 'acc_train', 'acc_val', 'Exactitud')]:
        ax.plot(h.epoca, h[a], '-o', color='#1786B3', label='Entrenamiento (evaluación)', ms=4)
        ax.plot(h.epoca, h[b], '--s', color='#EE5400', label='Validación', ms=4)
        if mejor is not None:
            ax.axvline(mejor, color='.4', ls='--', label=f'Checkpoint {mejor}')
        ax.set(xlabel='Época', ylabel=t)
        ax.legend(fontsize=8)
    fig.suptitle(titulo)
    fig.tight_layout()
    plt.show()


def diagnostico(resultado, nombres, imagenes=None, titulo='Prueba'):
    """Muestra en qué se equivoca un modelo ya elegido.

    ``resultado`` es lo que devuelve ``evaluar``; ``nombres``, la lista de clases.
    Imprime precisión, sensibilidad (recall) y F1 por clase, y dibuja la matriz de
    confusión: cada fila es la clase real y cada columna la predicha; el número de cada
    celda es la cantidad de imágenes y el color, la proporción dentro de la fila.

    Si se pasan ``imagenes``, muestra además los ocho errores con mayor probabilidad
    predicha. Devuelve la matriz de confusión.

    Sirve para analizar errores; no debe usarse para volver a elegir la estrategia.
    """
    y = resultado['y']
    p = resultado['probabilidades']
    pred = p.argmax(1)
    k = len(nombres)
    print(classification_report(y, pred, labels=np.arange(k), target_names=nombres, digits=4, zero_division=0))
    cm = confusion_matrix(y, pred, labels=np.arange(k))
    normal = cm / np.maximum(cm.sum(1, keepdims=True), 1)   # proporción por fila (clase real)
    fig, ax = plt.subplots(figsize=(7, 5.5))
    im = ax.imshow(normal, vmin=0, vmax=1, cmap='Blues')
    ax.set_xticks(range(k), nombres, rotation=45, ha='right', fontsize=10)
    ax.set_yticks(range(k), nombres, fontsize=10)
    ax.set(xlabel='Predicho', ylabel='Real', title=titulo + ' · proporción por clase')
    for i in range(k):
        for j in range(k):
            ax.text(j, i, str(cm[i, j]), ha='center', va='center', fontsize=11,
                    color='white' if normal[i, j] > .5 else 'black')
    ax.grid(False)
    fig.colorbar(im, ax=ax, label='Proporción de la clase real')
    fig.tight_layout()
    plt.show()
    err = np.flatnonzero(y != pred)   # posiciones de las imágenes mal clasificadas
    if imagenes is not None and len(err):
        confianza = p.max(1)
        orden = err[np.argsort(-confianza[err])][:8]   # los 8 errores más «seguros»
        fig, axes = plt.subplots(2, 4, figsize=(11, 5))
        for ax, idx in zip(axes.flat, orden):
            img = np.asarray(imagenes[idx])
            # Las imágenes vienen como (canales, alto, ancho); matplotlib espera (alto, ancho, canales).
            img = np.moveaxis(img, 0, -1) if img.ndim == 3 and img.shape[0] in [1, 3] else img
            ax.imshow(img.squeeze(), cmap='gray' if img.squeeze().ndim == 2 else None)
            ax.set_title(f'Real: {nombres[y[idx]]}\nPred: {nombres[pred[idx]]} · p={confianza[idx]:.2f}', fontsize=8)
            ax.axis('off')
        for ax in list(axes.flat)[len(orden):]:
            ax.axis('off')
        fig.suptitle('Errores de mayor probabilidad predicha · no es una medida de calibración')
        fig.tight_layout()
        plt.show()
    return cm


def guardar(modelo, fabrica, nombre, metadatos, val):
    """Guarda el modelo elegido y comprueba que se puede volver a cargar.

    - ``fabrica``: función sin argumentos que crea una red vacía con la misma
      estructura; se usa para recargar el archivo.
    - ``metadatos``: información para usar el modelo después (clases, tamaño de la
      imagen, normalización, arquitectura).
    - ``val``: conjunto de validación con el que se comprueba la recarga.

    Crea ``<nombre>.pt`` con el estado del modelo y los metadatos, lo recarga en una red
    nueva y verifica que ambas dan las mismas puntuaciones en un lote de validación.
    Después escribe ``<nombre>-metadatos.json`` con los metadatos, la huella SHA-256 y el
    tamaño del archivo. Devuelve el modelo recargado.
    """
    ruta = ARTEFACTOS / f'{nombre}.pt'
    estado = {k: v.detach().cpu().clone() for k, v in modelo.state_dict().items()}
    torch.save({'state_dict': estado, 'metadatos': metadatos, 'torch': str(torch.__version__)}, ruta)
    cargado = fabrica().to(DEVICE)
    paquete = torch.load(ruta, map_location=DEVICE, weights_only=True)
    cargado.load_state_dict(paquete['state_dict'])
    cargado.eval()
    # Comprobar la recarga con validación; la prueba nunca se usa para elegir.
    muestra = next(iter(cargador(val, 32)))[0].to(DEVICE)
    modelo.eval()
    with torch.inference_mode():
        torch.testing.assert_close(modelo(muestra), cargado(muestra), rtol=1e-5, atol=1e-6)
    metadatos = dict(metadatos, sha256=hashlib.sha256(ruta.read_bytes()).hexdigest(), bytes=ruta.stat().st_size)
    (ARTEFACTOS / f'{nombre}-metadatos.json').write_text(json.dumps(metadatos, indent=2, ensure_ascii=False),
                                                         encoding='utf-8')
    print('Recarga verificada:', ruta, 'bytes:', ruta.stat().st_size)
    return cargado