# Configuración general y librerías

import warnings
warnings.filterwarnings('ignore')

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

from sklearn.datasets import fetch_covtype
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import SGDClassifier
from sklearn.ensemble import RandomForestClassifier, HistGradientBoostingClassifier
from sklearn.metrics import (
    accuracy_score, balanced_accuracy_score, f1_score,
    precision_score, recall_score, log_loss,
    confusion_matrix, ConfusionMatrixDisplay, classification_report
)

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

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

SEED = 42
np.random.seed(SEED)
torch.manual_seed(SEED)

# Carga del dataset Covertype
cov = fetch_covtype(as_frame=True)
df = cov.frame.copy()

# Ajustamos etiquetas de clase a rango [0..K-1] para PyTorch (CrossEntropyLoss)
df['Cover_Type'] = df['Cover_Type'] - 1

print('Shape total:', df.shape)
print('Número de clases:', df['Cover_Type'].nunique())
df.head()

# Información general y revisión de nulos
print(df.info())
print('\nNulos por columna:')
print(df.isnull().sum().sum())

# Distribución de clases
class_counts = df['Cover_Type'].value_counts().sort_index()
class_props = (class_counts / class_counts.sum()).round(4)

display(pd.DataFrame({'count': class_counts, 'proportion': class_props}))

plt.figure(figsize=(8, 4))
sns.barplot(x=class_counts.index, y=class_counts.values, palette='viridis')
plt.title('Distribución de clases (Cover_Type)')
plt.xlabel('Clase')
plt.ylabel('Número de muestras')
plt.tight_layout()
plt.show()

# Para EDA visual rápido, tomamos una muestra moderada
eda_sample = df.sample(5000, random_state=SEED)

num_cols = ['Elevation', 'Aspect', 'Slope', 'Horizontal_Distance_To_Hydrology',
            'Vertical_Distance_To_Hydrology', 'Horizontal_Distance_To_Roadways',
            'Hillshade_9am', 'Hillshade_Noon', 'Hillshade_3pm',
            'Horizontal_Distance_To_Fire_Points']

_ = eda_sample[num_cols + ['Cover_Type']].hist(bins=30, figsize=(14, 10), edgecolor='black')
plt.suptitle('Distribuciones (muestra EDA)', y=1.02)
plt.tight_layout()
plt.show()

# Matriz de correlación en variables numéricas principales
corr = eda_sample[num_cols + ['Cover_Type']].corr(numeric_only=True)

plt.figure(figsize=(10, 8))
sns.heatmap(corr, cmap='coolwarm', center=0)
plt.title('Correlación (muestra EDA)')
plt.tight_layout()
plt.show()

# Muestra estratificada para entrenamiento docente
N_SAMPLE = 120_000
sample_df = (
    df.groupby('Cover_Type', group_keys=False)
      .apply(lambda x: x.sample(int(np.floor(len(x) / len(df) * N_SAMPLE)), random_state=SEED))
      .sample(frac=1, random_state=SEED)
)

# Ajuste por redondeo para llegar exactamente a N_SAMPLE
if len(sample_df) < N_SAMPLE:
    faltan = N_SAMPLE - len(sample_df)
    extra = df.drop(sample_df.index).sample(faltan, random_state=SEED)
    sample_df = pd.concat([sample_df, extra], axis=0)

sample_df = sample_df.sample(frac=1, random_state=SEED).reset_index(drop=True)

X = sample_df.drop(columns=['Cover_Type']).values
y = sample_df['Cover_Type'].values.astype(int)

X_train_full, X_test, y_train_full, y_test = train_test_split(
    X, y, test_size=0.15, random_state=SEED, stratify=y
)
X_train, X_val, y_train, y_val = train_test_split(
    X_train_full, y_train_full, test_size=0.1765, random_state=SEED, stratify=y_train_full
)

print('Train:', X_train.shape, 'Val:', X_val.shape, 'Test:', X_test.shape)

# Escalado para lineal y redes
scaler = StandardScaler()
X_train_sc = scaler.fit_transform(X_train)
X_val_sc = scaler.transform(X_val)
X_test_sc = scaler.transform(X_test)

