# Instalación opcional (descomenta si tu entorno no tiene dependencias)
# !pip install -q tensorflow scikit-learn matplotlib seaborn

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

# Imports principales
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

from sklearn.cluster import KMeans
from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score
from sklearn.metrics import confusion_matrix

# Configuración visual y semilla para reproducibilidad
sns.set(style="whitegrid")
np.random.seed(42)
tf.random.set_seed(42)

print("TensorFlow:", tf.__version__)

# Cargamos MNIST desde Keras
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

print("Train:", x_train.shape, y_train.shape)
print("Test :", x_test.shape, y_test.shape)

# Estadísticos básicos de los píxeles (antes de normalizar)
print("Rango de pixel train:", x_train.min(), "->", x_train.max())
print("Media train:", x_train.mean(), "Desviación:", x_train.std())

# Distribución de clases (solo EDA; no se usan etiquetas para entrenar el VAE)
unique, counts = np.unique(y_train, return_counts=True)
for d, c in zip(unique, counts):
    print(f"Dígito {d}: {c} muestras")

# Visualizamos ejemplos por clase para entender variabilidad intra-clase
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
axes = axes.ravel()

for digit in range(10):
    idx = np.where(y_train == digit)[0][0]
    axes[digit].imshow(x_train[idx], cmap='gray')
    axes[digit].set_title(f"Clase {digit}")
    axes[digit].axis('off')

plt.suptitle("Primer ejemplo encontrado por clase")
plt.tight_layout()
plt.show()

# Histograma de intensidades para entender contraste global del dataset
plt.figure(figsize=(8,4))
plt.hist(x_train.ravel(), bins=50, color='slateblue', alpha=0.85)
plt.title("Distribución de intensidades de píxel (train)")
plt.xlabel("Intensidad")
plt.ylabel("Frecuencia")
plt.show()

# Preprocesado para la red:
# 1) normalizar a [0,1]
# 2) aplanar a vector de 784
x_train_f = x_train.astype("float32") / 255.0
x_test_f  = x_test.astype("float32") / 255.0

x_train_f = x_train_f.reshape((-1, 784))
x_test_f = x_test_f.reshape((-1, 784))

# Separamos validación desde train
x_val_f = x_train_f[-10000:]
y_val = y_train[-10000:]
x_train_f = x_train_f[:-10000]
y_train_small = y_train[:-10000]

print("Train final:", x_train_f.shape)
print("Val final  :", x_val_f.shape)
print("Test final :", x_test_f.shape)

# Capa de muestreo (reparameterization trick)
class Sampling(layers.Layer):
    def call(self, inputs):
        z_mean, z_log_var = inputs
        epsilon = tf.random.normal(shape=tf.shape(z_mean))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

latent_dim = 2
input_dim = 784
hidden_dim = 256

# Encoder
encoder_inputs = keras.Input(shape=(input_dim,))
x = layers.Dense(hidden_dim, activation="relu")(encoder_inputs)
x = layers.Dense(hidden_dim // 2, activation="relu")(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])

encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
encoder.summary()

# Decoder
latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Dense(hidden_dim // 2, activation="relu")(latent_inputs)
x = layers.Dense(hidden_dim, activation="relu")(x)
decoder_outputs = layers.Dense(input_dim, activation="sigmoid")(x)

decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
decoder.summary()

# Modelo VAE custom para controlar losses y métricas de train/val
class VAE(keras.Model):
    def __init__(self, encoder, decoder, beta=1.0, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.beta = beta

        # Métricas rastreadas por Keras
        self.total_loss_tracker = keras.metrics.Mean(name="loss")
        self.rec_loss_tracker = keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")
        self.bin_acc = keras.metrics.BinaryAccuracy(name="binary_accuracy")

    @property
    def metrics(self):
        return [self.total_loss_tracker, self.rec_loss_tracker, self.kl_loss_tracker, self.bin_acc]

    def _compute_losses(self, data):
        """Calcula las tres pérdidas del VAE sobre un batch."""
        z_mean, z_log_var, z = self.encoder(data)
        reconstruction = self.decoder(z)

        # BCE por píxel: binary_crossentropy ya promedia sobre la última dimensión,
        # así que multiplicamos por el nº de features para obtener la suma sobre píxeles.
        bce_per_sample = keras.losses.binary_crossentropy(data, reconstruction)  # (batch,)
        n_features = tf.cast(tf.shape(data)[-1], tf.float32)
        rec_loss = tf.reduce_mean(bce_per_sample) * n_features

        kl_loss = -0.5 * tf.reduce_mean(
            tf.reduce_sum(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var), axis=1)
        )
        total_loss = rec_loss + self.beta * kl_loss
        return total_loss, rec_loss, kl_loss, reconstruction

    def train_step(self, data):
        # Keras puede pasar (x, None) cuando se usa validation_data=(x, None)
        if isinstance(data, tuple):
            data = data[0]

        with tf.GradientTape() as tape:
            total_loss, rec_loss, kl_loss, reconstruction = self._compute_losses(data)

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))

        self.total_loss_tracker.update_state(total_loss)
        self.rec_loss_tracker.update_state(rec_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        self.bin_acc.update_state(data, reconstruction)

        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.rec_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
            "binary_accuracy": self.bin_acc.result(),
        }

    def test_step(self, data):
        if isinstance(data, tuple):
            data = data[0]

        total_loss, rec_loss, kl_loss, reconstruction = self._compute_losses(data)

        self.total_loss_tracker.update_state(total_loss)
        self.rec_loss_tracker.update_state(rec_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        self.bin_acc.update_state(data, reconstruction)

        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.rec_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
            "binary_accuracy": self.bin_acc.result(),
        }

