import json
import numpy as np
from datasets import load_dataset
from sklearn.metrics import accuracy_score, f1_score
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from src.labels import LABELS

# Modelo base en español (BETO)
BASE_MODEL = "dccuchile/bert-base-spanish-wwm-cased"
MODEL_DIR  = "out_emotions"

# Mapeos etiqueta <-> id
label2id = {lab: i for i, lab in enumerate(LABELS)}
id2label = {i: lab for i, lab in enumerate(LABELS)}

def compute_metrics(p):
    preds = np.argmax(p.predictions, axis=1)
    labels = p.label_ids
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average="weighted")
    return {"accuracy": acc, "f1": f1}

def main():
    # Carga dataset JSONL
    dataset = load_dataset("json", data_files={
        "train": "data/train.jsonl",
        "validation": "data/valid.jsonl",
        "test": "data/test.jsonl",
    })

    # Tokenizador
    tok = AutoTokenizer.from_pretrained(BASE_MODEL)

    # Tokenización
    def encode(batch):
        return tok(batch["text"], truncation=True, padding="max_length", max_length=128)

    dataset = dataset.map(encode, batched=True)

    # Renombra 'label' -> 'labels' y convierte string -> id (entero)
    dataset = dataset.rename_column("label", "labels")

    def to_ids(ex):
        ex["labels"] = int(label2id[ex["labels"]])  # 'entusiasta' -> 2, etc.
        return ex

    dataset = dataset.map(to_ids)

    # Formato tensores
    dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

    # Modelo con metadatos de etiquetas
    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=len(LABELS),
        id2label=id2label,
        label2id=label2id,
    )

    # Args de entrenamiento (4.56 usa 'eval_strategy'; 'evaluation_strategy' quedará deprecado)
    args = TrainingArguments(
        output_dir=MODEL_DIR,
        eval_strategy="epoch",
        save_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=8,
        num_train_epochs=5,
        weight_decay=0.01,
        logging_dir="./logs",
        logging_steps=50,
        load_best_model_at_end=True,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        tokenizer=tok,               # aviso deprecado, pero sigue funcionando
        compute_metrics=compute_metrics,
    )

    trainer.train()
    # Guarda mejor modelo + tokenizer en out_emotions/
    trainer.save_model(MODEL_DIR)
    tok.save_pretrained(MODEL_DIR)

if __name__ == "__main__":
    main()