def classification_metrics(y_true, y_pred, y_proba, labels):
    """Métricas principales para clasificación multiclase."""
    return {
        'Accuracy': accuracy_score(y_true, y_pred),
        'BalancedAcc': balanced_accuracy_score(y_true, y_pred),
        'F1_macro': f1_score(y_true, y_pred, average='macro'),
        'Precision_macro': precision_score(y_true, y_pred, average='macro', zero_division=0),
        'Recall_macro': recall_score(y_true, y_pred, average='macro', zero_division=0),
        'LogLoss': log_loss(y_true, y_proba, labels=labels)
    }


def plot_curves(train_values, val_values, title, ylabel):
    epochs = np.arange(1, len(train_values) + 1)
    plt.figure(figsize=(8, 4.5))
    plt.plot(epochs, train_values, label='Train')
    plt.plot(epochs, val_values, label='Validación')
    plt.title(title)
    plt.xlabel('Época / iteración')
    plt.ylabel(ylabel)
    plt.legend()
    plt.tight_layout()
    plt.show()

labels = np.unique(y_train)

sgd_clf = SGDClassifier(
    loss='log_loss',
    penalty='l2',
    alpha=1e-4,
    learning_rate='optimal',
    random_state=SEED
)

n_epochs = 35
sgd_train_loss, sgd_val_loss = [], []
sgd_train_acc, sgd_val_acc = [], []

for epoch in range(n_epochs):
    sgd_clf.partial_fit(X_train_sc, y_train, classes=labels)

    # Probabilidades y predicciones para train/val
    train_proba = sgd_clf.predict_proba(X_train_sc)
    val_proba = sgd_clf.predict_proba(X_val_sc)
    train_pred = np.argmax(train_proba, axis=1)
    val_pred = np.argmax(val_proba, axis=1)

    sgd_train_loss.append(log_loss(y_train, train_proba, labels=labels))
    sgd_val_loss.append(log_loss(y_val, val_proba, labels=labels))
    sgd_train_acc.append(accuracy_score(y_train, train_pred))
    sgd_val_acc.append(accuracy_score(y_val, val_pred))

plot_curves(sgd_train_loss, sgd_val_loss, 'SGDClassifier - Loss (log-loss)', 'Log-loss')
plot_curves(sgd_train_acc, sgd_val_acc, 'SGDClassifier - Accuracy', 'Accuracy')

sgd_test_proba = sgd_clf.predict_proba(X_test_sc)
sgd_test_pred = np.argmax(sgd_test_proba, axis=1)
sgd_metrics = classification_metrics(y_test, sgd_test_pred, sgd_test_proba, labels)
print('Métricas test SGDClassifier:', sgd_metrics)

rf = RandomForestClassifier(
    n_estimators=10,
    max_depth=None,
    min_samples_leaf=1,
    random_state=SEED,
    n_jobs=-1,
    warm_start=True
)

rf_train_loss, rf_val_loss = [], []
rf_train_acc, rf_val_acc = [], []

n_trees_list = list(range(20, 221, 20))

for n_trees in n_trees_list:
    rf.set_params(n_estimators=n_trees)
    rf.fit(X_train, y_train)

    train_proba = rf.predict_proba(X_train)
    val_proba = rf.predict_proba(X_val)
    train_pred = np.argmax(train_proba, axis=1)
    val_pred = np.argmax(val_proba, axis=1)

    rf_train_loss.append(log_loss(y_train, train_proba, labels=labels))
    rf_val_loss.append(log_loss(y_val, val_proba, labels=labels))
    rf_train_acc.append(accuracy_score(y_train, train_pred))
    rf_val_acc.append(accuracy_score(y_val, val_pred))

plot_curves(rf_train_loss, rf_val_loss, 'Random Forest - Evolución de log-loss', 'Log-loss')
plot_curves(rf_train_acc, rf_val_acc, 'Random Forest - Evolución de accuracy', 'Accuracy')

rf_test_proba = rf.predict_proba(X_test)
rf_test_pred = np.argmax(rf_test_proba, axis=1)
rf_metrics = classification_metrics(y_test, rf_test_pred, rf_test_proba, labels)
print('Métricas test Random Forest:', rf_metrics)

hgb = HistGradientBoostingClassifier(
    loss='log_loss',
    learning_rate=0.08,
    max_depth=8,
    max_iter=1,
    random_state=SEED,
    warm_start=True
)

hgb_train_loss, hgb_val_loss = [], []
hgb_train_acc, hgb_val_acc = [], []

iters = list(range(10, 131, 10))

