# ==============================
# 1) Imports 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 random
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.metrics import (
    confusion_matrix,
    classification_report,
    roc_curve,
    auc,
    precision_recall_fscore_support,
)

import tensorflow as tf
from tensorflow import keras

# Semillas para reproducibilidad (hasta donde permite el backend).
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
tf.random.set_seed(SEED)

print('TensorFlow version:', tf.__version__)

# ==============================
# 2) Carga de IMDb
# ==============================

TOP_WORDS = 10_000
MAX_LEN = 200

# Cargamos dataset ya tokenizado (enteros).
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=TOP_WORDS)

print('Número de muestras train:', len(x_train))
print('Número de muestras test :', len(x_test))
print('Ejemplo secuencia (primeros 20 tokens):', x_train[0][:20])
print('Etiqueta ejemplo:', y_train[0])

# ==============================
# 3) EDA - Balance de clases
# ==============================

unique, counts = np.unique(y_train, return_counts=True)
class_dist = dict(zip(unique, counts))
print('Distribución de clases (train):', class_dist)

plt.figure(figsize=(5, 3))
plt.bar(['Negativa (0)', 'Positiva (1)'], [class_dist.get(0, 0), class_dist.get(1, 0)])
plt.title('Balance de clases en train')
plt.ylabel('Número de reseñas')
plt.show()

# ==============================
# 4) EDA - Longitudes de secuencia
# ==============================

train_lengths = [len(seq) for seq in x_train]

print('Longitud media:', np.mean(train_lengths).round(2))
print('Mediana      :', np.median(train_lengths))
print('P90          :', np.percentile(train_lengths, 90).round(2))
print('Máxima       :', np.max(train_lengths))

plt.figure(figsize=(7, 4))
plt.hist(train_lengths, bins=50, color='#E17055', edgecolor='black', alpha=0.8)
plt.axvline(MAX_LEN, color='blue', linestyle='--', label=f'MAX_LEN={MAX_LEN}')
plt.title('Distribución de longitud de reseñas (tokens)')
plt.xlabel('Longitud')
plt.ylabel('Frecuencia')
plt.legend()
plt.show()

# ==============================
# 5) EDA - Decodificar una reseña
# ==============================

# IMDb incluye un mapeo palabra -> índice. Lo invertimos para inspeccionar texto.
word_index = keras.datasets.imdb.get_word_index()
index_to_word = {idx + 3: word for word, idx in word_index.items()}
index_to_word[0] = '<PAD>'
index_to_word[1] = '<START>'
index_to_word[2] = '<UNK>'
index_to_word[3] = '<UNUSED>'

# Reconstruimos parcialmente una reseña de ejemplo.
example_decoded = ' '.join(index_to_word.get(token, '?') for token in x_train[0][:80])
print('Reseña reconstruida (fragmento):')
print(example_decoded)
print('Etiqueta real:', y_train[0], '(1=positiva, 0=negativa)')

# ==============================
# 6) Padding y tipos
# ==============================

x_train_pad = keras.preprocessing.sequence.pad_sequences(
    x_train,
    maxlen=MAX_LEN,
    padding='post',
    truncating='post'
)

x_test_pad = keras.preprocessing.sequence.pad_sequences(
    x_test,
    maxlen=MAX_LEN,
    padding='post',
    truncating='post'
)

y_train = np.array(y_train, dtype=np.int32)
y_test = np.array(y_test, dtype=np.int32)

print('Shape x_train_pad:', x_train_pad.shape)
print('Shape x_test_pad :', x_test_pad.shape)

# ==============================
# 7) Modelo GRU en Keras
# ==============================

EMBED_DIM = 128
GRU_UNITS_1 = 64
GRU_UNITS_2 = 32

model_gru = keras.Sequential([
    keras.layers.Embedding(input_dim=TOP_WORDS, output_dim=EMBED_DIM, input_length=MAX_LEN),
    keras.layers.Bidirectional(
        keras.layers.GRU(GRU_UNITS_1, return_sequences=True, dropout=0.2)
    ),
    keras.layers.GRU(GRU_UNITS_2, dropout=0.2),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dropout(0.3),
    keras.layers.Dense(1, activation='sigmoid')
])

model_gru.compile(
    optimizer=keras.optimizers.Adam(learning_rate=1e-3),
    loss='binary_crossentropy',
    metrics=[
        'accuracy',
        keras.metrics.Precision(name='precision'),
        keras.metrics.Recall(name='recall'),
        keras.metrics.AUC(name='auc')
    ]
)

model_gru.summary()

# ==============================
# 8) Comparación de parámetros con LSTM equivalente
# ==============================

