# Librerías base
import os
import random
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms

from sklearn.metrics import (
    roc_auc_score,
    average_precision_score,
    confusion_matrix,
    classification_report,
    roc_curve,
    precision_recall_curve,
)

# Configuración estética para gráficas
sns.set(style="whitegrid", context="notebook")
plt.rcParams["figure.figsize"] = (8, 4)

# Reproducibilidad
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

# Dispositivo de cómputo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Dispositivo en uso: {device}")

# Transformación a tensor en rango [0, 1]
transform = transforms.ToTensor()

train_raw = datasets.FashionMNIST(root="./data", train=True, download=True, transform=transform)
test_raw = datasets.FashionMNIST(root="./data", train=False, download=True, transform=transform)

class_names = train_raw.classes
print("Clases:", class_names)
print(f"Tamaño train: {len(train_raw)} | Tamaño test: {len(test_raw)}")

# Distribución de clases en train
train_targets = torch.tensor(train_raw.targets)
counts = torch.bincount(train_targets)

plt.figure(figsize=(10, 4))
plt.bar(class_names, counts.numpy())
plt.xticks(rotation=45, ha="right")
plt.title("Distribución de clases - FashionMNIST (train)")
plt.ylabel("Número de muestras")
plt.show()

for i, c in enumerate(counts.tolist()):
    print(f"Clase {i} ({class_names[i]}): {c}")

# Visualización de ejemplos por clase
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
axes = axes.flatten()

for cls in range(10):
    idx = (train_targets == cls).nonzero(as_tuple=False)[0].item()
    img, label = train_raw[idx]
    axes[cls].imshow(img.squeeze(), cmap="gray")
    axes[cls].set_title(class_names[label])
    axes[cls].axis("off")

plt.suptitle("Un ejemplo por clase")
plt.tight_layout()
plt.show()

# Estadísticas globales de intensidad de píxel
all_pixels = train_raw.data.float() / 255.0
print(f"Media global de píxel: {all_pixels.mean().item():.4f}")
print(f"Desviación estándar global: {all_pixels.std().item():.4f}")
print(f"Mínimo: {all_pixels.min().item():.4f} | Máximo: {all_pixels.max().item():.4f}")

# Parámetros del experimento
NORMAL_CLASS = 0
BATCH_SIZE = 128

# Convertimos imágenes a float [0,1] y luego a vector (784)
X_train = train_raw.data.float() / 255.0
y_train = train_targets
X_test = test_raw.data.float() / 255.0
y_test = torch.tensor(test_raw.targets)

X_train_flat = X_train.view(-1, 28 * 28)
X_test_flat = X_test.view(-1, 28 * 28)

# Índices por clase
train_normal_idx = (y_train == NORMAL_CLASS).nonzero(as_tuple=False).squeeze()
train_anom_idx = (y_train != NORMAL_CLASS).nonzero(as_tuple=False).squeeze()

# Split de normales train/val
n_norm = len(train_normal_idx)
perm = torch.randperm(n_norm)
train_norm_idx = train_normal_idx[perm[: int(0.8 * n_norm)]]
val_norm_idx = train_normal_idx[perm[int(0.8 * n_norm):]]

# Tomamos anomalías para val y para monitor train
n_val_norm = len(val_norm_idx)
perm_anom = train_anom_idx[torch.randperm(len(train_anom_idx))]
val_anom_idx = perm_anom[:n_val_norm]
monitor_anom_idx = perm_anom[n_val_norm:n_val_norm + 2000]
monitor_norm_idx = train_norm_idx[:2000]

# Tensores finales
X_train_norm = X_train_flat[train_norm_idx]
X_val_norm = X_train_flat[val_norm_idx]

X_val_mix = torch.cat([X_train_flat[val_norm_idx], X_train_flat[val_anom_idx]], dim=0)
y_val_mix = torch.cat([
    torch.zeros(len(val_norm_idx), dtype=torch.long),
    torch.ones(len(val_anom_idx), dtype=torch.long)
], dim=0)

X_monitor_mix = torch.cat([X_train_flat[monitor_norm_idx], X_train_flat[monitor_anom_idx]], dim=0)
y_monitor_mix = torch.cat([
    torch.zeros(len(monitor_norm_idx), dtype=torch.long),
    torch.ones(len(monitor_anom_idx), dtype=torch.long)
], dim=0)

# Test mixto completo
test_norm_idx = (y_test == NORMAL_CLASS).nonzero(as_tuple=False).squeeze()
test_anom_idx = (y_test != NORMAL_CLASS).nonzero(as_tuple=False).squeeze()

