import numpy as np

np.random.seed(42)

n_neurons = 256
n_layers = 4

# Simulamos la propagación de varianza con distintas inicializaciones
print("=" * 65)
print("Propagación de varianza a través de 4 capas (256 neuronas/capa)")
print("=" * 65)

# Señal de entrada simulada (por ejemplo, píxeles normalizados)
x = np.random.randn(1000, n_neurons) * 0.5  # Var ≈ 0.25

strategies = {
    "Normal(0,1)":   lambda n_in: np.random.randn(n_in, n_in) * 1.0,
    "Uniform[-1,1]": lambda n_in: np.random.uniform(-1, 1, (n_in, n_in)),
    "Xavier":        lambda n_in: np.random.randn(n_in, n_in) * np.sqrt(2.0 / (n_in + n_in)),
    "He/Kaiming":    lambda n_in: np.random.randn(n_in, n_in) * np.sqrt(2.0 / n_in),
}

for name, weight_fn in strategies.items():
    a = x.copy()
    print(f"\n📊 {name}:")
    print(f"   Entrada  → Var = {np.var(a):.6f}")
    for l in range(n_layers):
        W = weight_fn(n_neurons)
        z = a @ W.T
        a = np.maximum(0, z)  # ReLU
        var = np.var(a)
        status = "✅" if 0.01 < var < 100 else ("💥 EXPLOTA" if var > 100 else "📉 Se desvanece")
        print(f"   Capa {l+1}   → Var = {var:.6f}  {status}")

print("\n" + "=" * 65)
print("💡 Solo He/Kaiming mantiene la varianza estable con ReLU")
print("=" * 65)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import copy
import warnings
warnings.filterwarnings('ignore')

# Reproducibilidad
torch.manual_seed(42)
np.random.seed(42)

# Dispositivo
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Dispositivo: {device}")
print(f"PyTorch: {torch.__version__}")

# Transformación: normalizar a [0,1] y aplanar a vector de 784
transform = transforms.Compose([
    transforms.ToTensor(),  # [0,255] → [0,1]
])

# Descargar datasets
train_dataset = datasets.FashionMNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.FashionMNIST(
    root='./data', train=False, download=True, transform=transform
)

# DataLoaders
BATCH_SIZE = 256
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# Clases
class_names = ['Camiseta', 'Pantalón', 'Jersey', 'Vestido', 'Abrigo',
               'Sandalia', 'Camisa', 'Zapatilla', 'Bolso', 'Bota']

print(f"Train: {len(train_dataset)} imágenes")
print(f"Test:  {len(test_dataset)} imágenes")
print(f"Clases: {len(class_names)}")
print(f"Forma de una imagen: {train_dataset[0][0].shape}")

fig, axes = plt.subplots(2, 8, figsize=(14, 4))
for i, ax in enumerate(axes.flat):
    img, label = train_dataset[i]
    ax.imshow(img.squeeze(), cmap='gray')
    ax.set_title(class_names[label], fontsize=8)
    ax.axis('off')
plt.suptitle('Muestras de FashionMNIST', fontsize=12, fontweight='bold')
plt.tight_layout()
plt.show()

class MLP(nn.Module):
    """MLP con capas ocultas y captura de activaciones."""

    def __init__(self, input_dim=784, hidden_dim=256, output_dim=10, n_hidden=4):
        super().__init__()

        layers = []
        # Primera capa oculta
        layers.append(nn.Linear(input_dim, hidden_dim))
        layers.append(nn.ReLU())

        # Capas ocultas intermedias
        for _ in range(n_hidden - 1):
            layers.append(nn.Linear(hidden_dim, hidden_dim))
            layers.append(nn.ReLU())

        # Capa de salida (sin activación — se usa CrossEntropyLoss)
        layers.append(nn.Linear(hidden_dim, output_dim))

        self.network = nn.Sequential(*layers)

        # Almacén de activaciones para análisis
        self.activations = {}
        self._register_hooks()

    def _register_hooks(self):
        """Registra hooks para capturar activaciones de cada capa ReLU."""
        layer_idx = 0
        for module in self.network:
            if isinstance(module, nn.ReLU):
                name = f'relu_{layer_idx}'
                module.register_forward_hook(self._make_hook(name))
                layer_idx += 1

    def _make_hook(self, name):
        def hook(module, input, output):
            self.activations[name] = output.detach().cpu()
        return hook

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Aplanar: [B, 1, 28, 28] → [B, 784]
        return self.network(x)