model_lstm_ref = keras.Sequential([
    keras.layers.Embedding(input_dim=TOP_WORDS, output_dim=EMBED_DIM, input_length=MAX_LEN),
    keras.layers.Bidirectional(
        keras.layers.LSTM(GRU_UNITS_1, return_sequences=True, dropout=0.2)
    ),
    keras.layers.LSTM(GRU_UNITS_2, dropout=0.2),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dropout(0.3),
    keras.layers.Dense(1, activation='sigmoid')
])

model_gru.build(input_shape=(None, MAX_LEN))
model_lstm_ref.build(input_shape=(None, MAX_LEN))

gru_params = model_gru.count_params()
lstm_params = model_lstm_ref.count_params()

print(f'Parámetros GRU : {gru_params:,}')
print(f'Parámetros LSTM: {lstm_params:,}')
print(f'Reducción relativa aprox. con GRU: {(1 - gru_params/lstm_params)*100:.2f}%')

# ==============================
# 9) Callbacks y entrenamiento
# ==============================

BATCH_SIZE = 128
EPOCHS = 8

callbacks = [
    keras.callbacks.EarlyStopping(monitor='val_loss', patience=2, restore_best_weights=True),
    keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss', factor=0.5, patience=1, min_lr=1e-5, verbose=1
    )
]

history = model_gru.fit(
    x_train_pad,
    y_train,
    validation_split=0.2,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    callbacks=callbacks,
    verbose=1
)

# ==============================
# 10) Curvas de entrenamiento
# ==============================

hist = history.history
epochs_r = range(1, len(hist['loss']) + 1)

plt.figure(figsize=(13, 4))

plt.subplot(1, 2, 1)
plt.plot(epochs_r, hist['loss'], marker='o', label='Train loss')
plt.plot(epochs_r, hist['val_loss'], marker='o', label='Val loss')
plt.title('Loss vs Epoch')
plt.xlabel('Epoch')
plt.ylabel('Binary crossentropy')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(epochs_r, hist['accuracy'], marker='o', label='Train acc')
plt.plot(epochs_r, hist['val_accuracy'], marker='o', label='Val acc')
plt.title('Accuracy vs Epoch')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

plt.tight_layout()
plt.show()

# ==============================
# 11) Evaluación final en test
# ==============================

test_metrics = model_gru.evaluate(x_test_pad, y_test, verbose=0)
metric_names = model_gru.metrics_names

print('Métricas en test:')
for name, value in zip(metric_names, test_metrics):
    print(f'  {name:>10s}: {value:.4f}')

# ==============================
# 12) Predicciones y métricas detalladas
# ==============================

y_prob = model_gru.predict(x_test_pad, verbose=0).ravel()
y_pred = (y_prob >= 0.5).astype(int)

cm = confusion_matrix(y_test, y_pred)

plt.figure(figsize=(5, 4))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
plt.title('Matriz de confusión (test)')
plt.xlabel('Predicción')
plt.ylabel('Real')
plt.show()

print(classification_report(y_test, y_pred, digits=4))

precision, recall, f1, _ = precision_recall_fscore_support(y_test, y_pred, average='binary')
print(f'Precision: {precision:.4f}')
print(f'Recall   : {recall:.4f}')
print(f'F1-score : {f1:.4f}')

# ==============================
# 13) Curva ROC
# ==============================

fpr, tpr, _ = roc_curve(y_test, y_prob)
roc_auc = auc(fpr, tpr)

plt.figure(figsize=(6, 5))
plt.plot(fpr, tpr, label=f'ROC (AUC={roc_auc:.4f})')
plt.plot([0, 1], [0, 1], linestyle='--', label='Azar')
plt.title('Curva ROC - Modelo GRU')
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.legend()
plt.grid(alpha=0.3)
plt.show()

# ==============================
# 14) Ejemplos de predicción
# ==============================

def decode_review(token_ids, index_to_word_map, max_tokens=60):
    # Decodifica una secuencia de índices a texto legible (recortado).
    return ' '.join(index_to_word_map.get(t, '?') for t in token_ids[:max_tokens])

sample_idx = np.random.choice(len(x_test), size=5, replace=False)

for i, idx in enumerate(sample_idx, start=1):
    prob = y_prob[idx]
    pred = int(prob >= 0.5)
    real = int(y_test[idx])
    text_preview = decode_review(x_test[idx], index_to_word)

    print(f'--- Ejemplo {i} ---')
    print(f'Real={real} | Pred={pred} | Prob positiva={prob:.4f}')
    print('Texto (fragmento):', text_preview)
    print()