X_test_mix = torch.cat([X_test_flat[test_norm_idx], X_test_flat[test_anom_idx]], dim=0)
y_test_mix = torch.cat([
    torch.zeros(len(test_norm_idx), dtype=torch.long),
    torch.ones(len(test_anom_idx), dtype=torch.long)
], dim=0)

# DataLoaders
train_loader = DataLoader(TensorDataset(X_train_norm), batch_size=BATCH_SIZE, shuffle=True)
val_norm_loader = DataLoader(TensorDataset(X_val_norm), batch_size=BATCH_SIZE, shuffle=False)
val_mix_loader = DataLoader(TensorDataset(X_val_mix, y_val_mix), batch_size=BATCH_SIZE, shuffle=False)
monitor_loader = DataLoader(TensorDataset(X_monitor_mix, y_monitor_mix), batch_size=BATCH_SIZE, shuffle=False)
test_mix_loader = DataLoader(TensorDataset(X_test_mix, y_test_mix), batch_size=BATCH_SIZE, shuffle=False)

print(f"Train normal: {len(X_train_norm)}")
print(f"Val normal: {len(X_val_norm)}")
print(f"Val mixto (normal+anómalo): {len(X_val_mix)}")
print(f"Monitor train mixto: {len(X_monitor_mix)}")
print(f"Test mixto: {len(X_test_mix)}")

class AutoencoderMLP(nn.Module):
    def __init__(self, input_dim=784, latent_dim=32):
        super().__init__()
        # Encoder: comprime de 784 a latent_dim
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, latent_dim),
        )
        # Decoder: reconstruye de latent_dim a 784
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, input_dim),
            nn.Sigmoid(),
        )

    def forward(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat


model = AutoencoderMLP(input_dim=28*28, latent_dim=32).to(device)
print(model)

# Hiperparámetros
EPOCHS = 20
LR = 1e-3
WEIGHT_DECAY_L2 = 1e-5  # castigo por complejidad

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY_L2)


def reconstruction_errors(model, loader):
    """Devuelve errores de reconstrucción por muestra para un loader."""
    model.eval()
    errs, labels = [], []
    with torch.no_grad():
        for batch in loader:
            # Soporta loaders con (X,) o (X,y)
            if len(batch) == 1:
                x = batch[0].to(device)
                y = None
            else:
                x, y = batch
                x, y = x.to(device), y.to(device)

            x_hat = model(x)
            # Error medio por muestra (promedio sobre pixeles)
            e = ((x - x_hat) ** 2).mean(dim=1)
            errs.append(e.cpu())
            if y is not None:
                labels.append(y.cpu())

    errs = torch.cat(errs).numpy()
    labels = torch.cat(labels).numpy() if labels else None
    return errs, labels


def detection_accuracy(errors, labels, threshold):
    preds = (errors > threshold).astype(int)
    return (preds == labels).mean()

history = {
    'train_loss': [],
    'val_loss': [],
    'train_det_acc': [],
    'val_det_acc': [],
    'threshold': [],
}

for epoch in range(1, EPOCHS + 1):
    # ===== Entrenamiento =====
    model.train()
    train_losses = []

    for (x_batch,) in train_loader:
        x_batch = x_batch.to(device)

        optimizer.zero_grad()
        x_hat = model(x_batch)
        loss = criterion(x_hat, x_batch)
        loss.backward()
        optimizer.step()

        train_losses.append(loss.item())

    mean_train_loss = float(np.mean(train_losses))

    # ===== Validación (solo normales para loss) =====
    model.eval()
    val_losses = []
    with torch.no_grad():
        for (x_val,) in val_norm_loader:
            x_val = x_val.to(device)
            x_val_hat = model(x_val)
            vloss = criterion(x_val_hat, x_val)
            val_losses.append(vloss.item())

    mean_val_loss = float(np.mean(val_losses))

    # ===== Umbral + accuracy de detección =====
    # Umbral a partir del percentil 95 de error en datos normales de entrenamiento
    train_norm_errors, _ = reconstruction_errors(model, train_loader)
    threshold = np.percentile(train_norm_errors, 95)

    monitor_errors, monitor_labels = reconstruction_errors(model, monitor_loader)
    val_errors, val_labels = reconstruction_errors(model, val_mix_loader)

    train_acc = detection_accuracy(monitor_errors, monitor_labels, threshold)
    val_acc = detection_accuracy(val_errors, val_labels, threshold)

    # Guardamos histórico
    history['train_loss'].append(mean_train_loss)
    history['val_loss'].append(mean_val_loss)
    history['train_det_acc'].append(train_acc)
    history['val_det_acc'].append(val_acc)
    history['threshold'].append(threshold)

    print(
        f"Epoch {epoch:02d}/{EPOCHS} | "
        f"train_loss={mean_train_loss:.5f} | val_loss={mean_val_loss:.5f} | "
        f"train_acc={train_acc:.4f} | val_acc={val_acc:.4f} | "
        f"threshold={threshold:.5f}"
    )

