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

# ── Cargar el dataset ───────────────────────────────────
diabetes = load_diabetes()
X = diabetes.data       # (442, 10)
y = diabetes.target      # (442,)

df = pd.DataFrame(X, columns=diabetes.feature_names)
df["target"] = y

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

print(df.describe())

df.hist(bins=15, figsize=(15, 10))
plt.suptitle("Distribución de cada variable", fontsize=14, y=1.01)
plt.tight_layout()
plt.show()

# Mapa de calor de la matriz de correlación
corr_matrix = df.corr()
plt.figure(figsize=(12, 10))
sns.heatmap(corr_matrix, annot=True, cmap="coolwarm", fmt=".2f")
plt.title("Matriz de correlación")
plt.show()

# Top correlaciones con target
print("\nCorrelación con la variable objetivo:")
print(corr_matrix["target"].sort_values(ascending=False).to_string())

# ── Funciones de activación ────────────────────────────────
def relu(x):
    """ReLU: max(0, x) — introduce no linealidad."""
    return np.maximum(0, x)

def relu_derivative(x):
    """Derivada de ReLU: 1 si x > 0, 0 en otro caso."""
    return (x > 0).astype(float)


# ── Red neuronal implementada con NumPy ──────────────────
class NeuralNetwork:
    """
    MLP de 2 capas ocultas para regresión.
    Arquitectura: input → 32 → 16 → 1
    """
    def __init__(self, input_dim, lr=0.01, seed=42):
        np.random.seed(seed)
        self.lr = lr
        
        # Inicialización aleatoria de pesos (distribución normal escalada)
        self.W1 = np.random.randn(input_dim, 32) * 0.01
        self.b1 = np.zeros((1, 32))
        self.W2 = np.random.randn(32, 16) * 0.01
        self.b2 = np.zeros((1, 16))
        self.W3 = np.random.randn(16, 1) * 0.01
        self.b3 = np.zeros((1, 1))
    
    def forward(self, X):
        """Forward pass: propaga los datos capa a capa."""
        self.z1 = X @ self.W1 + self.b1          # Pre-activación capa 1
        self.a1 = relu(self.z1)                    # Activación ReLU
        self.z2 = self.a1 @ self.W2 + self.b2     # Pre-activación capa 2
        self.a2 = relu(self.z2)                    # Activación ReLU
        self.z3 = self.a2 @ self.W3 + self.b3     # Salida (lineal, regresión)
        return self.z3
    
    def compute_loss(self, y_true, y_pred):
        """MSE Loss: L = (1/2m) * Σ(y - ŷ)²"""
        m = y_true.shape[0]
        return np.mean((y_true - y_pred) ** 2) / 2
    
    def backward(self, X, y_true, y_pred):
        """Backward pass: calcula gradientes por regla de la cadena."""
        m = X.shape[0]
        
        # Gradiente de la pérdida respecto a la salida
        dz3 = (y_pred - y_true) / m
        
        # Capa 3 (salida)
        self.dW3 = self.a2.T @ dz3
        self.db3 = np.sum(dz3, axis=0, keepdims=True)
        
        # Capa 2
        da2 = dz3 @ self.W3.T
        dz2 = da2 * relu_derivative(self.z2)
        self.dW2 = self.a1.T @ dz2
        self.db2 = np.sum(dz2, axis=0, keepdims=True)
        
        # Capa 1
        da1 = dz2 @ self.W2.T
        dz1 = da1 * relu_derivative(self.z1)
        self.dW1 = X.T @ dz1
        self.db1 = np.sum(dz1, axis=0, keepdims=True)
    
    def update_parameters(self):
        """Descenso del gradiente: W ← W - α·∂L/∂W"""
        self.W1 -= self.lr * self.dW1
        self.b1 -= self.lr * self.db1
        self.W2 -= self.lr * self.dW2
        self.b2 -= self.lr * self.db2
        self.W3 -= self.lr * self.dW3
        self.b3 -= self.lr * self.db3
    
    def train(self, X, y, epochs=1000, X_test=None, y_test=None, verbose=True):
        """Bucle de entrenamiento completo."""
        train_losses = []
        test_losses = []
        
        for epoch in range(epochs):
            # Forward
            y_pred = self.forward(X)
            loss = self.compute_loss(y, y_pred)
            train_losses.append(loss)
            
            # Backward + update
            self.backward(X, y, y_pred)
            self.update_parameters()
            
            # Evaluar en test
            if X_test is not None:
                y_test_pred = self.forward(X_test)
                test_loss = self.compute_loss(y_test, y_test_pred)
                test_losses.append(test_loss)
                # Re-forward en train para restaurar activaciones
                self.forward(X)
            
            if verbose and (epoch + 1) % 200 == 0:
                msg = f"Epoch {epoch+1}/{epochs} — Train Loss: {loss:.4f}"
                if test_losses:
                    msg += f" | Test Loss: {test_losses[-1]:.4f}"
                print(msg)
        
        return train_losses, test_losses
    
    def predict(self, X):
        """Genera predicciones (forward pass sin gradientes)."""
        return self.forward(X)

