# Si trabajas en un entorno limpio, descomenta la línea siguiente:
# !pip install -q transformers sentencepiece torch seaborn matplotlib pandas

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

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, set_seed

# Configuración visual y de reproducibilidad
sns.set_theme(style="whitegrid", context="notebook")
set_seed(42)
np.random.seed(42)
torch.manual_seed(42)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Usando dispositivo: {DEVICE}")

MODEL_NAME = "google/mt5-small"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(DEVICE)
model.eval()

print("Modelo y tokenizer cargados correctamente.")
print(f"Vocabulario tokenizer: {tokenizer.vocab_size} tokens")

# Banco de frases en español para análisis didáctico
sample_sentences_es = [
    "La inteligencia artificial está transformando la medicina.",
    "Ayer no pude ir al trabajo porque llovía mucho.",
    "Si estudias con constancia, aprenderás más rápido.",
    "El gato duerme debajo de la mesa mientras yo leo.",
    "¿Podrías explicarme cómo funciona la atención en Transformers?",
]


def build_translation_prompt(spanish_text: str) -> str:
    """Construye el prompt text-to-text para traducción es->en."""
    return f"translate Spanish to English: {spanish_text}"


@torch.no_grad()
def translate_es_en(spanish_text: str, max_new_tokens: int = 64, num_beams: int = 4):
    """Genera traducción con beam search y devuelve texto + tensores de entrada/salida."""
    prompt = build_translation_prompt(spanish_text)
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

    generated_ids = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        num_beams=num_beams,
        early_stopping=True
    )

    translation = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return translation, inputs, generated_ids


rows = []
for s in sample_sentences_es:
    tr, _, _ = translate_es_en(s)
    rows.append({"es": s, "en_pred": tr})

pd.DataFrame(rows)

analysis_sentence = "La inteligencia artificial está transformando la medicina."
prompt = build_translation_prompt(analysis_sentence)

# Tokenizamos la entrada (fuente)
enc_inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

# Generamos traducción para obtener la secuencia objetivo predicha
with torch.no_grad():
    generated_ids = model.generate(**enc_inputs, max_new_tokens=64, num_beams=4)

translation_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print("Entrada (ES):", analysis_sentence)
print("Traducción predicha:", translation_text)

# Para inspección de atención, usamos decoder_input_ids como la salida generada sin el último token
# (teacher forcing parcial para obtener distribuciones por cada paso del decoder)
decoder_input_ids = generated_ids[:, :-1]

with torch.no_grad():
    outputs = model(
        input_ids=enc_inputs["input_ids"],
        attention_mask=enc_inputs["attention_mask"],
        decoder_input_ids=decoder_input_ids,
        output_attentions=True,
        return_dict=True
    )

# Inspección rápida de estructuras
print("# capas encoder attn:", len(outputs.encoder_attentions))
print("# capas decoder self-attn:", len(outputs.decoder_attentions))
print("# capas cross-attn:", len(outputs.cross_attentions))
print("Shape ejemplo encoder_attn capa 0:", tuple(outputs.encoder_attentions[0].shape))
print("Shape ejemplo cross_attn capa 0:", tuple(outputs.cross_attentions[0].shape))

# Extraemos longitudes útiles
src_len = enc_inputs["input_ids"].shape[1]
tgt_len = decoder_input_ids.shape[1]

enc_shape = outputs.encoder_attentions[0].shape
dec_shape = outputs.decoder_attentions[0].shape
cross_shape = outputs.cross_attentions[0].shape

assert enc_shape[2] == src_len and enc_shape[3] == src_len, "Forma inesperada en encoder attention"
assert dec_shape[2] == tgt_len and dec_shape[3] == tgt_len, "Forma inesperada en decoder self-attention"
assert cross_shape[2] == tgt_len and cross_shape[3] == src_len, "Forma inesperada en cross-attention"

print("✅ Shapes de atención validadas correctamente.")

def ids_to_tokens(ids_tensor):
    """Convierte ids (1D tensor) a lista de tokens."""
    return tokenizer.convert_ids_to_tokens(ids_tensor.tolist())


def plot_attention_map(attn_matrix, x_tokens, y_tokens, title, figsize=(11, 7), cmap="magma"):
    """Dibuja heatmap de una matriz de atención 2D."""
    plt.figure(figsize=figsize)
    sns.heatmap(attn_matrix, xticklabels=x_tokens, yticklabels=y_tokens, cmap=cmap)
    plt.xticks(rotation=75, ha="right", fontsize=9)
    plt.yticks(fontsize=9)
    plt.title(title)
    plt.xlabel("Tokens clave/valor")
    plt.ylabel("Tokens consulta")
    plt.tight_layout()
    plt.show()


# Tokens fuente y objetivo
src_tokens = ids_to_tokens(enc_inputs["input_ids"][0].cpu())
tgt_tokens = ids_to_tokens(decoder_input_ids[0].cpu())

print("Tokens fuente (encoder):")
print(src_tokens)
print("\nTokens objetivo parciales (decoder):")
print(tgt_tokens)

layer_idx = 0
head_idx = 0

# Tensor: (batch, heads, src_len, src_len)
enc_attn = outputs.encoder_attentions[layer_idx][0, head_idx].detach().cpu().numpy()

