from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from datasets import Dataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"


def split_indices(labels: np.ndarray, seed: int = 2026) -> dict[str, np.ndarray]:
    indices = np.arange(len(labels))
    train, remainder = train_test_split(
        indices,
        test_size=0.30,
        stratify=labels,
        random_state=seed,
    )
    validation, test = train_test_split(
        remainder,
        test_size=0.50,
        stratify=labels[remainder],
        random_state=seed,
    )
    return {"train": train, "validation": validation, "test": test}


def main() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    digits = load_digits()
    images = digits.images.astype(np.float32)
    labels = digits.target.astype(np.int64)
    manifest = {}
    for name, indices in split_indices(labels).items():
        dataset = Dataset.from_dict(
            {
                "image": [image.reshape(-1).tolist() for image in images[indices]],
                "label": labels[indices].tolist(),
            }
        )
        path = DATA_DIR / f"{name}.parquet"
        dataset.to_parquet(path)
        counts = np.bincount(labels[indices], minlength=10)
        manifest[name] = {
            "rows": len(indices),
            "class_counts": counts.tolist(),
            "path": path.name,
        }
    (DATA_DIR / "manifest.json").write_text(
        json.dumps(manifest, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()