print("✅ Clase NeuralNetwork definida")

# ── Preparación de datos ──────────────────────────────────
X = diabetes.data
y = diabetes.target.reshape(-1, 1)  # Columna para compatibilidad matricial

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

# Estandarizar X e y
scaler_X = StandardScaler()
scaler_y = StandardScaler()

X_train_s = scaler_X.fit_transform(X_train)
X_test_s = scaler_X.transform(X_test)
y_train_s = scaler_y.fit_transform(y_train)
y_test_s = scaler_y.transform(y_test)

print(f"Train: {X_train_s.shape[0]} muestras | Test: {X_test_s.shape[0]} muestras")
print(f"Features shape: {X_train_s.shape[1]} | Target shape: {y_train_s.shape[1]}")

# ── Crear y entrenar la red ───────────────────────────────
nn = NeuralNetwork(input_dim=X_train_s.shape[1], lr=0.1)
train_losses, test_losses = nn.train(
    X_train_s, y_train_s,
    epochs=1000,
    X_test=X_test_s, y_test=y_test_s
)

# ── Curvas de pérdida ──────────────────────────────────────
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label="Train Loss", color="#6C5CE7", lw=2)
plt.plot(test_losses, label="Test Loss", color="#E17055", lw=2, linestyle="--")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.title("Curvas de entrenamiento y validación")
plt.legend()
plt.grid(alpha=0.3)
plt.show()

print(f"\nLoss final — Train: {train_losses[-1]:.4f} | Test: {test_losses[-1]:.4f}")

# ── Predicciones en escala original ───────────────────────
y_train_pred = scaler_y.inverse_transform(nn.predict(X_train_s))
y_test_pred = scaler_y.inverse_transform(nn.predict(X_test_s))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

# Train
ax1.scatter(y_train, y_train_pred, alpha=0.6, color="#6C5CE7", s=30)
lims = [y_train.min(), y_train.max()]
ax1.plot(lims, lims, "k--", lw=1.5, label="Predicción perfecta")
ax1.set_xlabel("Valor real")
ax1.set_ylabel("Predicción")
ax1.set_title("Entrenamiento")
ax1.legend()
ax1.grid(alpha=0.3)

# Test
ax2.scatter(y_test, y_test_pred, alpha=0.6, color="#E17055", s=30)
lims = [y_test.min(), y_test.max()]
ax2.plot(lims, lims, "k--", lw=1.5, label="Predicción perfecta")
ax2.set_xlabel("Valor real")
ax2.set_ylabel("Predicción")
ax2.set_title("Test")
ax2.legend()
ax2.grid(alpha=0.3)

plt.suptitle("Predicciones vs. Valores Reales (escala original)", fontsize=14)
plt.tight_layout()
plt.show()