vae = VAE(encoder, decoder, beta=1.0)
vae.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-3))

# Entrenamos el VAE
history = vae.fit(
    x_train_f,
    epochs=20,
    batch_size=128,
    validation_data=(x_val_f, None),
    verbose=1
)

# Curvas de entrenamiento: loss y "accuracy" (binary_accuracy)
h = history.history
epochs = range(1, len(h['loss']) + 1)

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

# Loss total
axes[0].plot(epochs, h['loss'], label='train_loss')
axes[0].plot(epochs, h['val_loss'], label='val_loss')
axes[0].set_title('Loss total vs epoch')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss')
axes[0].legend()

# Accuracy de reconstrucción
axes[1].plot(epochs, h['binary_accuracy'], label='train_binary_accuracy')
axes[1].plot(epochs, h['val_binary_accuracy'], label='val_binary_accuracy')
axes[1].set_title('Binary accuracy vs epoch')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Accuracy')
axes[1].legend()

plt.tight_layout()
plt.show()

# Curvas separadas de reconstrucción y KL para entender "castigos y recompensas"
fig, axes = plt.subplots(1, 2, figsize=(14, 4))

axes[0].plot(epochs, h['reconstruction_loss'], label='train_rec_loss')
axes[0].plot(epochs, h['val_reconstruction_loss'], label='val_rec_loss')
axes[0].set_title('Término de reconstrucción (recompensa por fidelidad)')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Recon loss')
axes[0].legend()

axes[1].plot(epochs, h['kl_loss'], label='train_kl_loss')
axes[1].plot(epochs, h['val_kl_loss'], label='val_kl_loss')
axes[1].set_title('Término KL (castigo por alejarse de N(0,I))')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('KL loss')
axes[1].legend()

plt.tight_layout()
plt.show()

# Obtenemos representaciones latentes medias (z_mean) para train/val/test
z_train_mean, _, _ = encoder.predict(x_train_f, batch_size=256, verbose=0)
z_val_mean, _, _ = encoder.predict(x_val_f, batch_size=256, verbose=0)
z_test_mean, _, _ = encoder.predict(x_test_f, batch_size=256, verbose=0)

print(z_train_mean.shape, z_val_mean.shape, z_test_mean.shape)

# Entrenamos KMeans sin etiquetas sobre el espacio latente
kmeans = KMeans(n_clusters=10, random_state=42, n_init=20)
train_clusters = kmeans.fit_predict(z_train_mean)
val_clusters = kmeans.predict(z_val_mean)
test_clusters = kmeans.predict(z_test_mean)

# Mapeo cluster->etiqueta por mayoría en train
cluster_to_label = {}
for c in range(10):
    idx = np.where(train_clusters == c)[0]
    if len(idx) == 0:
        cluster_to_label[c] = 0
    else:
        vals, cnts = np.unique(y_train_small[idx], return_counts=True)
        cluster_to_label[c] = vals[np.argmax(cnts)]

# Predicción final en val/test usando el mapeo aprendido
val_pred = np.array([cluster_to_label[c] for c in val_clusters])
test_pred = np.array([cluster_to_label[c] for c in test_clusters])

val_acc = np.mean(val_pred == y_val)
test_acc = np.mean(test_pred == y_test)

