import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import r2_score, mean_absolute_error

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# ── Cargar el dataset de diabetes ──────────────────────────
diabetes = load_diabetes()
X = diabetes.data       # (442, 10) — 10 características clínicas
y = diabetes.target      # (442,)   — progresión de la enfermedad

# Convertimos a DataFrame para facilitar la exploración
df = pd.DataFrame(X, columns=diabetes.feature_names)
df["target"] = y

print(f"Dataset: {X.shape[0]} pacientes, {X.shape[1]} características")
print(f"Variable objetivo — min: {y.min():.0f}, max: {y.max():.0f}, media: {y.mean():.1f}")

print(df.describe())

df.hist(bins=15, figsize=(15,10))
plt.tight_layout()
plt.show()

# Correlación con la variable objetivo
corr_matrix = df.corr()
plt.figure(figsize=(12,10))
sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
plt.show()

# Correlación específica con la variable objetivo
corr_target = corr_matrix['target'].sort_values(ascending=False)
print(corr_target)

# ── Preparación de datos para entrenamiento ────────────────

# Separar características (X) y variable objetivo (y)
X = df.drop("target", axis=1).values
y = df["target"].values.reshape(-1, 1)  # Reshape a columna para compatibilidad

# Split 80% entrenamiento / 20% test
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Estandarización: media=0, std=1
# IMPORTANTE: fit SOLO en train, transform en ambos (evita data leakage)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Convertir arrays NumPy → tensores PyTorch
X_train_tensor = torch.FloatTensor(X_train)
y_train_tensor = torch.FloatTensor(y_train)
X_test_tensor = torch.FloatTensor(X_test)
y_test_tensor = torch.FloatTensor(y_test)

# DataLoaders: permiten iterar en mini-batches durante el entrenamiento
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

print(f"Train: {len(train_dataset)} muestras | Test: {len(test_dataset)} muestras")

# ── Definición de la red neuronal ───────────────────────────
class DiabetesNet(nn.Module):
    """
    Red fully-connected parametrizable.
    - n1: neuronas en la primera capa oculta
    - n2: neuronas en la segunda capa oculta
    Variando n1 y n2 controlamos la capacidad del modelo.
    """
    def __init__(self, n1=64, n2=32):
        super(DiabetesNet, self).__init__()
        self.fc1 = nn.Linear(10, n1)    # Capa 1: 10 entradas → n1 neuronas
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(n1, n2)    # Capa 2: n1 → n2 neuronas
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(n2, 1)     # Salida: n2 → 1 (regresión)
    
    def forward(self, x):
        out = self.relu1(self.fc1(x))   # Forward pass capa 1
        out = self.relu2(self.fc2(out)) # Forward pass capa 2
        out = self.fc3(out)             # Predicción final (sin activación → regresión)
        return out

# Instanciamos el modelo con capacidad intermedia (64, 32)
model = DiabetesNet(64, 32)
print(f"Parámetros totales: {sum(p.numel() for p in model.parameters()):,}")

# ── Función de entrenamiento reutilizable ──────────────────
def train_model(model, num_epochs=1000, lr=0.01, verbose=False):
    """
    Entrena el modelo y devuelve las curvas de pérdida (train y test).
    """
    train_losses = []
    test_losses = []
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    for epoch in range(num_epochs):
        # ── Fase de entrenamiento ──
        model.train()
        running_loss = 0.0
        for features, labels in train_loader:
            optimizer.zero_grad()          # Limpiar gradientes acumulados
            outputs = model(features)       # Forward pass
            loss = criterion(outputs, labels)  # Calcular pérdida
            loss.backward()                 # Backpropagation
            optimizer.step()                # Actualizar pesos
            running_loss += loss.item() * features.size(0)
        
        epoch_loss = running_loss / len(train_loader.dataset)
        train_losses.append(epoch_loss)
        
        # ── Fase de evaluación (sin gradientes) ──
        model.eval()
        with torch.no_grad():
            test_running_loss = 0.0
            for features, labels in test_loader:
                outputs = model(features)
                loss = criterion(outputs, labels)
                test_running_loss += loss.item() * features.size(0)
            test_epoch_loss = test_running_loss / len(test_loader.dataset)
            test_losses.append(test_epoch_loss)
        
        if (epoch + 1) % 100 == 0 and verbose:
            print(f"Epoch [{epoch+1}/{num_epochs}] — Train Loss: {epoch_loss:.4f} | Val Loss: {test_epoch_loss:.4f}")
    
    return train_losses, test_losses

# Entrenamos el modelo durante 100 epochs
train_losses, test_losses = train_model(model, num_epochs=100, verbose=True)

