# Librerías y configuración

import os
# Forzar ejecución en CPU (evita errores del autotuner XLA/Triton en GPU)
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import time
import warnings
warnings.filterwarnings('ignore')

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

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

sns.set_theme(style='whitegrid', context='notebook')
plt.rcParams['figure.figsize'] = (9, 5)

SEED = 42
np.random.seed(SEED)
tf.random.set_seed(SEED)

# Carga de FashionMNIST
(X_train_full, y_train_full), (X_test, y_test) = keras.datasets.fashion_mnist.load_data()

# Split train/val
X_train, X_val = X_train_full[:54000], X_train_full[54000:]
y_train, y_val = y_train_full[:54000], y_train_full[54000:]

class_names = [
    'T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
    'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'
]

print('Train:', X_train.shape, 'Val:', X_val.shape, 'Test:', X_test.shape)
print('Rango píxeles:', X_train.min(), 'a', X_train.max())

# EDA: distribución de clases
counts = pd.Series(y_train).value_counts().sort_index()

plt.figure(figsize=(10, 4))
sns.barplot(x=[class_names[i] for i in counts.index], y=counts.values, palette='viridis')
plt.xticks(rotation=30, ha='right')
plt.title('Distribución de clases en train')
plt.ylabel('Número de muestras')
plt.tight_layout()
plt.show()

# EDA: muestras visuales
plt.figure(figsize=(12, 6))
for i in range(20):
    ax = plt.subplot(4, 5, i+1)
    plt.imshow(X_train[i], cmap='gray')
    plt.title(class_names[y_train[i]], fontsize=9)
    plt.axis('off')
plt.suptitle('Ejemplos de FashionMNIST')
plt.tight_layout()
plt.show()

# Preprocesado: normalización y flatten implícito en el modelo
X_train = X_train.astype('float32') / 255.0
X_val = X_val.astype('float32') / 255.0
X_test = X_test.astype('float32') / 255.0

def build_mlp(input_shape=(28, 28), n_classes=10):
    """MLP sencillo para clasificación multiclase en FashionMNIST."""
    model = keras.Sequential([
        layers.Input(shape=input_shape),
        layers.Flatten(),
        layers.Dense(256, activation='relu'),
        layers.Dense(128, activation='relu'),
        layers.Dense(n_classes, activation='softmax')
    ])
    return model


def compile_model(model, learning_rate):
    """Compila el modelo con SGD puro (sin momento) y sparse categorical crossentropy.
    
    Usamos SGD en lugar de Adam para que el efecto del learning rate
    y batch size sea claramente visible en las curvas de entrenamiento.
    """
    model.compile(
        optimizer=keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.0),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )
    return model


def train_once(learning_rate, batch_size, epochs=15, verbose=0):
    """Entrena un modelo y devuelve historial + métricas + tiempos."""
    model = build_mlp()
    model = compile_model(model, learning_rate)

    t0 = time.perf_counter()
    history = model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        epochs=epochs,
        batch_size=batch_size,
        verbose=verbose
    )
    train_time = time.perf_counter() - t0

    # Inferencia en test
    t1 = time.perf_counter()
    y_prob = model.predict(X_test, verbose=0)
    infer_time = time.perf_counter() - t1

    y_pred = np.argmax(y_prob, axis=1)

    metrics = {
        'Accuracy': accuracy_score(y_test, y_pred),
        'Precision_macro': precision_score(y_test, y_pred, average='macro', zero_division=0),
        'Recall_macro': recall_score(y_test, y_pred, average='macro', zero_division=0),
        'F1_macro': f1_score(y_test, y_pred, average='macro', zero_division=0),
        'Train_time_s': train_time,
        'Infer_time_s': infer_time,
    }

    return model, history.history, metrics