# Curvas de entrenamiento: loss y accuracy de detección
epochs = np.arange(1, EPOCHS + 1)

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

axes[0].plot(epochs, history['train_loss'], marker='o', label='Train loss')
axes[0].plot(epochs, history['val_loss'], marker='o', label='Val loss (normales)')
axes[0].set_title('Loss de reconstrucción vs época')
axes[0].set_xlabel('Época')
axes[0].set_ylabel('MSE')
axes[0].legend()

axes[1].plot(epochs, history['train_det_acc'], marker='o', label='Train detection accuracy (monitor)')
axes[1].plot(epochs, history['val_det_acc'], marker='o', label='Val detection accuracy')
axes[1].set_title('Accuracy de detección vs época')
axes[1].set_xlabel('Época')
axes[1].set_ylabel('Accuracy')
axes[1].set_ylim(0.0, 1.0)
axes[1].legend()

plt.tight_layout()
plt.show()

# Errores y etiquetas en test
test_errors, test_labels = reconstruction_errors(model, test_mix_loader)

# Umbral final
final_threshold = history['threshold'][-1]
test_preds = (test_errors > final_threshold).astype(int)

# Métricas
test_acc = (test_preds == test_labels).mean()
roc_auc = roc_auc_score(test_labels, test_errors)
pr_auc = average_precision_score(test_labels, test_errors)

print(f"Threshold final: {final_threshold:.5f}")
print(f"Test accuracy: {test_acc:.4f}")
print(f"ROC-AUC: {roc_auc:.4f}")
print(f"PR-AUC: {pr_auc:.4f}")

# Curvas ROC y Precision-Recall
fpr, tpr, _ = roc_curve(test_labels, test_errors)
precision, recall, _ = precision_recall_curve(test_labels, test_errors)

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

axes[0].plot(fpr, tpr, label=f'ROC-AUC = {roc_auc:.3f}')
axes[0].plot([0, 1], [0, 1], '--', color='gray')
axes[0].set_title('Curva ROC')
axes[0].set_xlabel('False Positive Rate')
axes[0].set_ylabel('True Positive Rate')
axes[0].legend()

axes[1].plot(recall, precision, label=f'PR-AUC = {pr_auc:.3f}')
axes[1].set_title('Curva Precision-Recall')
axes[1].set_xlabel('Recall')
axes[1].set_ylabel('Precision')
axes[1].legend()

plt.tight_layout()
plt.show()

# Matriz de confusión + informe
cm = confusion_matrix(test_labels, test_preds)

plt.figure(figsize=(5, 4))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=['Pred normal', 'Pred anómalo'],
            yticklabels=['Real normal', 'Real anómalo'])
plt.title('Matriz de confusión - Test')
plt.show()

print(classification_report(test_labels, test_preds, target_names=['normal', 'anomalia']))

# Seleccionamos ejemplos aleatorios de test mixto para visualizar decisiones
n_examples = 12
indices = np.random.choice(len(X_test_mix), size=n_examples, replace=False)

x_samples = X_test_mix[indices].to(device)
y_samples = y_test_mix[indices].numpy()

model.eval()
with torch.no_grad():
    x_hat_samples = model(x_samples)

errs = ((x_samples - x_hat_samples) ** 2).mean(dim=1).cpu().numpy()
preds = (errs > final_threshold).astype(int)

fig, axes = plt.subplots(n_examples, 2, figsize=(6, 2.2 * n_examples))

for i in range(n_examples):
    original = x_samples[i].cpu().view(28, 28)
    recon = x_hat_samples[i].cpu().view(28, 28)

    axes[i, 0].imshow(original, cmap='gray')
    axes[i, 0].axis('off')
    axes[i, 0].set_title(f"Original | y={y_samples[i]}")

    axes[i, 1].imshow(recon, cmap='gray')
    axes[i, 1].axis('off')
    axes[i, 1].set_title(
        f"Recon | err={errs[i]:.4f} | pred={preds[i]}"
    )

plt.tight_layout()
plt.show()

# Resumen textual de la demostración
for i in range(n_examples):
    real = 'anomalia' if y_samples[i] == 1 else 'normal'
    pred = 'anomalia' if preds[i] == 1 else 'normal'
    print(
        f"Ejemplo {i:02d} -> real: {real:9s} | pred: {pred:9s} | error={errs[i]:.5f}"
    )