# Curvas de entrenamiento
plt.figure(figsize=(10,5))
plt.plot(range(1, len(train_losses)+1), train_losses, label='Pérdida de entrenamiento')
plt.plot(range(1, len(test_losses)+1), test_losses, label='Pérdida de validación')
plt.xlabel('Épocas')
plt.ylabel('Pérdida (MSE)')
plt.title('Curvas de entrenamiento y validación')
plt.legend()
plt.show()

# ── Función de evaluación reutilizable ─────────────────────
def get_metrics(model, X_tensor, y_true, verbose=True):
    """Calcula R², MAE y MSE para un modelo dado."""
    model.eval()
    with torch.no_grad():
        predictions = model(X_tensor).numpy()
        
        r2 = r2_score(y_true, predictions)
        mae = mean_absolute_error(y_true, predictions)
        criterion = nn.MSELoss()
        mse = criterion(torch.FloatTensor(predictions), torch.FloatTensor(y_true)).item()

    if verbose:
        print(f"R² (coef. determinación): {r2:.4f}")
        print(f"MAE (error absoluto medio): {mae:.4f}")
        print(f"MSE (error cuadrático medio): {mse:.4f}")
    return r2, mae, mse

get_metrics(model, X_test_tensor, y_test)

predictions = model(X_test_tensor).detach().numpy()
true_values = y_test
plt.figure(figsize=(10,5))
plt.scatter(true_values, predictions, alpha=0.7)
plt.xlabel('Valores reales')
plt.ylabel('Predicciones')
plt.title('Comparación de valores reales vs predicciones')
plt.plot([true_values.min(), true_values.max()], [true_values.min(), true_values.max()], 'k--')
plt.show()

# ── Preparar datos con split más grande para test ──────────
X = df.drop("target", axis=1).values
y = df["target"].values.reshape(-1, 1)

# 60% train / 40% test — más datos de test para medir varianza
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=42)

scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

X_train_tensor = torch.FloatTensor(X_train)
y_train_tensor = torch.FloatTensor(y_train)
X_test_tensor = torch.FloatTensor(X_test)
y_test_tensor = torch.FloatTensor(y_test)

train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

print(f"Train: {len(train_dataset)} | Test: {len(test_dataset)} (split 60/40)")

n_epochs = 600

# Modelo súper simple
model_0 = DiabetesNet(4,2)
train_losses_0, test_losses_0 = train_model(model_0, num_epochs=n_epochs, verbose=False)
r20, mae0, mse0 = get_metrics(model_0, X_test_tensor, y_test, verbose=False)
r20_train, mae0_train, mse0_train = get_metrics(model_0, X_train_tensor, y_train, verbose=False)

# Modelo intermedio
model_1 = DiabetesNet(16,8)
train_losses_1, test_losses_1 = train_model(model_1, num_epochs=n_epochs, verbose=False)
r21, mae1, mse1 = get_metrics(model_1, X_test_tensor, y_test, verbose=False)
r21_train, mae1_train, mse1_train = get_metrics(model_1, X_train_tensor, y_train, verbose=False)

# Modelo complejo
model_2 = DiabetesNet(64,32)
train_losses_2, test_losses_2 = train_model(model_2, num_epochs=n_epochs, verbose=False)
r22, mae2, mse2 = get_metrics(model_2, X_test_tensor, y_test, verbose=False)
r22_train, mae2_train, mse2_train = get_metrics(model_2, X_train_tensor, y_train, verbose=False)

# Modelo super complejo
model_3 = DiabetesNet(128,64)
train_losses_3, test_losses_3 = train_model(model_3, num_epochs=n_epochs, verbose=False)
r23, mae3, mse3 = get_metrics(model_3, X_test_tensor, y_test, verbose=False)
r23_train, mae3_train, mse3_train = get_metrics(model_3, X_train_tensor, y_train, verbose=False)

plt.figure(figsize=(10,5))
plt.xlabel('Épocas')
plt.ylabel('Pérdida (MSE)')
plt.title('Curvas de validación para diferentes modelos')
tl0 = test_losses_0[::1] # smooth out the curve
tl1 = test_losses_1[::1]
tl2 = test_losses_2[::1]
tl3 = test_losses_3[::1]
tr0 = train_losses_0[::1]
tr1 = train_losses_1[::1]
tr2 = train_losses_2[::1]
tr3 = train_losses_3[::1]
# plot train as solid lines and test as dashed
plt.plot(range(1, len(tr0)+1), tr0, label='Entrenamiento (4,2)', color='blue')
plt.plot(range(1, len(tl0)+1), tl0, label='Validación (4,2)', color='blue', linestyle='--')
plt.plot(range(1, len(tr1)+1), tr1, label='Entrenamiento (16,8)', color='red')
plt.plot(range(1, len(tl1)+1), tl1, label='Validación (16,8)', color='red', linestyle='--')
plt.plot(range(1, len(tr2)+1), tr2, label='Entrenamiento (64,32)', color='green')
plt.plot(range(1, len(tl2)+1), tl2, label='Validación (64,32)', color='green', linestyle='--')
plt.plot(range(1, len(tr3)+1), tr3, label='Entrenamiento (128,64)', color='purple')
plt.plot(range(1, len(tl3)+1), tl3, label='Validación (128,64)', color='purple', linestyle='--')

