import csv
from pathlib import Path

from datasets import load_dataset
from tqdm import tqdm

# Конфигурация корневой директории локального датасета
DATASET_DIR = Path("../dataset/EmoSet-118K")

def process_and_save_split(dataset_split, split_name: str, output_dir: Path):
    # Подготовка структуры директорий для текущей выборки
    split_dir = output_dir / split_name
    img_dir = split_dir / "images"
    img_dir.mkdir(parents=True, exist_ok=True)

    labels_path = split_dir / "labels.csv"

    print(f"Обработка выборки: {split_name}...")

    # Открытие файла разметки перед циклом для минимизации I/O операций диска
    with open(labels_path, mode="w", newline="", encoding="utf-8") as csv_file:
        writer = csv.writer(csv_file)
        writer.writerow(["filename", "label"])

        for example in tqdm(dataset_split, desc=split_name):
            img = example["image"]
            emotion_label = example["emotion"]
            img_id = example["image_id"]

            file_name = f"{img_id}.jpg"
            
            # Принудительная конвертация в RGB для безопасного сохранения в JPEG-формате
            if img.mode != "RGB":
                img = img.convert("RGB")
                
            img.save(img_dir / file_name, format="JPEG")
            
            writer.writerow([file_name, emotion_label])

if __name__ == "__main__":
    DATASET_DIR.mkdir(exist_ok=True, parents=True)

    # Инициализация подключения к Hugging Face Hub
    print("Загрузка метаданных EmoSet-118K...")
    raw_dataset = load_dataset("Woleek/EmoSet-118K")

    # Итеративная выгрузка размеченных данных
    for split_key in ["train", "val", "test"]:
        if split_key in raw_dataset:
            process_and_save_split(
                dataset_split=raw_dataset[split_key],
                split_name=split_key,
                output_dir=DATASET_DIR
            )

    print("Экспорт датасета завершен.")