import os
import torch
import pandas as pd
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding
)

# 1. Configuration
# Pour le Français, on utilise un modèle multilingue (XLM-RoBERTa) ou dédié (CamemBERT)
MODEL_NAME = "xlm-roberta-base" 
OUTPUT_DIR = "./results"
SAVE_DIR = "./fakenews_model_fr"

def train():
    # 2. Chargement des données (Exemple avec un CSV fictif)
    # Assurez-vous d'avoir un CSV avec au moins 'text' et 'label' (0=Real, 1=Fake)
    if not os.path.exists("data.csv"):
        print("Erreur: data.csv non trouvé. Création d'un exemple factice...")
        data = {
            "text": [
                "Le président a inauguré une nouvelle école ce matin à Paris.", 
                "Des extraterrestres ont envahi la tour Eiffel hier soir !"
            ],
            "label": [0, 1]
        }
        pd.DataFrame(data).to_csv("data.csv", index=False)

    df = pd.read_csv("data.csv")
    
    # Nettoyage des données : Supprimer les NaNs et s'assurer que 'text' est de type string
    df = df.dropna(subset=['text', 'label'])
    df['text'] = df['text'].astype(str)
    df = df[df['text'].str.len() > 10]

    # OPTIMISATION HACKATHON 🚀 : Échantillon aléatoire pour aller plus vite
    MAX_SAMPLES = 2000
    if len(df) > MAX_SAMPLES:
        print(f"Extraction d'un échantillon aléatoire de {MAX_SAMPLES} sur {len(df)} lignes...")
        df = df.sample(MAX_SAMPLES, random_state=42).reset_index(drop=True)

    dataset = Dataset.from_pandas(df.reset_index(drop=True))

    # 3. Tokenization
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    def preprocess_function(examples):
        return tokenizer(examples["text"], truncation=True, padding=True)

    tokenized_dataset = dataset.map(preprocess_function, batched=True)

    # 4. Chargement du modèle
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)

    # 5. Arguments d'entraînement
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        learning_rate=2e-5,
        per_device_train_batch_size=8, # Réduit pour économiser la RAM
        num_train_epochs=1, # 1 seule passe pour le hackathon
        weight_decay=0.01,
        save_strategy="epoch",
    )

    # 6. Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        tokenizer=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
    )

    # 7. Lancement de l'entraînement
    print("Début de l'entraînement...")
    trainer.train()

    # 8. Sauvegarde
    model.save_pretrained(SAVE_DIR)
    tokenizer.save_pretrained(SAVE_DIR)
    print(f"Modèle sauvegardé dans {SAVE_DIR}")

if __name__ == "__main__":
    train()