#!/usr/bin/env python3
"""Utilidades para cargar, convertir y dividir datasets.

Requisitos: uv pip install datasets pandas pyarrow
Uso:        python3 utilidades_datos.py      (descarga un dataset pequeño y hace una demostración)
"""

import os
from pathlib import Path

from datasets import load_dataset


def cargar(nombre, config=None, split="train", streaming=False):
    """Carga un dataset del Hugging Face Hub (se guarda en caché tras la primera descarga)."""
    return load_dataset(nombre, config, split=split, streaming=streaming)


def dividir(dataset, prop_val=0.1, prop_test=0.1, semilla=42):
    """Devuelve (train, val, test) con una división reproducible."""
    primera = dataset.train_test_split(test_size=prop_test, seed=semilla)
    resto = primera["train"]
    segunda = resto.train_test_split(test_size=prop_val / (1 - prop_test), seed=semilla)
    return segunda["train"], segunda["test"], primera["test"]


def guardar_en_formatos(dataset, carpeta="datos"):
    """Guarda el dataset en CSV, JSON y Parquet y devuelve el tamaño de cada archivo en MB."""
    carpeta = Path(carpeta)
    carpeta.mkdir(exist_ok=True)
    rutas = {
        "csv": carpeta / "datos.csv",
        "json": carpeta / "datos.json",
        "parquet": carpeta / "datos.parquet",
    }
    dataset.to_csv(str(rutas["csv"]))
    dataset.to_json(str(rutas["json"]))
    dataset.to_parquet(str(rutas["parquet"]))
    return {formato: os.path.getsize(ruta) / 1e6 for formato, ruta in rutas.items()}


def resumen(dataset, nombre="dataset"):
    print(f"{nombre}: {len(dataset)} ejemplos · columnas: {', '.join(dataset.column_names)}")


def main():
    print("Descargando un dataset pequeño de reseñas (rotten_tomatoes)…\n")
    ds = cargar("cornell-movie-review-data/rotten_tomatoes")
    resumen(ds, "Completo")
    print(f"Primer ejemplo: {ds[0]}\n")

    train, val, test = dividir(ds)
    resumen(train, "Entrenamiento")
    resumen(val, "Validación")
    resumen(test, "Test")

    print("\nTamaño en disco según el formato:")
    for formato, mb in guardar_en_formatos(ds).items():
        print(f"  {formato:<8} {mb:6.2f} MB")
    print("\nFíjate en lo poco que ocupa Parquet frente a CSV y JSON.")


if __name__ == "__main__":
    main()