for it in iters:
    hgb.set_params(max_iter=it)
    hgb.fit(X_train, y_train)

    train_proba = hgb.predict_proba(X_train)
    val_proba = hgb.predict_proba(X_val)
    train_pred = np.argmax(train_proba, axis=1)
    val_pred = np.argmax(val_proba, axis=1)

    hgb_train_loss.append(log_loss(y_train, train_proba, labels=labels))
    hgb_val_loss.append(log_loss(y_val, val_proba, labels=labels))
    hgb_train_acc.append(accuracy_score(y_train, train_pred))
    hgb_val_acc.append(accuracy_score(y_val, val_pred))

plot_curves(hgb_train_loss, hgb_val_loss, 'HistGradientBoosting - Evolución de log-loss', 'Log-loss')
plot_curves(hgb_train_acc, hgb_val_acc, 'HistGradientBoosting - Evolución de accuracy', 'Accuracy')

hgb_test_proba = hgb.predict_proba(X_test)
hgb_test_pred = np.argmax(hgb_test_proba, axis=1)
hgb_metrics = classification_metrics(y_test, hgb_test_pred, hgb_test_proba, labels)
print('Métricas test HistGradientBoosting:', hgb_metrics)

# Conversión a tensores
X_train_t = torch.tensor(X_train_sc, dtype=torch.float32)
y_train_t = torch.tensor(y_train, dtype=torch.long)
X_val_t = torch.tensor(X_val_sc, dtype=torch.float32)
y_val_t = torch.tensor(y_val, dtype=torch.long)
X_test_t = torch.tensor(X_test_sc, dtype=torch.float32)
y_test_t = torch.tensor(y_test, dtype=torch.long)

train_loader = DataLoader(TensorDataset(X_train_t, y_train_t), batch_size=256, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val_t, y_val_t), batch_size=512, shuffle=False)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
num_classes = len(labels)
print('Dispositivo:', device, '| Clases:', num_classes)

class SimpleClassifier(nn.Module):
    def __init__(self, in_features, n_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, 64),
            nn.ReLU(),
            nn.Linear(64, n_classes)
        )

    def forward(self, x):
        return self.net(x)


class DeepClassifier(nn.Module):
    def __init__(self, in_features, n_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.25),

            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.20),

            nn.Linear(128, 64),
            nn.ReLU(),

            nn.Linear(64, n_classes)
        )

    def forward(self, x):
        return self.net(x)

def train_torch_classifier(model, train_loader, val_loader, epochs=30, lr=1e-3):
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    hist = {
        'train_loss': [], 'val_loss': [],
        'train_acc': [], 'val_acc': []
    }

    for epoch in range(epochs):
        # Entrenamiento
        model.train()
        train_losses = []
        train_correct = 0
        train_total = 0

        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)

            optimizer.zero_grad()
            logits = model(xb)
            loss = criterion(logits, yb)
            loss.backward()
            optimizer.step()

            train_losses.append(loss.item())
            preds = torch.argmax(logits, dim=1)
            train_correct += (preds == yb).sum().item()
            train_total += yb.size(0)

        # Validación
        model.eval()
        val_losses = []
        val_correct = 0
        val_total = 0

        with torch.no_grad():
            for xb, yb in val_loader:
                xb, yb = xb.to(device), yb.to(device)
                logits = model(xb)
                loss = criterion(logits, yb)

                val_losses.append(loss.item())
                preds = torch.argmax(logits, dim=1)
                val_correct += (preds == yb).sum().item()
                val_total += yb.size(0)

        hist['train_loss'].append(float(np.mean(train_losses)))
        hist['val_loss'].append(float(np.mean(val_losses)))
        hist['train_acc'].append(train_correct / train_total)
        hist['val_acc'].append(val_correct / val_total)

    return model, hist


def predict_torch_classifier(model, X_tensor):
    model.eval()
    with torch.no_grad():
        logits = model(X_tensor.to(device))
        proba = torch.softmax(logits, dim=1).cpu().numpy()
        pred = np.argmax(proba, axis=1)
    return pred, proba

simple_nn = SimpleClassifier(in_features=X_train_sc.shape[1], n_classes=num_classes)
simple_nn, simple_hist = train_torch_classifier(simple_nn, train_loader, val_loader, epochs=35, lr=1e-3)

plot_curves(simple_hist['train_loss'], simple_hist['val_loss'], 'NN simple - Loss (CrossEntropy)', 'Loss')
plot_curves(simple_hist['train_acc'], simple_hist['val_acc'], 'NN simple - Accuracy', 'Accuracy')

