"""
Download FairFace dataset from Hugging Face
"""

from datasets import load_dataset
import pandas as pd
from pathlib import Path
from tqdm import tqdm

# Mapping constants
RACE_LABELS = ['East Asian', 'Indian', 'Black', 'White', 'Middle Eastern', 
               'Latino_Hispanic', 'Southeast Asian']
GENDER_LABELS = {0: 'Male', 1: 'Female'}


def process_split(dataset_split, output_dir, split_name):
    """Process and save dataset split"""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    data = []
    for idx, sample in enumerate(tqdm(dataset_split, desc=f"Processing {split_name}")):
        # Save image
        img_filename = f"{split_name}_{idx}.jpg"
        sample['image'].save(output_dir / img_filename)
        
        # Prepare CSV row
        data.append({
            'file': img_filename,
            'age': sample['age'],
            'gender': GENDER_LABELS[sample['gender']],
            'race': RACE_LABELS[sample['race']],
            'service_test': sample.get('service_test', False)
        })
    
    return pd.DataFrame(data)


def main():
    print("FairFace Dataset Downloader")
    print("-" * 50)
    
    # Setup directories
    data_dir = Path("data/raw/fairface")
    data_dir.mkdir(parents=True, exist_ok=True)
    
    # Load dataset
    print("\nLoading dataset from Hugging Face...")
    dataset = load_dataset("HuggingFaceM4/FairFace", "0.25")
    
    print(f"Training samples: {len(dataset['train'])}")
    print(f"Validation samples: {len(dataset['validation'])}")
    
    # Process validation set
    print("\nDownloading validation set...")
    val_df = process_split(dataset['validation'], data_dir / 'val', 'val')
    val_df.to_csv(data_dir / 'fairface_label_val.csv', index=False)
    print(f"Validation complete: {len(val_df)} images")
    
    # Process training set
    print("\nDownloading training set...")
    train_df = process_split(dataset['train'], data_dir / 'train', 'train')
    train_df.to_csv(data_dir / 'fairface_label_train.csv', index=False)
    print(f"Training complete: {len(train_df)} images")
    
    # Summary
    print("\n" + "-" * 50)
    print("Download complete")
    print(f"Location: {data_dir}")


if __name__ == "__main__":
    main()