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

# 1. Configuration pour Colab / GPU
MODEL_NAME = "xlm-roberta-base" # Multilingue pour supporter le Français
OUTPUT_DIR = "./ai_results"
SAVE_DIR = "./ai_detector_model"

def train():
    # 2. Chargement des données
    # Sur Colab, vous pouvez charger directement depuis Hugging Face :
    # Exemple avec HC3 (Human-ChatGPT Comparison Corpus) en Français
    print("📥 Chargement du dataset HC3 (Français)...")
    dataset = load_dataset("Hello-SimpleAI/HC3", "french", split="train")
    
    # Transformation du format HC3 en classification binaire
    # HC3 a 'human_answers' et 'chatgpt_answers'
    # On va créer un dataset plat : [texte, label]
    texts = []
    labels = []
    for item in dataset:
        for ans in item['human_answers']:
            texts.append(ans)
            labels.append(0) # Humain
        for ans in item['chatgpt_answers']:
            texts.append(ans)
            labels.append(1) # AI
            
    df = pd.DataFrame({"text": texts, "label": labels})
    
    # Échantillonnage pour la démo Hackathon (rapide)
    MAX_SAMPLES = 5000
    if len(df) > MAX_SAMPLES:
        df = df.sample(MAX_SAMPLES, random_state=42).reset_index(drop=True)
    
    print(f"✅ Dataset prêt : {len(df)} exemples.")
    hf_dataset = Dataset.from_pandas(df)

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

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

    tokenized_dataset = hf_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 optimisés pour Colab (GPU)
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        learning_rate=2e-5,
        per_device_train_batch_size=8,
        num_train_epochs=1,
        weight_decay=0.01,
        save_strategy="epoch",
        fp16=torch.cuda.is_available(), # Active l'accélération si GPU présent
        push_to_hub=False,
    )

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

    # 7. Lancement
    print("🚀 Début de l'entraînement sur Colab...")
    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()