# Verificar arquitectura
model_test = MLP()
print(model_test)
n_params = sum(p.numel() for p in model_test.parameters())
print(f"\nTotal parámetros: {n_params:,}")

def init_zeros(model):
    """Inicializa todos los pesos a cero."""
    for name, param in model.named_parameters():
        nn.init.zeros_(param)
    return model


def init_ones(model):
    """Inicializa todos los pesos a uno."""
    for name, param in model.named_parameters():
        nn.init.ones_(param)
    return model


def init_normal(model):
    """Inicializa pesos con N(0, 1) — varianza demasiado grande."""
    for name, param in model.named_parameters():
        if 'weight' in name:
            nn.init.normal_(param, mean=0.0, std=1.0)
        elif 'bias' in name:
            nn.init.zeros_(param)
    return model


def init_uniform(model):
    """Inicializa pesos con U(-1, 1) — varianza no adaptada."""
    for name, param in model.named_parameters():
        if 'weight' in name:
            nn.init.uniform_(param, a=-1.0, b=1.0)
        elif 'bias' in name:
            nn.init.zeros_(param)
    return model


def init_xavier(model):
    """Inicialización Xavier/Glorot Normal — ideal para tanh/sigmoid."""
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:
            nn.init.xavier_normal_(param)
        elif 'bias' in name:
            nn.init.zeros_(param)
    return model


def init_he(model):
    """Inicialización He/Kaiming Normal — ideal para ReLU."""
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:
            nn.init.kaiming_normal_(param, mode='fan_in', nonlinearity='relu')
        elif 'bias' in name:
            nn.init.zeros_(param)
    return model


# Diccionario con todas las estrategias
INIT_STRATEGIES = {
    'Ceros':        init_zeros,
    'Unos':         init_ones,
    'Normal(0,1)':  init_normal,
    'Uniform[-1,1]':init_uniform,
    'Xavier':       init_xavier,
    'He/Kaiming':   init_he,
}

print(f"Estrategias definidas: {list(INIT_STRATEGIES.keys())}")

def capture_activations(model, data_loader):
    """Pasa un batch por el modelo y devuelve las activaciones capturadas."""
    model.eval()
    images, _ = next(iter(data_loader))
    images = images.to(device)
    with torch.no_grad():
        _ = model(images)
    return {k: v.numpy() for k, v in model.activations.items()}


# Capturar activaciones para cada inicialización
fig, axes = plt.subplots(len(INIT_STRATEGIES), 4, figsize=(16, 14))

for row, (name, init_fn) in enumerate(INIT_STRATEGIES.items()):
    # Crear modelo fresco e inicializar
    torch.manual_seed(42)
    model = MLP().to(device)
    init_fn(model)

    # Capturar activaciones
    acts = capture_activations(model, train_loader)

    for col, (layer_name, layer_acts) in enumerate(acts.items()):
        ax = axes[row, col]
        flat = layer_acts.flatten()

        # Limitar rango para visualización
        valid = flat[np.isfinite(flat)]
        if len(valid) > 0 and np.std(valid) > 0:
            clip_val = min(np.percentile(np.abs(valid), 99), 50)
            clipped = np.clip(valid, -clip_val, clip_val)
            ax.hist(clipped, bins=50, color='#0984E3', alpha=0.7, density=True)
            ax.set_title(f'{layer_name}', fontsize=8)
            ax.text(0.95, 0.95, f'μ={np.mean(valid):.2f}\nσ={np.std(valid):.2f}',
                    transform=ax.transAxes, ha='right', va='top', fontsize=7,
                    bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8))
        else:
            ax.text(0.5, 0.5, 'Todo ceros\no NaN/Inf',
                    transform=ax.transAxes, ha='center', va='center', fontsize=9)
            ax.set_title(f'{layer_name}', fontsize=8)

        if col == 0:
            ax.set_ylabel(name, fontsize=9, fontweight='bold')
        ax.tick_params(labelsize=6)

plt.suptitle('Distribución de activaciones ANTES del entrenamiento (por capa)',
             fontsize=13, fontweight='bold', y=1.01)
plt.tight_layout()
plt.show()

