# ============================================================
# 0) Imports y configuración global
# ============================================================

import csv
import io
import os
import random
import re
import time
from collections import Counter, OrderedDict
from pathlib import Path

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 Dataset, DataLoader

from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
)

sns.set_theme(style="whitegrid")
plt.rcParams["figure.figsize"] = (10, 4)


# ---------- Reemplazos ligeros de torchtext (deprecado) ----------

def basic_english_tokenizer(text: str) -> list[str]:
    """Tokenizador básico: minúsculas, separa puntuación, divide por espacios."""
    text = text.lower()
    text = re.sub(r"([.!?,;:\"'()\[\]{}\-/])", r" \1 ", text)
    return text.split()

tokenizer = basic_english_tokenizer


class SimpleVocab:
    """Vocabulario ligero compatible con la interfaz usada en el notebook."""

    def __init__(self, counter: Counter, min_freq: int = 1, specials: list[str] | None = None):
        self.itos = list(specials or [])
        self.stoi = {tok: i for i, tok in enumerate(self.itos)}
        for tok, freq in counter.most_common():
            if freq >= min_freq and tok not in self.stoi:
                self.stoi[tok] = len(self.itos)
                self.itos.append(tok)
        self._default_index = 0

    def set_default_index(self, idx: int):
        self._default_index = idx

    def __getitem__(self, token: str) -> int:
        return self.stoi.get(token, self._default_index)

    def __len__(self) -> int:
        return len(self.itos)


def download_ag_news(root: str = "./data/ag_news"):
    """Descarga AG News CSV si no existe y devuelve rutas (train, test)."""
    root = Path(root)
    root.mkdir(parents=True, exist_ok=True)
    base_url = "https://raw.githubusercontent.com/mhjabreel/CharCnn_Keras/master/data/ag_news_csv"
    paths = {}
    for split in ("train", "test"):
        fpath = root / f"{split}.csv"
        if not fpath.exists():
            import urllib.request
            url = f"{base_url}/{split}.csv"
            print(f"Descargando {url} ...")
            urllib.request.urlretrieve(url, fpath)
            print(f"  -> {fpath}")
        paths[split] = fpath
    return paths["train"], paths["test"]


def load_ag_news_csv(csv_path: str) -> list[tuple[int, str]]:
    """Lee CSV de AG News y devuelve lista de (label_0indexed, text)."""
    data = []
    with open(csv_path, "r", encoding="utf-8") as f:
        reader = csv.reader(f)
        for row in reader:
            label = int(row[0]) - 1  # 1-4 → 0-3
            text = row[1] + " " + row[2]  # title + description
            data.append((label, text))
    return data


# ---------- Reproducibilidad ----------

