import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import time
from sklearn.datasets import fetch_california_housing
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

# ── Cargar dataset ────────────────────────────────────────
data = fetch_california_housing()
df = pd.DataFrame(data.data, columns=data.feature_names)
df['MedHouseVal'] = data.target

print(f"Dataset: {df.shape[0]:,} muestras, {df.shape[1]-1} features")
print(f"\nEstadísticas principales:")
df.describe().round(2)

# ── Matriz de correlación ────────────────────────────────
corr = df.corr()

plt.figure(figsize=(10, 8))
sns.heatmap(corr, annot=True, fmt='.2f', cmap='RdBu_r', center=0,
            square=True, linewidths=0.5)
plt.title('Correlación entre variables', fontsize=14)
plt.tight_layout()
plt.show()

# ── Histogramas ────────────────────────────────────────────
fig, axes = plt.subplots(3, 3, figsize=(14, 10))
for i, col in enumerate(df.columns):
    ax = axes.flat[i]
    ax.hist(df[col], bins=40, alpha=0.7, color='#6C5CE7', edgecolor='white')
    ax.set_title(col, fontsize=11)
    ax.grid(alpha=0.3)
plt.suptitle('Distribución de cada variable', fontsize=14)
plt.tight_layout()
plt.show()

# ── Preparación ───────────────────────────────────────────
X, y = data.data, data.target.reshape(-1, 1)

# Estandarizar features y target
scaler_X = StandardScaler()
scaler_y = StandardScaler()
X = scaler_X.fit_transform(X)
y = scaler_y.fit_transform(y)

# Split train/test
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=11
)

print(f"Train: {X_train.shape} | Test: {X_test.shape}")

class NeuralNetwork:
    """
    MLP de 2 capas ocultas (64 → 32 → 1) implementado con NumPy.
    Usa ReLU, inicialización He e incluye forward, backward y update.
    """
    def __init__(self, input_size, hidden1=64, hidden2=32, output_size=1, lr=0.01, seed=42):
        np.random.seed(seed)
        # Inicialización He (óptima para ReLU)
        self.W1 = np.random.randn(input_size, hidden1) * np.sqrt(2. / input_size)
        self.b1 = np.zeros((1, hidden1))
        self.W2 = np.random.randn(hidden1, hidden2) * np.sqrt(2. / hidden1)
        self.b2 = np.zeros((1, hidden2))
        self.W3 = np.random.randn(hidden2, output_size) * np.sqrt(2. / hidden2)
        self.b3 = np.zeros((1, output_size))
        self.lr = lr

    def relu(self, Z):
        return np.maximum(0, Z)
    
    def relu_derivative(self, Z):
        return (Z > 0).astype(float)
    
    def forward(self, X):
        # Capa 1: entrada → 64 neuronas
        self.Z1 = X @ self.W1 + self.b1
        self.A1 = self.relu(self.Z1)
        # Capa 2: 64 → 32 neuronas
        self.Z2 = self.A1 @ self.W2 + self.b2
        self.A2 = self.relu(self.Z2)
        # Salida: 32 → 1 (regresión lineal)
        self.Z3 = self.A2 @ self.W3 + self.b3
        return self.Z3
    
    def compute_loss(self, y_pred, y_true):
        return np.mean((y_pred - y_true) ** 2)
    
    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        # Gradiente de MSE: dL/dy_pred = 2(y_pred - y_true)/m
        dZ3 = (y_pred - y_true) * (2 / m)
        dW3 = self.A2.T @ dZ3
        db3 = np.sum(dZ3, axis=0, keepdims=True)
        
        dA2 = dZ3 @ self.W3.T
        dZ2 = dA2 * self.relu_derivative(self.Z2)
        dW2 = self.A1.T @ dZ2
        db2 = np.sum(dZ2, axis=0, keepdims=True)
        
        dA1 = dZ2 @ self.W2.T
        dZ1 = dA1 * self.relu_derivative(self.Z1)
        dW1 = X.T @ dZ1
        db1 = np.sum(dZ1, axis=0, keepdims=True)
        
        return {'dW1': dW1, 'db1': db1, 'dW2': dW2, 'db2': db2, 'dW3': dW3, 'db3': db3}
    
    def update_parameters(self, grads):
        for param in ['W1', 'b1', 'W2', 'b2', 'W3', 'b3']:
            setattr(self, param, getattr(self, param) - self.lr * grads[f'd{param}'])

print('✅ Clase NeuralNetwork definida')
print(f'   Arquitectura: 8 → 64 (ReLU) → 32 (ReLU) → 1')
print(f'   Parámetros: {8*64+64 + 64*32+32 + 32*1+1:,}')

def train_model(model, X, y, epochs=50, batch_size=32):
    """Entrena el modelo con el batch_size especificado."""
    history_epoch = []   # Pérdida media por época
    history_batch = []   # Pérdida por batch (para ver estabilidad)
    m = X.shape[0]
    
    for epoch in range(epochs):
        # Barajar datos cada época
        idx = np.random.permutation(m)
        X_shuf, y_shuf = X[idx], y[idx]
        
        epoch_loss = 0
        batch_losses = []
        
        for i in range(0, m, batch_size):
            X_b = X_shuf[i:i+batch_size]
            y_b = y_shuf[i:i+batch_size]
            
            y_pred = model.forward(X_b)           # Forward
            loss = model.compute_loss(y_pred, y_b) # Pérdida
            grads = model.backward(X_b, y_b, y_pred) # Backward
            model.update_parameters(grads)         # Actualizar
            
            epoch_loss += loss * X_b.shape[0]
            batch_losses.append(loss)
        
        epoch_loss /= m
        history_epoch.append(epoch_loss)
        history_batch.extend(batch_losses)
        
        if (epoch + 1) % 2 == 0:
            print(f'  Época {epoch+1:>3}/{epochs} — MSE: {epoch_loss:.4f}')
    
    return history_epoch, history_batch