simple_test_pred, simple_test_proba = predict_torch_classifier(simple_nn, X_test_t)
simple_metrics = classification_metrics(y_test, simple_test_pred, simple_test_proba, labels)
print('Métricas test NN simple:', simple_metrics)

deep_nn = DeepClassifier(in_features=X_train_sc.shape[1], n_classes=num_classes)
deep_nn, deep_hist = train_torch_classifier(deep_nn, train_loader, val_loader, epochs=45, lr=8e-4)

plot_curves(deep_hist['train_loss'], deep_hist['val_loss'], 'NN profunda - Loss (CrossEntropy)', 'Loss')
plot_curves(deep_hist['train_acc'], deep_hist['val_acc'], 'NN profunda - Accuracy', 'Accuracy')

deep_test_pred, deep_test_proba = predict_torch_classifier(deep_nn, X_test_t)
deep_metrics = classification_metrics(y_test, deep_test_pred, deep_test_proba, labels)
print('Métricas test NN profunda:', deep_metrics)

results = pd.DataFrame([
    {'Modelo': 'SGDClassifier (lineal)', **sgd_metrics},
    {'Modelo': 'Random Forest', **rf_metrics},
    {'Modelo': 'HistGradientBoosting', **hgb_metrics},
    {'Modelo': 'NN simple (PyTorch)', **simple_metrics},
    {'Modelo': 'NN profunda (PyTorch)', **deep_metrics},
]).sort_values('Accuracy', ascending=False)

results

fig, axes = plt.subplots(1, 3, figsize=(18, 4.5))

sns.barplot(data=results, x='Accuracy', y='Modelo', ax=axes[0], palette='Blues')
axes[0].set_title('Accuracy (mayor es mejor)')

sns.barplot(data=results, x='F1_macro', y='Modelo', ax=axes[1], palette='Greens')
axes[1].set_title('F1 macro (mayor es mejor)')

sns.barplot(data=results, x='LogLoss', y='Modelo', ax=axes[2], palette='Reds_r')
axes[2].set_title('Log-loss (menor es mejor)')

plt.tight_layout()
plt.show()

best_model_name = results.iloc[0]['Modelo']
print('Mejor modelo según Accuracy:', best_model_name)

pred_map = {
    'SGDClassifier (lineal)': (sgd_test_pred, sgd_test_proba),
    'Random Forest': (rf_test_pred, rf_test_proba),
    'HistGradientBoosting': (hgb_test_pred, hgb_test_proba),
    'NN simple (PyTorch)': (simple_test_pred, simple_test_proba),
    'NN profunda (PyTorch)': (deep_test_pred, deep_test_proba),
}

best_pred, best_proba = pred_map[best_model_name]

# Matriz de confusión
disp = ConfusionMatrixDisplay(
    confusion_matrix=confusion_matrix(y_test, best_pred, labels=labels),
    display_labels=labels
)
fig, ax = plt.subplots(figsize=(7, 7))
disp.plot(ax=ax, cmap='Blues', colorbar=False)
plt.title(f'Matriz de confusión - {best_model_name}')
plt.tight_layout()
plt.show()

# Reporte de clasificación
print(classification_report(y_test, best_pred, digits=4))

# Confianza máxima de la predicción
best_conf = best_proba.max(axis=1)
correct = (best_pred == y_test)

plt.figure(figsize=(9, 4.5))
sns.kdeplot(best_conf[correct], label='Predicciones correctas', fill=True)
sns.kdeplot(best_conf[~correct], label='Predicciones incorrectas', fill=True)
plt.title('Distribución de confianza del mejor modelo')
plt.xlabel('Probabilidad máxima predicha')
plt.ylabel('Densidad')
plt.legend()
plt.tight_layout()
plt.show()

assert all(np.isfinite(results['Accuracy'])), 'Accuracy no finita detectada'
assert all(np.isfinite(results['F1_macro'])), 'F1 no finita detectada'
assert all(np.isfinite(results['LogLoss'])), 'LogLoss no finita detectada'
assert (results['Accuracy'] >= 0).all() and (results['Accuracy'] <= 1).all(), 'Accuracy fuera de rango'

print('Accuracy mejor modelo:', float(results.iloc[0]['Accuracy']))
print('LogLoss mejor modelo:', float(results.sort_values('LogLoss').iloc[0]['LogLoss']))
print('✅ Sanity checks completados correctamente')