import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.datasets import load_wine
from sklearn.preprocessing import StandardScaler

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

print(f"Dataset: {df.shape[0]} muestras, {df.shape[1]-1} features, {len(data.target_names)} clases")
print(f"Clases: {list(data.target_names)}")
print(f"\nMuestras por clase:")
print(df['target'].value_counts().sort_index().to_string())

# ── Matriz de correlación ────────────────────────────────
plt.figure(figsize=(10, 8))
sns.heatmap(df.corr(), annot=True, fmt='.1f', cmap='RdBu_r', center=0,
            square=True, linewidths=0.5, annot_kws={'size': 7})
plt.title('Correlación entre variables', fontsize=14)
plt.tight_layout()
plt.show()

# ── Distribuciones por clase ──────────────────────────────
fig, axes = plt.subplots(4, 4, figsize=(16, 12))
for i, col in enumerate(data.feature_names):
    ax = axes.flat[i]
    for cls in range(3):
        mask = df['target'] == cls
        ax.hist(df.loc[mask, col], bins=15, alpha=0.5,
                label=data.target_names[cls])
    ax.set_title(col, fontsize=9)
    ax.legend(fontsize=6)
    ax.grid(alpha=0.3)
# Ocultar subplot extra
for j in range(len(data.feature_names), 16):
    axes.flat[j].set_visible(False)
plt.suptitle('Distribución de cada variable por clase', fontsize=14)
plt.tight_layout()
plt.show()

# ── Funciones de activación y derivadas ───────────────────
def sigmoid(x):
    return 1 / (1 + np.exp(-np.clip(x, -500, 500)))

def sigmoid_derivative(x):
    s = sigmoid(x)
    return s * (1 - s)

def tanh_act(x):
    return np.tanh(x)

def tanh_derivative(x):
    return 1 - np.tanh(x)**2

def relu(x):
    return np.maximum(0, x)

def relu_derivative(x):
    return (x > 0).astype(float)

def leaky_relu(x, alpha=0.01):
    return np.where(x > 0, x, alpha * x)

def leaky_relu_derivative(x, alpha=0.01):
    return np.where(x > 0, 1.0, alpha)

# ── Visualizar las funciones ──────────────────────────────
x = np.linspace(-5, 5, 400)
fig, axes = plt.subplots(2, 4, figsize=(16, 6))

funcs = [
    ('Sigmoid', sigmoid, sigmoid_derivative),
    ('Tanh', tanh_act, tanh_derivative),
    ('ReLU', relu, relu_derivative),
    ('Leaky ReLU', leaky_relu, leaky_relu_derivative),
]

for i, (name, f, df) in enumerate(funcs):
    # Función
    axes[0, i].plot(x, f(x), lw=2, color='#6C5CE7')
    axes[0, i].set_title(f'{name}', fontsize=12)
    axes[0, i].axhline(0, color='gray', lw=0.5)
    axes[0, i].axvline(0, color='gray', lw=0.5)
    axes[0, i].grid(alpha=0.3)
    # Derivada
    axes[1, i].plot(x, df(x), lw=2, color='#00B894')
    axes[1, i].set_title(f"{name}' (derivada)", fontsize=11)
    axes[1, i].axhline(0, color='gray', lw=0.5)
    axes[1, i].axvline(0, color='gray', lw=0.5)
    axes[1, i].grid(alpha=0.3)

plt.suptitle('Funciones de activación y sus derivadas', fontsize=14)
plt.tight_layout()
plt.show()

