from datasets import load_dataset, DatasetDict, Dataset
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    Trainer, TrainingArguments, DataCollatorWithPadding
)
import numpy as np
from utils import compute_metrics  
import os

def load_ag_news():
    """
    Charge le jeu de données AG News via Hugging Face datasets.
    
    Returns:
        DatasetDict: Contenant les splits train/test.
    """
    dataset = load_dataset("ag_news")
    return dataset


def get_balanced_subset(dataset_split, n_per_class=1000):
    """
    Crée un sous-ensemble équilibré contenant `n_per_class` exemples par classe.
    
    Args:
        dataset_split (Dataset): Split de type train ou test.
        n_per_class (int): Nombre d'exemples à garder par classe.
    
    Returns:
        Dataset: Sous-ensemble équilibré.
    """
    subsets = []

    for label in range(4):
        #Filtrage des exemples correspondant à la classe `label`
        filtered = dataset_split.filter(lambda example: example['label'] == label)
        #Sélection des n premiers exemples (ou tous s’il y en a moins)
        subsets.append(filtered.select(range(min(n_per_class, len(filtered)))))

    #Fusionner les sous-ensembles
    combined_dict = {
        key: sum([subset[key] for subset in subsets], []) for key in subsets[0].features.keys()
    }

    return Dataset.from_dict(combined_dict)


def preprocess_data(dataset, tokenizer):
    """
    Tokenise le jeu de données avec troncature et padding.

    Args:
        dataset (DatasetDict): Données d'entraînement et de test.
        tokenizer (AutoTokenizer): Tokenizer à utiliser.

    Returns:
        DatasetDict: Données tokenisées.
    """
    def preprocess(batch):
        return tokenizer(batch["text"], truncation=True, padding=True)

    return dataset.map(preprocess, batched=True)


def main():
    """
    Lance le fine-tuning du modèle BERT sur AG News et sauvegarde le modèle.
    """
    #Création des dossiers de sortie
    os.makedirs("../models/fine_tuned_model", exist_ok=True)
    os.makedirs("../logs", exist_ok=True)

    #Chargement du jeu de données
    dataset = load_ag_news()

    #Création de sous-ensembles équilibrés (entraînement/test)
    train_subset = get_balanced_subset(dataset["train"], n_per_class=3000)
    test_subset = get_balanced_subset(dataset["test"], n_per_class=1000)

    dataset_small = DatasetDict({
        "train": train_subset,
        "test": test_subset
    })

    #Chargement du tokenizer
    model_name = "bert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    #Prétraitement (tokenisation)
    encoded = preprocess_data(dataset_small, tokenizer)

    #Préparation des entrées avec padding dynamique
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    #Chargement du modèle BERT pour classification
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=4  #AG News contient 4 classes
    )

    #Configuration de l'entraînement
    training_args = TrainingArguments(
        output_dir="../models/fine_tuned_model",
        eval_strategy="epoch",
        save_strategy="epoch",
        num_train_epochs=3,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,
        load_best_model_at_end=True,
        metric_for_best_model="accuracy",
        logging_dir="../logs",
        seed=42
    )

    #Définition du trainer Hugging Face
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=encoded["train"],
        eval_dataset=encoded["test"],
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=lambda p: compute_metrics(
            np.argmax(p.predictions, axis=1), p.label_ids
        )
    )

    #Lancement de l'entraînement
    trainer.train()

    #Sauvegarde finale du modèle
    trainer.save_model("../models/fine_tuned_model")
    print("✅ Modèle sauvegardé dans ../models/fine_tuned_model")


if __name__ == "__main__":
    main()