print('✅ Función train_model definida')

# ── Configuración ─────────────────────────────────────────
epochs = 10
lr = 0.001

# ── 1) Batch completo (GD clásico) ───────────────────────
print('═' * 50)
print('MÉTODO 1: Batch completo (GD)')
print('═' * 50)
model_full = NeuralNetwork(input_size=X_train.shape[1], lr=lr)
t0 = time.time()
hist_full, batches_full = train_model(model_full, X_train, y_train, epochs, batch_size=X_train.shape[0])
time_full = time.time() - t0
print(f'  ⏱ Tiempo: {time_full:.3f}s\n')

# ── 2) SGD puro (batch_size = 1) ─────────────────────────
print('═' * 50)
print('MÉTODO 2: SGD (batch_size = 1)')
print('═' * 50)
model_sgd = NeuralNetwork(input_size=X_train.shape[1], lr=lr)
t0 = time.time()
hist_sgd, batches_sgd = train_model(model_sgd, X_train, y_train, epochs, batch_size=1)
time_sgd = time.time() - t0
print(f'  ⏱ Tiempo: {time_sgd:.3f}s\n')

# ── 3) Mini-batch (batch_size = 32) ──────────────────────
print('═' * 50)
print('MÉTODO 3: Mini-batch (batch_size = 32)')
print('═' * 50)
model_mini = NeuralNetwork(input_size=X_train.shape[1], lr=lr)
t0 = time.time()
hist_mini, batches_mini = train_model(model_mini, X_train, y_train, epochs, batch_size=32)
time_mini = time.time() - t0
print(f'  ⏱ Tiempo: {time_mini:.3f}s')

# ── Pérdida por época ─────────────────────────────────────
plt.figure(figsize=(10, 5))
plt.plot(hist_full, 'o-', label=f'Batch completo ({time_full:.2f}s)', lw=2)
plt.plot(hist_sgd, 's-', label=f'SGD, bs=1 ({time_sgd:.2f}s)', lw=2)
plt.plot(hist_mini, '^-', label=f'Mini-batch, bs=32 ({time_mini:.2f}s)', lw=2)
plt.xlabel('Época', fontsize=12)
plt.ylabel('MSE', fontsize=12)
plt.title('Evolución de la pérdida por época', fontsize=14)
plt.legend(fontsize=11)
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# ── Pérdida batch a batch ────────────────────────────────
# Interpolamos para que los 3 tengan la misma longitud
max_len = len(batches_sgd)

# GD: 1 punto por época → repetir para igualar
full_interp = []
pts_per_epoch = max_len // epochs
for val in batches_full:
    full_interp.extend([val] * pts_per_epoch)

# Mini-batch: interpolar
mini_interp = []
pts_mini_epoch = max_len // epochs
batches_per_epoch_mini = len(batches_mini) // epochs
for ep in range(epochs):
    ep_batches = batches_mini[ep*batches_per_epoch_mini:(ep+1)*batches_per_epoch_mini]
    ratio = pts_mini_epoch // max(len(ep_batches), 1)
    for v in ep_batches:
        mini_interp.extend([v] * ratio)

# Submuestrear para la visualización
step = max(1, max_len // 500)

plt.figure(figsize=(12, 5))
plt.plot(full_interp[::step], label='Batch completo', alpha=0.9, lw=1.5)
plt.plot(batches_sgd[::step], label='SGD (bs=1)', alpha=0.5, lw=0.5)
plt.plot(mini_interp[::step], label='Mini-batch (bs=32)', alpha=0.7, lw=1)
plt.ylim(0, 4)
plt.xlabel('Paso de actualización', fontsize=12)
plt.ylabel('MSE', fontsize=12)
plt.title('Estabilidad del entrenamiento (pérdida por batch)', fontsize=14)
plt.legend(fontsize=11)
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# ── Tiempo de entrenamiento ──────────────────────────────
methods = ['Batch completo', 'SGD (bs=1)', 'Mini-batch (bs=32)']
times = [time_full, time_sgd, time_mini]
colors = ['#6C5CE7', '#00B894', '#FDCB6E']

fig, ax = plt.subplots(figsize=(8, 4))
bars = ax.barh(methods, times, color=colors, edgecolor='white', height=0.5)
for bar, t in zip(bars, times):
    ax.text(bar.get_width() + 0.05, bar.get_y() + bar.get_height()/2,
            f'{t:.2f}s', va='center', fontsize=12)
ax.set_xlabel('Tiempo (segundos)', fontsize=12)
ax.set_title('Tiempo de entrenamiento por método', fontsize=14)
ax.grid(axis='x', alpha=0.3)
plt.tight_layout()
plt.show()