def train_model(model, train_loader, test_loader, epochs=15, lr=0.01):
    """Entrena un modelo y devuelve métricas por época."""
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

    history = {
        'train_loss': [],
        'test_loss': [],
        'train_acc': [],
        'test_acc': [],
    }

    for epoch in range(epochs):
        # ── Entrenamiento ──
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Proteger contra NaN
            if torch.isnan(loss) or torch.isinf(loss):
                history['train_loss'].append(float('nan'))
                history['test_loss'].append(float('nan'))
                history['train_acc'].append(0.0)
                history['test_acc'].append(0.0)
                print(f"  Epoch {epoch+1}: ⚠️ NaN/Inf detectado — entrenamiento abortado")
                return history

            loss.backward()
            # Gradient clipping para evitar explosión total
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
            optimizer.step()

            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        train_loss = running_loss / total
        train_acc = 100.0 * correct / total

        # ── Evaluación ──
        model.eval()
        test_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                test_loss += loss.item() * images.size(0)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        test_loss /= total
        test_acc = 100.0 * correct / total

        history['train_loss'].append(train_loss)
        history['test_loss'].append(test_loss)
        history['train_acc'].append(train_acc)
        history['test_acc'].append(test_acc)

        print(f"  Epoch {epoch+1:2d}/{epochs} | "
              f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.1f}% | "
              f"Test Acc: {test_acc:.1f}%")

    return history

print("Función de entrenamiento definida ✅")

# ── Entrenar con cada estrategia ──
EPOCHS = 15
results = {}

for name, init_fn in INIT_STRATEGIES.items():
    print(f"\n{'='*60}")
    print(f"🔧 Inicialización: {name}")
    print(f"{'='*60}")

    # Crear modelo fresco con la misma seed
    torch.manual_seed(42)
    model = MLP().to(device)
    init_fn(model)

    # Verificar distribución de pesos inicial
    all_weights = []
    for pname, param in model.named_parameters():
        if 'weight' in pname:
            all_weights.append(param.data.cpu().numpy().flatten())
    all_w = np.concatenate(all_weights)
    print(f"  Pesos iniciales: μ={np.mean(all_w):.4f}, σ={np.std(all_w):.4f}, "
          f"min={np.min(all_w):.4f}, max={np.max(all_w):.4f}")

    # Entrenar
    history = train_model(model, train_loader, test_loader, epochs=EPOCHS)
    results[name] = history

print(f"\n{'='*60}")
print("✅ Entrenamiento completado para todas las estrategias")
print(f"{'='*60}")

# Colores para cada estrategia
colors = {
    'Ceros':         '#e74c3c',   # rojo
    'Unos':          '#e67e22',   # naranja
    'Normal(0,1)':   '#9b59b6',   # púrpura
    'Uniform[-1,1]': '#f39c12',   # amarillo
    'Xavier':        '#2ecc71',   # verde
    'He/Kaiming':    '#0984E3',   # azul
}

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

for name, hist in results.items():
    epochs_range = range(1, len(hist['train_loss']) + 1)
    # Filtrar NaN para el plot
    train_loss = [v if not np.isnan(v) else None for v in hist['train_loss']]

    ax1.plot(epochs_range, train_loss, label=name, color=colors[name],
             linewidth=2, marker='o', markersize=3)

    test_loss = [v if not np.isnan(v) else None for v in hist['test_loss']]
    ax2.plot(epochs_range, test_loss, label=name, color=colors[name],
             linewidth=2, marker='o', markersize=3)

ax1.set_title('Loss de Entrenamiento', fontsize=12, fontweight='bold')
ax1.set_xlabel('Época')
ax1.set_ylabel('Cross-Entropy Loss')
ax1.legend(fontsize=8)
ax1.set_ylim(bottom=0)
ax1.grid(True, alpha=0.3)

ax2.set_title('Loss de Test', fontsize=12, fontweight='bold')
ax2.set_xlabel('Época')
ax2.set_ylabel('Cross-Entropy Loss')
ax2.legend(fontsize=8)
ax2.set_ylim(bottom=0)
ax2.grid(True, alpha=0.3)