class NeuralNetwork:
    """
    MLP de 2 capas ocultas con activación configurable.
    Softmax en la salida para clasificación multiclase.
    """
    def __init__(self, input_size, hidden_sizes, output_size,
                 activation='sigmoid', learning_rate=0.01, leaky_alpha=0.01):
        self.lr = learning_rate
        self.leaky_alpha = leaky_alpha
        
        # Inicialización He
        self.W1 = np.random.randn(input_size, hidden_sizes[0]) * np.sqrt(2. / input_size)
        self.b1 = np.zeros((1, hidden_sizes[0]))
        self.W2 = np.random.randn(hidden_sizes[0], hidden_sizes[1]) * np.sqrt(2. / hidden_sizes[0])
        self.b2 = np.zeros((1, hidden_sizes[1]))
        self.W3 = np.random.randn(hidden_sizes[1], output_size) * np.sqrt(2. / hidden_sizes[1])
        self.b3 = np.zeros((1, output_size))
        
        # Seleccionar función de activación
        act_map = {
            'sigmoid': (sigmoid, sigmoid_derivative),
            'tanh': (tanh_act, tanh_derivative),
            'relu': (relu, relu_derivative),
            'leaky_relu': (
                lambda x: leaky_relu(x, self.leaky_alpha),
                lambda x: leaky_relu_derivative(x, self.leaky_alpha)
            ),
        }
        if activation not in act_map:
            raise ValueError(f"Activación '{activation}' no reconocida")
        self.activation, self.activation_derivative = act_map[activation]
    
    def softmax(self, x):
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)
    
    def forward(self, X):
        self.Z1 = X @ self.W1 + self.b1
        self.A1 = self.activation(self.Z1)
        self.Z2 = self.A1 @ self.W2 + self.b2
        self.A2 = self.activation(self.Z2)
        self.Z3 = self.A2 @ self.W3 + self.b3
        self.A3 = self.softmax(self.Z3)
        return self.A3
    
    def compute_loss(self, y_true, y_pred):
        m = y_true.shape[0]
        log_lik = -np.log(y_pred[range(m), y_true] + 1e-9)
        return np.sum(log_lik) / m
    
    def backward(self, X, y_true):
        m = X.shape[0]
        # One-hot
        y_oh = np.zeros_like(self.A3)
        y_oh[np.arange(m), y_true] = 1
        
        dZ3 = self.A3 - y_oh
        dW3 = self.A2.T @ dZ3 / m
        db3 = np.sum(dZ3, axis=0, keepdims=True) / m
        
        dA2 = dZ3 @ self.W3.T
        dZ2 = dA2 * self.activation_derivative(self.Z2)
        dW2 = self.A1.T @ dZ2 / m
        db2 = np.sum(dZ2, axis=0, keepdims=True) / m
        
        dA1 = dZ2 @ self.W2.T
        dZ1 = dA1 * self.activation_derivative(self.Z1)
        dW1 = X.T @ dZ1 / m
        db1 = np.sum(dZ1, axis=0, keepdims=True) / m
        
        # Actualizar pesos
        self.W3 -= self.lr * dW3; self.b3 -= self.lr * db3
        self.W2 -= self.lr * dW2; self.b2 -= self.lr * db2
        self.W1 -= self.lr * dW1; self.b1 -= self.lr * db1
    
    def train(self, X, y, epochs=20, batch_size=32, verbose=True):
        history = {'loss': [], 'accuracy': []}
        for epoch in range(epochs):
            idx = np.random.permutation(X.shape[0])
            X_s, y_s = X[idx], y[idx]
            
            epoch_loss = 0
            for i in range(0, X.shape[0], batch_size):
                X_b, y_b = X_s[i:i+batch_size], y_s[i:i+batch_size]
                y_pred = self.forward(X_b)
                loss = self.compute_loss(y_b, y_pred)
                epoch_loss += loss * X_b.shape[0]
                self.backward(X_b, y_b)
            
            epoch_loss /= X.shape[0]
            acc = np.mean(self.predict(X) == y)
            history['loss'].append(epoch_loss)
            history['accuracy'].append(acc)
            if verbose and (epoch + 1) % 5 == 0:
                print(f"  Epoch {epoch+1:>3}/{epochs} — Loss: {epoch_loss:.4f} — Acc: {acc*100:.1f}%")
        return history
    
    def predict(self, X):
        return np.argmax(self.forward(X), axis=1)

print('✅ Clase NeuralNetwork definida (13 → 32 → 16 → 3)')

# ── Preparación de datos ─────────────────────────────────
X = data.data
y = data.target

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# ── Entrenamiento ─────────────────────────────────────────
activations = ['sigmoid', 'tanh', 'relu', 'leaky_relu']
histories = {}
models = {}

for act in activations:
    print(f"\n{'═'*50}")
    print(f"  Activación: {act.upper()}")
    print(f"{'═'*50}")
    np.random.seed(42)  # Misma inicialización para todos
    model = NeuralNetwork(
        input_size=X.shape[1],
        hidden_sizes=[32, 16],
        output_size=len(np.unique(y)),
        activation=act,
        learning_rate=0.05
    )
    history = model.train(X_scaled, y, epochs=30, batch_size=16)
    histories[act] = history
    models[act] = model

# ── Curvas de pérdida ─────────────────────────────────────
colors = {'sigmoid': '#E17055', 'tanh': '#6C5CE7', 'relu': '#00B894', 'leaky_relu': '#FDCB6E'}

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Loss
for act in activations:
    axes[0].plot(histories[act]['loss'], label=act, lw=2, color=colors[act])
axes[0].set_xlabel('Epoch', fontsize=12)
axes[0].set_ylabel('Cross-Entropy Loss', fontsize=12)
axes[0].set_title('Evolución de la pérdida', fontsize=14)
axes[0].legend(fontsize=11)
axes[0].grid(alpha=0.3)

# Accuracy
for act in activations:
    axes[1].plot(histories[act]['accuracy'], label=act, lw=2, color=colors[act])
axes[1].set_xlabel('Epoch', fontsize=12)
axes[1].set_ylabel('Accuracy', fontsize=12)
axes[1].set_title('Evolución de la accuracy', fontsize=14)
axes[1].legend(fontsize=11)
axes[1].grid(alpha=0.3)
axes[1].set_ylim(0, 1.05)

plt.tight_layout()
plt.show()

# ── Resumen final ─────────────────────────────────────────
print(f"{'Activación':<15} {'Loss final':>12} {'Accuracy':>10}")
print('─' * 40)
for act in activations:
    loss = histories[act]['loss'][-1]
    acc = histories[act]['accuracy'][-1]
    print(f"{act:<15} {loss:>12.4f} {acc*100:>9.1f}%")

# Bar plot de accuracy final
fig, ax = plt.subplots(figsize=(8, 4))
accs = [histories[a]['accuracy'][-1] * 100 for a in activations]
bars = ax.bar(activations, accs, color=[colors[a] for a in activations],
              edgecolor='white', width=0.5)
for bar, acc in zip(bars, accs):
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
            f'{acc:.1f}%', ha='center', fontsize=12)
ax.set_ylabel('Accuracy (%)', fontsize=12)
ax.set_title('Accuracy final por función de activación', fontsize=14)
ax.set_ylim(0, 105)
ax.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.show()