def set_seed(seed: int = 42):
    """Fija semillas para reproducibilidad."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


set_seed(42)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Dispositivo:", device)

# ============================================================
# 1) Carga de AG_NEWS
# ============================================================

LABELS = ["World", "Sports", "Business", "Sci/Tech"]

train_csv, test_csv = download_ag_news()

train_raw = load_ag_news_csv(train_csv)
test_raw = load_ag_news_csv(test_csv)

print(f"Train size: {len(train_raw):,}")
print(f"Test  size: {len(test_raw):,}")
print("Ejemplo:")
print(train_raw[0][0], "->", train_raw[0][1][:140], "...")

# ============================================================
# 2.1 Distribución de clases
# ============================================================

train_labels = [y for y, _ in train_raw]
class_counts = Counter(train_labels)

plt.figure(figsize=(7, 4))
plt.bar([LABELS[i] for i in range(4)], [class_counts[i] for i in range(4)], color=["#4c78a8", "#f58518", "#54a24b", "#e45756"])
plt.title("Distribución de clases en train")
plt.ylabel("Número de titulares")
plt.xticks(rotation=15)
plt.show()

print("Conteos por clase:")
for i, c in class_counts.items():
    print(f"  {LABELS[i]}: {c:,}")

# ============================================================
# 2.2 Longitud de secuencias (número de tokens por titular)
# ============================================================

train_lengths = [len(tokenizer(text)) for _, text in train_raw]

plt.figure(figsize=(8, 4))
plt.hist(train_lengths, bins=50, color="#4c78a8", alpha=0.8)
plt.title("Distribución de longitudes de titulares (tokens)")
plt.xlabel("Longitud")
plt.ylabel("Frecuencia")
plt.show()

for q in [50, 75, 90, 95, 99]:
    print(f"Percentil {q:>2}: {np.percentile(train_lengths, q):.1f} tokens")

# ============================================================
# 2.3 Palabras más frecuentes (aprox.)
# ============================================================

word_counter = Counter()
for _, text in train_raw[:20000]:  # muestreo para acelerar el conteo
    word_counter.update(tokenizer(text))

print("Top 20 tokens más frecuentes:")
print(word_counter.most_common(20))

# ============================================================
# 3.1 Construcción de vocabulario
# ============================================================

# Para demo rápida puedes limitar la cantidad de ejemplos de entrenamiento
MAX_TRAIN_SAMPLES = 50000  # sube a 120000 para usar todo AG News

train_used = train_raw[:MAX_TRAIN_SAMPLES]

specials = ["<pad>", "<unk>"]

counter = Counter()
for _, text in train_used:
    counter.update(tokenizer(text))

# min_freq evita vocabulario enorme y reduce ruido
tok_vocab = SimpleVocab(counter, min_freq=2, specials=specials)
tok_vocab.set_default_index(tok_vocab["<unk>"])

PAD_IDX = tok_vocab["<pad>"]
UNK_IDX = tok_vocab["<unk>"]
VOCAB_SIZE = len(tok_vocab)

print("Vocab size:", VOCAB_SIZE)
print("PAD_IDX:", PAD_IDX, "UNK_IDX:", UNK_IDX)

# ============================================================
# 3.2 Dataset PyTorch: tokenización + numerización + padding
# ============================================================

MAX_LEN = 35  # cercano al p95 observado normalmente en titulares AG News


def encode_text(text: str, max_len: int = MAX_LEN):
    """Tokeniza, convierte a ids y aplica truncado/padding."""
    ids = [tok_vocab[token] for token in tokenizer(text)]
    ids = ids[:max_len]
    if len(ids) < max_len:
        ids += [PAD_IDX] * (max_len - len(ids))
    return ids


class AGNewsDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        label, text = self.data[idx]
        x = torch.tensor(encode_text(text), dtype=torch.long)
        y = torch.tensor(label, dtype=torch.long)
        return x, y


# Split train/val manual (reproducible)
val_ratio = 0.15
n_total = len(train_used)
indices = np.arange(n_total)
np.random.shuffle(indices)

n_val = int(n_total * val_ratio)
val_idx = indices[:n_val]
train_idx = indices[n_val:]

train_split = [train_used[i] for i in train_idx]
val_split = [train_used[i] for i in val_idx]

train_ds = AGNewsDataset(train_split)
val_ds = AGNewsDataset(val_split)

# Test completo (se podría submuestrear en entornos lentos)
test_ds = AGNewsDataset(test_raw)

BATCH_SIZE = 128

train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False)
test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False)

print(f"Train batches: {len(train_loader)}")
print(f"Val batches:   {len(val_loader)}")
print(f"Test batches:  {len(test_loader)}")

# ============================================================
# 4) Definición del modelo GRU
# ============================================================

class GRUClassifier(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int,
        hidden_dim: int,
        num_classes: int,
        num_layers: int = 2,
        dropout: float = 0.3,
        bidirectional: bool = True,
        pad_idx: int = 0,
    ):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        self.gru = nn.GRU(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0,
            bidirectional=bidirectional,
        )

        self.dropout = nn.Dropout(dropout)
        self.bidirectional = bidirectional
        direction_factor = 2 if bidirectional else 1
        self.classifier = nn.Linear(hidden_dim * direction_factor, num_classes)

    def forward(self, x):
        # x: (batch, seq_len)
        emb = self.dropout(self.embedding(x))  # (batch, seq_len, embed_dim)

        # GRU devuelve output y h_n
        _, h_n = self.gru(emb)

        # h_n shape: (num_layers * num_directions, batch, hidden_dim)
        if self.bidirectional:
            # Última capa forward y backward
            h_last = torch.cat([h_n[-2], h_n[-1]], dim=1)
        else:
            h_last = h_n[-1]

        logits = self.classifier(self.dropout(h_last))
        return logits


model = GRUClassifier(
    vocab_size=VOCAB_SIZE,
    embed_dim=128,
    hidden_dim=128,
    num_classes=4,
    num_layers=2,
    dropout=0.3,
    bidirectional=True,
    pad_idx=PAD_IDX,
).to(device)

n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(model)
print(f"Parámetros entrenables: {n_params:,}")

# ============================================================
# 5) Bucle de entrenamiento y validación
# ============================================================

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)


def run_epoch(model, loader, optimizer=None, clip_grad=1.0):
    """Ejecuta una época de train o validación.
    Si optimizer es None, se ejecuta en modo evaluación.
    """
    is_train = optimizer is not None
    model.train() if is_train else model.eval()

    losses, preds_all, y_all = [], [], []

    for x_batch, y_batch in loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        with torch.set_grad_enabled(is_train):
            logits = model(x_batch)
            loss = criterion(logits, y_batch)

            if is_train:
                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip_grad)
                optimizer.step()

        losses.append(loss.item())
        preds_all.extend(torch.argmax(logits, dim=1).detach().cpu().numpy())
        y_all.extend(y_batch.detach().cpu().numpy())

    epoch_loss = float(np.mean(losses))
    epoch_acc = accuracy_score(y_all, preds_all)
    return epoch_loss, epoch_acc


EPOCHS = 12
PATIENCE = 3

history = {
    "train_loss": [],
    "val_loss": [],
    "train_acc": [],
    "val_acc": [],
}

best_val_loss = float("inf")
best_state = None
patience_counter = 0

start_time = time.time()

for epoch in range(1, EPOCHS + 1):
    train_loss, train_acc = run_epoch(model, train_loader, optimizer=optimizer)
    val_loss, val_acc = run_epoch(model, val_loader, optimizer=None)

    history["train_loss"].append(train_loss)
    history["val_loss"].append(val_loss)
    history["train_acc"].append(train_acc)
    history["val_acc"].append(val_acc)

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

    # Early stopping por val_loss
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
        patience_counter = 0
    else:
        patience_counter += 1
        if patience_counter >= PATIENCE:
            print("Early stopping activado.")
            break

elapsed = time.time() - start_time
print(f"Tiempo total de entrenamiento: {elapsed/60:.1f} min")

# Restauramos mejor estado validación
if best_state is not None:
    model.load_state_dict(best_state)
    print("Modelo restaurado al mejor checkpoint de validación.")

# ============================================================
# 5.1 Curvas de entrenamiento: loss y accuracy
# ============================================================

epoch_axis = np.arange(1, len(history["train_loss"]) + 1)

fig, ax = plt.subplots(1, 2, figsize=(13, 4))

ax[0].plot(epoch_axis, history["train_loss"], marker="o", label="train")
ax[0].plot(epoch_axis, history["val_loss"], marker="o", label="val")
ax[0].set_title("Loss vs Epoch")
ax[0].set_xlabel("Epoch")
ax[0].set_ylabel("Cross-Entropy Loss")
ax[0].legend()

ax[1].plot(epoch_axis, history["train_acc"], marker="o", label="train")
ax[1].plot(epoch_axis, history["val_acc"], marker="o", label="val")
ax[1].set_title("Accuracy vs Epoch")
ax[1].set_xlabel("Epoch")
ax[1].set_ylabel("Accuracy")
ax[1].legend()

plt.tight_layout()
plt.show()

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

model.eval()
all_preds, all_targets = [], []

with torch.no_grad():
    for x_batch, y_batch in test_loader:
        x_batch = x_batch.to(device)
        logits = model(x_batch)
        preds = torch.argmax(logits, dim=1).cpu().numpy()

        all_preds.extend(preds)
        all_targets.extend(y_batch.numpy())

test_acc = accuracy_score(all_targets, all_preds)
test_f1_macro = f1_score(all_targets, all_preds, average="macro")

print(f"Test Accuracy : {test_acc:.4f}")
print(f"Test Macro-F1 : {test_f1_macro:.4f}")

print("Classification report:")
print(classification_report(all_targets, all_preds, target_names=LABELS, digits=4))

# ============================================================
# 6.1 Matriz de confusión
# ============================================================

cm = confusion_matrix(all_targets, all_preds)

plt.figure(figsize=(6, 5))
sns.heatmap(
    cm,
    annot=True,
    fmt="d",
    cmap="Blues",
    xticklabels=LABELS,
    yticklabels=LABELS,
)
plt.title("Matriz de confusión (test)")
plt.xlabel("Predicción")
plt.ylabel("Etiqueta real")
plt.tight_layout()
plt.show()

# ============================================================
# 7) Mostrar ejemplos acertados/erróneos
# ============================================================

# Construimos predicciones sobre un subconjunto del test con texto original
sample_size = 300
sample_data = test_raw[:sample_size]

x_sample = torch.tensor([encode_text(text) for _, text in sample_data], dtype=torch.long).to(device)
y_sample = np.array([label for label, _ in sample_data])
texts_sample = [text for _, text in sample_data]

with torch.no_grad():
    logits = model(x_sample)
    pred_sample = torch.argmax(logits, dim=1).cpu().numpy()

correct_idx = np.where(pred_sample == y_sample)[0]
wrong_idx = np.where(pred_sample != y_sample)[0]

print(f"Aciertos en muestra: {len(correct_idx)}/{sample_size}")
print(f"Errores  en muestra: {len(wrong_idx)}/{sample_size}")

print("--- Ejemplos de aciertos ---")
for i in correct_idx[:3]:
    print(f"Real={LABELS[y_sample[i]]} | Pred={LABELS[pred_sample[i]]}")
    print(texts_sample[i][:180], "\n")

print("--- Ejemplos de errores ---")
for i in wrong_idx[:3]:
    print(f"Real={LABELS[y_sample[i]]} | Pred={LABELS[pred_sample[i]]}")
    print(texts_sample[i][:180], "\n")

# ============================================================
# 8) Experimento opcional (rápido): cambiar bidirectional=False
# ============================================================

# Esta celda está pensada como ejercicio.
# Para ahorrar tiempo por defecto no se ejecuta entrenamiento completo aquí.
# Puedes descomentar y lanzar una versión corta (3-4 épocas).

# quick_model = GRUClassifier(
#     vocab_size=VOCAB_SIZE,
#     embed_dim=128,
#     hidden_dim=128,
#     num_classes=4,
#     num_layers=2,
#     dropout=0.3,
#     bidirectional=False,
#     pad_idx=PAD_IDX,
# ).to(device)
#
# quick_opt = torch.optim.Adam(quick_model.parameters(), lr=1e-3)
# for e in range(4):
#     tr_l, tr_a = run_epoch(quick_model, train_loader, optimizer=quick_opt)
#     va_l, va_a = run_epoch(quick_model, val_loader, optimizer=None)
#     print(f"[UNI-GRU] epoch={e+1} tr_acc={tr_a:.4f} val_acc={va_a:.4f}")