val_nmi = normalized_mutual_info_score(y_val, val_clusters)
test_nmi = normalized_mutual_info_score(y_test, test_clusters)

val_ari = adjusted_rand_score(y_val, val_clusters)
test_ari = adjusted_rand_score(y_test, test_clusters)

print(f"Val  ACC (mapeo mayoría): {val_acc:.4f}")
print(f"Test ACC (mapeo mayoría): {test_acc:.4f}")
print(f"Val  NMI: {val_nmi:.4f} | Test NMI: {test_nmi:.4f}")
print(f"Val  ARI: {val_ari:.4f} | Test ARI: {test_ari:.4f}")

# Visualización del espacio latente 2D con color por clase real (solo para análisis)
plt.figure(figsize=(8,6))
scatter = plt.scatter(z_test_mean[:,0], z_test_mean[:,1], c=y_test, s=8, cmap='tab10', alpha=0.7)
plt.colorbar(scatter, ticks=range(10), label='Etiqueta real')
plt.title('Espacio latente (z_mean) del test set')
plt.xlabel('z[0]')
plt.ylabel('z[1]')
plt.show()

# Matriz de confusión para inspeccionar qué dígitos se confunden más
cm = confusion_matrix(y_test, test_pred)
plt.figure(figsize=(8,6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.title('Matriz de confusión (clasificación no supervisada + mapeo por mayoría)')
plt.xlabel('Predicción')
plt.ylabel('Real')
plt.show()

# Generación de nuevas imágenes muestreando z ~ N(0, I)
n = 12
z_samples = np.random.normal(size=(n, latent_dim)).astype('float32')
generated = decoder.predict(z_samples, verbose=0).reshape(-1, 28, 28)

fig, axes = plt.subplots(2, 6, figsize=(10,4))
for i, ax in enumerate(axes.ravel()):
    ax.imshow(generated[i], cmap='gray')
    ax.axis('off')
plt.suptitle('Nuevos dígitos generados por el VAE')
plt.tight_layout()
plt.show()

# Elegimos dos imágenes reales de test y obtenemos sus z_mean
idx_a, idx_b = 0, 1
x_a = x_test_f[idx_a:idx_a+1]
x_b = x_test_f[idx_b:idx_b+1]

za, _, _ = encoder.predict(x_a, verbose=0)
zb, _, _ = encoder.predict(x_b, verbose=0)

# Interpolación lineal en latent space
steps = 12
alphas = np.linspace(0, 1, steps)
z_interp = np.array([(1-a)*za[0] + a*zb[0] for a in alphas], dtype='float32')

x_interp = decoder.predict(z_interp, verbose=0).reshape(-1, 28, 28)

fig, axes = plt.subplots(1, steps, figsize=(16,2.5))
for i, ax in enumerate(axes):
    ax.imshow(x_interp[i], cmap='gray')
    ax.set_title(f"{alphas[i]:.1f}", fontsize=8)
    ax.axis('off')

plt.suptitle(f"Interpolación latente entre test[{idx_a}] (label={y_test[idx_a]}) y test[{idx_b}] (label={y_test[idx_b]})")
plt.tight_layout()
plt.show()

# Demostración con muestras aleatorias de test
rng = np.random.default_rng(123)
indices = rng.choice(len(x_test_f), size=8, replace=False)

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

for row, idx in enumerate(indices):
    x = x_test_f[idx:idx+1]
    y_true = y_test[idx]

    z_m, _, z_s = encoder.predict(x, verbose=0)
    recon = decoder.predict(z_s, verbose=0).reshape(28, 28)

    cluster = kmeans.predict(z_m)[0]
    y_hat = cluster_to_label[cluster]

    axes[row, 0].imshow(x.reshape(28,28), cmap='gray')
    axes[row, 0].set_title(f"Original\nreal={y_true}")
    axes[row, 0].axis('off')

    axes[row, 1].imshow(recon, cmap='gray')
    axes[row, 1].set_title("Reconstrucción")
    axes[row, 1].axis('off')

    axes[row, 2].text(0.05, 0.8, f"z_mean = {np.round(z_m[0], 2)}", fontsize=10)
    axes[row, 2].text(0.05, 0.55, f"cluster = {cluster}", fontsize=10)
    axes[row, 2].text(0.05, 0.3, f"pred = {y_hat}", fontsize=12, weight='bold')
    axes[row, 2].axis('off')

plt.suptitle("Demostración del agente no supervisado entrenado", y=1.002)
plt.tight_layout()
plt.show()