import os
import pandas as pd
from datasets import load_dataset

def main():
    print("Chargement du dataset 'aadityaubhat/GPT-wiki-intro'...")
    # On charge les 2000 premières lignes
    dataset = load_dataset('aadityaubhat/GPT-wiki-intro', split='train[:2000]')
    
    # On convertit en dataframe pandas
    df = dataset.to_pandas()
    
    print("Préparation des données...")
    # Isoler les textes humains et attribuer le label 0
    df_human = pd.DataFrame({
        'text': df['wiki_intro'],
        'label': 0
    })
    
    # Isoler les textes générés par IA et attribuer le label 1
    df_ai = pd.DataFrame({
        'text': df['generated_intro'],
        'label': 1
    })
    
    # Concaténer
    df_final = pd.concat([df_human, df_ai], ignore_index=True)
    
    # Mélanger aléatoirement
    df_final = df_final.sample(frac=1, random_state=42).reset_index(drop=True)
    
    # S"assurer que le dossier data/ existe
    os.makedirs('data', exist_ok=True)
    
    output_path = 'data/dataset_clean.csv'
    print(f"Sauvegarde du dataset dans {output_path}...")
    df_final.to_csv(output_path, index=False, encoding='utf-8')
    
    print("Terminé avec succès !")
    print(f"Nombre total de lignes : {len(df_final)}")
    print(df_final.head())

if __name__ == '__main__':
    main()