plt.legend()

# Plot the mse values for train and test
plt.figure(figsize=(10,5))
plt.xlabel('Modelo')
plt.ylabel('MSE')
plt.title('Comparación de MSE para diferentes modelos')
# plot both train and test as bars with no fill
plt.bar([0,1,2,3], [mse0, mse1, mse2, mse3], label='Validación', color='blue', fill=True, lw=3, alpha=0.7)
plt.bar([0,1,2,3], [mse0_train, mse1_train, mse2_train, mse3_train], label='Entrenamiento', color='red', fill=False, lw=3, alpha=0.7)
plt.legend()


plt.show()

# ── Función para calcular varianza de los errores ──────────
def variance(model, X_test, y_test):
    """Varianza de los residuos (predicción - valor real)."""
    predictions = model(X_test).detach().numpy()
    true_values = np.array(y_test)
    return np.var(predictions - true_values)

vars = []
for model in [model_0, model_1, model_2, model_3]:
    var = variance(model, X_test_tensor, y_test)
    vars.append(var)

plt.figure(figsize=(10,5))
plt.xlabel('Modelo')
plt.ylabel('Varianza')
plt.title('Varianza de MSE para diferentes modelos')
plt.bar([0,1,2,3], vars, color='blue', fill=True, lw=3, alpha=0.7)
plt.show()

# ── Barrido de complejidad ──────────────────────────────────
model_complexity = np.arange(10, 200, 10, dtype=int)
vars_list = []
mse_test = []
mse_train = []

for i, n in enumerate(model_complexity):
    print(f"Entrenando modelo ({n}, {n//2}) — [{i+1}/{len(model_complexity)}]...")
    model = DiabetesNet(n, n // 2)
    train_losses, test_losses = train_model(model, num_epochs=400, verbose=False)
    
    # Métricas en test y train
    r2, mae, mse = get_metrics(model, X_test_tensor, y_test, verbose=False)
    r2train, maetrain, msetrain = get_metrics(model, X_train_tensor, y_train, verbose=False)
    
    vars_list.append(variance(model, X_test_tensor, y_test))
    mse_test.append(mse)
    mse_train.append(msetrain)

print("✅ Barrido completado")

# ── Gráfico 1: Varianza vs. Complejidad ────────────────────
plt.figure(figsize=(10, 5))
plt.plot(model_complexity, vars_list, color="blue", lw=3, marker="o", markersize=4)
plt.xlabel("Complejidad del modelo (neuronas en capa 1)")
plt.ylabel("Varianza de los residuos")
plt.title("Varianza vs. Complejidad del modelo")
plt.grid(alpha=0.3)
plt.show()

# ── Gráfico 2: MSE Train vs. Complejidad ──────────────────
plt.figure(figsize=(10, 5))
plt.plot(model_complexity, mse_train, label="MSE Entrenamiento", color="red", lw=3, marker="o", markersize=4)
plt.plot(model_complexity, mse_test, label="MSE Validación", color="blue", lw=3, marker="s", markersize=4)
plt.xlabel("Complejidad del modelo (neuronas en capa 1)")
plt.ylabel("MSE")
plt.title("MSE vs. Complejidad del modelo")
plt.legend()
plt.grid(alpha=0.3)
plt.show()

# ── Gráfico final: Sesgo vs. Varianza (normalizado) ────────
# Normalizamos ambas métricas para visualizarlas en la misma escala
vars_norm = np.array(vars_list) / np.max(vars_list)
mse_train_norm = np.array(mse_train) / np.max(mse_train)

plt.figure(figsize=(10, 5))
plt.plot(model_complexity, mse_train_norm, label="Sesgo (MSE train normalizado)", color="#E17055", lw=3)
plt.plot(model_complexity, vars_norm, label="Varianza (normalizada)", color="#6C5CE7", lw=3)
plt.fill_between(model_complexity, mse_train_norm, alpha=0.1, color="#E17055")
plt.fill_between(model_complexity, vars_norm, alpha=0.1, color="#6C5CE7")
plt.xlabel("Complejidad del modelo (neuronas en capa 1)")
plt.ylabel("Valor normalizado")
plt.title("Compromiso Sesgo-Varianza")
plt.legend(fontsize=12)
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()