def plot_train_val(history, title_prefix='Modelo'):
    """Pinta loss y accuracy train/val de un historial Keras."""
    # Loss
    plt.figure(figsize=(8, 4.5))
    plt.plot(history['loss'], label='Train')
    plt.plot(history['val_loss'], label='Validación')
    plt.title(f'{title_prefix} - Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.tight_layout()
    plt.show()

    # Accuracy
    plt.figure(figsize=(8, 4.5))
    plt.plot(history['accuracy'], label='Train')
    plt.plot(history['val_accuracy'], label='Validación')
    plt.title(f'{title_prefix} - Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.tight_layout()
    plt.show()

# Rango amplio de learning rates: desde muy lento hasta divergente
learning_rates = [1e-4, 1e-3, 1e-2, 1e-1, 1.0]
fixed_batch = 128
EPOCHS = 15

lr_results = []
lr_histories = {}

for lr in learning_rates:
    print(f'\nEntrenando con learning_rate={lr} y batch_size={fixed_batch}')
    _, hist, metrics = train_once(learning_rate=lr, batch_size=fixed_batch, epochs=EPOCHS, verbose=0)
    lr_histories[lr] = hist

    row = {'learning_rate': lr, 'batch_size': fixed_batch, **metrics}
    lr_results.append(row)

    # Curvas por entrenamiento
    plot_train_val(hist, title_prefix=f'LR={lr}')

lr_df = pd.DataFrame(lr_results).sort_values('Accuracy', ascending=False).reset_index(drop=True)
lr_df

# Comparativas visuales para learning rate
fig, axes = plt.subplots(1, 3, figsize=(18, 4.5))

sns.barplot(data=lr_df, x='learning_rate', y='Accuracy', ax=axes[0], palette='Blues')
axes[0].set_xscale('log')
axes[0].set_title('Accuracy test vs learning rate')

sns.barplot(data=lr_df, x='learning_rate', y='F1_macro', ax=axes[1], palette='Greens')
axes[1].set_xscale('log')
axes[1].set_title('F1 macro vs learning rate')

sns.barplot(data=lr_df, x='learning_rate', y='Train_time_s', ax=axes[2], palette='Reds')
axes[2].set_xscale('log')
axes[2].set_title('Tiempo entrenamiento vs learning rate')

plt.tight_layout()
plt.show()

# Curvas superpuestas para comparar estabilidad por learning rate
plt.figure(figsize=(9, 5))
for lr in learning_rates:
    plt.plot(lr_histories[lr]['val_loss'], label=f'lr={lr}')
plt.title('Validación Loss vs Epoch para distintos learning rates')
plt.xlabel('Epoch')
plt.ylabel('Val loss')
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(9, 5))
for lr in learning_rates:
    plt.plot(lr_histories[lr]['val_accuracy'], label=f'lr={lr}')
plt.title('Validación Accuracy vs Epoch para distintos learning rates')
plt.xlabel('Epoch')
plt.ylabel('Val accuracy')
plt.legend()
plt.tight_layout()
plt.show()

# Rango amplio de batch sizes: desde muy ruidoso hasta muy suave
batch_sizes = [8, 32, 128, 512, 2048]
fixed_lr = 0.01

bs_results = []
bs_histories = {}

for bs in batch_sizes:
    print(f'\nEntrenando con learning_rate={fixed_lr} y batch_size={bs}')
    _, hist, metrics = train_once(learning_rate=fixed_lr, batch_size=bs, epochs=EPOCHS, verbose=0)
    bs_histories[bs] = hist

    row = {'batch_size': bs, 'learning_rate': fixed_lr, **metrics}
    bs_results.append(row)

    # Curvas por entrenamiento
    plot_train_val(hist, title_prefix=f'Batch={bs}')

bs_df = pd.DataFrame(bs_results).sort_values('Accuracy', ascending=False).reset_index(drop=True)
bs_df

# Comparativas visuales para batch size
fig, axes = plt.subplots(1, 3, figsize=(18, 4.5))

sns.barplot(data=bs_df, x='batch_size', y='Accuracy', ax=axes[0], palette='Blues')
axes[0].set_title('Accuracy test vs batch size')

sns.barplot(data=bs_df, x='batch_size', y='F1_macro', ax=axes[1], palette='Greens')
axes[1].set_title('F1 macro vs batch size')

sns.barplot(data=bs_df, x='batch_size', y='Train_time_s', ax=axes[2], palette='Reds')
axes[2].set_title('Tiempo entrenamiento vs batch size')

plt.tight_layout()
plt.show()

# Curvas superpuestas para comparar estabilidad por batch size
plt.figure(figsize=(9, 5))
for bs in batch_sizes:
    plt.plot(bs_histories[bs]['val_loss'], label=f'batch={bs}')
plt.title('Validación Loss vs Epoch para distintos batch sizes')
plt.xlabel('Epoch')
plt.ylabel('Val loss')
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(9, 5))
for bs in batch_sizes:
    plt.plot(bs_histories[bs]['val_accuracy'], label=f'batch={bs}')
plt.title('Validación Accuracy vs Epoch para distintos batch sizes')
plt.xlabel('Epoch')
plt.ylabel('Val accuracy')
plt.legend()
plt.tight_layout()
plt.show()

lr_df['Experimento'] = 'learning_rate_sweep'
bs_df['Experimento'] = 'batch_size_sweep'

final_df = pd.concat([lr_df, bs_df], ignore_index=True, sort=False)
final_df

# Mejores configuraciones por cada barrido
best_lr = lr_df.sort_values('Accuracy', ascending=False).iloc[0]
best_bs = bs_df.sort_values('Accuracy', ascending=False).iloc[0]

print('Mejor learning rate (batch fijo=128):')
print(best_lr)
print('\nMejor batch size (lr fijo=1e-3):')
print(best_bs)

assert len(lr_df) == 5, 'Deben evaluarse 5 learning rates'
assert len(bs_df) == 5, 'Deben evaluarse 5 batch sizes'

for df_name, df in [('lr_df', lr_df), ('bs_df', bs_df)]:
    for col in ['Accuracy', 'Precision_macro', 'Recall_macro', 'F1_macro']:
        assert np.isfinite(df[col]).all(), f'{df_name}:{col} contiene no finitos'
        assert ((df[col] >= 0) & (df[col] <= 1)).all(), f'{df_name}:{col} fuera de [0,1]'

    assert (df['Train_time_s'] > 0).all(), f'{df_name}:Train_time_s debe ser > 0'
    assert (df['Infer_time_s'] > 0).all(), f'{df_name}:Infer_time_s debe ser > 0'

# Chequeo de longitud de curvas
for lr in learning_rates:
    assert len(lr_histories[lr]['loss']) == EPOCHS
for bs in batch_sizes:
    assert len(bs_histories[bs]['loss']) == EPOCHS

print('✅ Sanity checks completados correctamente')