plot_attention_map(
    enc_attn,
    x_tokens=src_tokens,
    y_tokens=src_tokens,
    title=f"Encoder self-attention | capa={layer_idx}, cabeza={head_idx}"
)

layer_idx = 0
head_idx = 0

# Tensor: (batch, heads, tgt_len, tgt_len)
dec_attn = outputs.decoder_attentions[layer_idx][0, head_idx].detach().cpu().numpy()

plot_attention_map(
    dec_attn,
    x_tokens=tgt_tokens,
    y_tokens=tgt_tokens,
    title=f"Decoder self-attention (masked) | capa={layer_idx}, cabeza={head_idx}",
    cmap="viridis"
)

layer_idx = 0
head_idx = 0

# Tensor: (batch, heads, tgt_len, src_len)
cross_attn = outputs.cross_attentions[layer_idx][0, head_idx].detach().cpu().numpy()

plot_attention_map(
    cross_attn,
    x_tokens=src_tokens,
    y_tokens=tgt_tokens,
    title=f"Cross-attention | capa={layer_idx}, cabeza={head_idx}",
    cmap="rocket_r"
)

def attention_entropy(prob_vector, eps=1e-12):
    """Entropía de Shannon de una distribución discreta."""
    p = np.clip(prob_vector, eps, 1.0)
    return -(p * np.log(p)).sum()


records = []

# Recorremos cross-attention para medir concentración media por cabeza
for layer_i, layer_tensor in enumerate(outputs.cross_attentions):
    # layer_tensor: (batch, heads, tgt_len, src_len)
    mat = layer_tensor[0].detach().cpu().numpy()
    n_heads = mat.shape[0]

    for head_i in range(n_heads):
        # Promedio de entropía por filas (tokens objetivo)
        row_entropies = [attention_entropy(row / row.sum()) for row in mat[head_i]]
        records.append({
            "layer": layer_i,
            "head": head_i,
            "mean_row_entropy": float(np.mean(row_entropies)),
        })

entropy_df = pd.DataFrame(records)
entropy_pivot = entropy_df.pivot(index="layer", columns="head", values="mean_row_entropy")

plt.figure(figsize=(10, 6))
sns.heatmap(entropy_pivot, annot=True, fmt=".2f", cmap="YlGnBu")
plt.title("Entropía media de cross-attention por capa/cabeza\n(menor = atención más concentrada)")
plt.xlabel("Cabeza")
plt.ylabel("Capa")
plt.tight_layout()
plt.show()

entropy_df.head()

@torch.no_grad()
def average_cross_attention_alignment(spanish_text: str, max_new_tokens: int = 64):
    """Devuelve traducción y alineaciones argmax usando cross-attention promediada."""
    prompt = build_translation_prompt(spanish_text)
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

    gen_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, num_beams=4)
    dec_ids = gen_ids[:, :-1]

    out = model(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        decoder_input_ids=dec_ids,
        output_attentions=True,
        return_dict=True,
    )

    src_toks = ids_to_tokens(inputs["input_ids"][0].cpu())
    tgt_toks = ids_to_tokens(dec_ids[0].cpu())

    # Promedio sobre capas y cabezas: (tgt_len, src_len)
    cross_stack = torch.stack(out.cross_attentions, dim=0)  # (layers, batch, heads, tgt, src)
    avg_cross = cross_stack[:, 0].mean(dim=(0, 1)).cpu().numpy()  # media en layers y heads

    best_src_idx = avg_cross.argmax(axis=1)
    alignments = [(tgt_toks[i], src_toks[j], float(avg_cross[i, j])) for i, j in enumerate(best_src_idx)]

    translation = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
    return translation, src_toks, tgt_toks, avg_cross, alignments


test_sentences = [
    "El profesor explicó el tema con ejemplos sencillos.",
    "Mañana viajaremos a Londres si el clima mejora.",
    "No me gusta llegar tarde a las reuniones.",
]

for sent in test_sentences:
    tr, src_toks, tgt_toks, avg_cross, aligns = average_cross_attention_alignment(sent)
    print("\n" + "=" * 100)
    print("ES:", sent)
    print("EN:", tr)
    print("Alineaciones aproximadas (token_target -> token_source):")
    for tgt_tok, src_tok, score in aligns[: min(len(aligns), 18)]:
        print(f"  {tgt_tok:>12} -> {src_tok:<18} (score={score:.3f})")

# Verificamos en la cross-attention promedio del análisis principal
cross_stack = torch.stack(outputs.cross_attentions, dim=0)  # (layers, batch, heads, tgt, src)
avg_cross_main = cross_stack[:, 0].mean(dim=(0, 1)).cpu().numpy()  # (tgt, src)

row_sums = avg_cross_main.sum(axis=1)
max_dev = np.max(np.abs(row_sums - 1.0))

print("Suma por fila (primeros 10):", np.round(row_sums[:10], 6))
print("Desviación máxima respecto a 1.0:", float(max_dev))

assert max_dev < 1e-3, "Las filas no parecen distribuciones válidas (revisar extracción de atención)."
print("✅ Test de normalización de atención superado.")