plt.suptitle('Evolución de la Loss según la Inicialización',
             fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

for name, hist in results.items():
    epochs_range = range(1, len(hist['train_acc']) + 1)
    ax1.plot(epochs_range, hist['train_acc'], label=name,
             color=colors[name], linewidth=2, marker='o', markersize=3)
    ax2.plot(epochs_range, hist['test_acc'], label=name,
             color=colors[name], linewidth=2, marker='o', markersize=3)

ax1.set_title('Accuracy de Entrenamiento', fontsize=12, fontweight='bold')
ax1.set_xlabel('Época')
ax1.set_ylabel('Accuracy (%)')
ax1.legend(fontsize=8)
ax1.set_ylim(0, 100)
ax1.axhline(y=10, color='gray', linestyle='--', alpha=0.5, label='Azar (10%)')
ax1.grid(True, alpha=0.3)

ax2.set_title('Accuracy de Test', fontsize=12, fontweight='bold')
ax2.set_xlabel('Época')
ax2.set_ylabel('Accuracy (%)')
ax2.legend(fontsize=8)
ax2.set_ylim(0, 100)
ax2.axhline(y=10, color='gray', linestyle='--', alpha=0.5, label='Azar (10%)')
ax2.grid(True, alpha=0.3)

plt.suptitle('Evolución del Accuracy según la Inicialización',
             fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()

# Accuracy final de test para cada estrategia
final_accs = {}
for name, hist in results.items():
    accs = hist['test_acc']
    final_accs[name] = accs[-1] if accs else 0.0

# Ordenar de mayor a menor
sorted_names = sorted(final_accs, key=final_accs.get, reverse=True)
sorted_accs = [final_accs[n] for n in sorted_names]
sorted_colors = [colors[n] for n in sorted_names]

fig, ax = plt.subplots(figsize=(10, 5))
bars = ax.barh(sorted_names, sorted_accs, color=sorted_colors, edgecolor='white',
               linewidth=0.5, height=0.6)

# Etiquetas en cada barra
for bar, acc in zip(bars, sorted_accs):
    ax.text(bar.get_width() + 0.5, bar.get_y() + bar.get_height()/2,
            f'{acc:.1f}%', va='center', fontweight='bold', fontsize=10)

ax.set_xlabel('Test Accuracy (%)', fontsize=11)
ax.set_title('Accuracy Final de Test por Estrategia de Inicialización',
             fontsize=13, fontweight='bold')
ax.set_xlim(0, 100)
ax.axvline(x=10, color='gray', linestyle='--', alpha=0.5)
ax.text(11, -0.5, 'Azar', color='gray', fontsize=8)
ax.grid(axis='x', alpha=0.3)
ax.invert_yaxis()
plt.tight_layout()
plt.show()

# Tabla resumen
print("\n" + "="*55)
print(f"{'Estrategia':<18} {'Test Acc':>10} {'Veredicto':>25}")
print("="*55)
for name in sorted_names:
    acc = final_accs[name]
    if acc < 15:
        verdict = '❌ No aprende'
    elif acc < 70:
        verdict = '⚠️ Aprende mal'
    elif acc < 85:
        verdict = '🔶 Subóptimo'
    else:
        verdict = '✅ Buen resultado'
    print(f"{name:<18} {acc:>9.1f}% {verdict:>25}")
print("="*55)

# Re-entrenar brevemente para capturar pesos finales
# (ya tenemos los históricos, pero necesitamos los modelos finales para pesos)
final_models = {}

for name, init_fn in INIT_STRATEGIES.items():
    torch.manual_seed(42)
    model = MLP().to(device)
    init_fn(model)

    # Entrenar silenciosamente
    model.train()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

    for epoch in range(EPOCHS):
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            if torch.isnan(loss) or torch.isinf(loss):
                break
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
            optimizer.step()

    final_models[name] = model

# Visualizar distribución de pesos de la primera capa
fig, axes = plt.subplots(2, 3, figsize=(15, 8))

for idx, (name, model) in enumerate(final_models.items()):
    ax = axes[idx // 3, idx % 3]

    # Obtener pesos de la primera capa lineal
    first_layer = model.network[0]
    weights = first_layer.weight.data.cpu().numpy().flatten()

    if np.all(np.isfinite(weights)) and np.std(weights) > 0:
        ax.hist(weights, bins=80, color=colors[name], alpha=0.7, density=True)
        ax.set_title(f'{name}\nμ={np.mean(weights):.4f}, σ={np.std(weights):.4f}',
                     fontsize=9, fontweight='bold')
    else:
        ax.text(0.5, 0.5, 'Pesos no finitos\no sin variación',
                transform=ax.transAxes, ha='center', va='center', fontsize=10)
        ax.set_title(f'{name}', fontsize=9, fontweight='bold')

    ax.tick_params(labelsize=7)
    ax.grid(True, alpha=0.2)

plt.suptitle('Distribución de pesos de la 1ª capa oculta DESPUÉS del entrenamiento (15 épocas)',
             fontsize=12, fontweight='bold', y=1.01)
plt.tight_layout()
plt.show()