"""Phase 5 — paires image-légende pour l'alignement image-texte.

Par défaut, on utilise le **générateur jouet maison** (`toy_images`) : 9 classes
forme × couleur, chacune décrite en français (« un cercle rouge »). Cela suffit à
valider un CLIP minimal sans télécharger COCO.

Option réelle : `--source hf --preset coco_fr` télécharge un sous-ensemble de
paires image-légende depuis le Hub (si disponible dans l'environnement).

Usage :

    python -m src.data.prepare_image_text --out-dir data/processed/toy_clip \\
        --n-train 1800 --n-val 180
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from src.data.toy_images import generate_dataset, to_arrays


def make_toy_pairs(n: int, grid: int = 32, seed: int = 0):
    """Retourne (images NCHW float32, captions list[str], labels list[str])."""
    items = generate_dataset(n, grid=grid, seed=seed)
    images, labels = to_arrays(items)
    captions = [it.caption for it in items]
    names = [it.label for it in items]
    return images, captions, names


def prepare_toy(out_dir: Path, n_train: int = 1800, n_val: int = 180, grid: int = 32,
                seed: int = 0) -> dict:
    """Écrit images/captions train+val en `.npy`/`.json`."""
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    tr_imgs, tr_caps, tr_names = make_toy_pairs(n_train, grid=grid, seed=seed)
    va_imgs, va_caps, va_names = make_toy_pairs(n_val, grid=grid, seed=seed + 999)

    np.save(out_dir / "train_images.npy", tr_imgs)
    np.save(out_dir / "val_images.npy", va_imgs)
    (out_dir / "train_captions.json").write_text(json.dumps(tr_caps, ensure_ascii=False), encoding="utf-8")
    (out_dir / "val_captions.json").write_text(json.dumps(va_caps, ensure_ascii=False), encoding="utf-8")
    (out_dir / "train_labels.json").write_text(json.dumps(tr_names, ensure_ascii=False), encoding="utf-8")
    (out_dir / "val_labels.json").write_text(json.dumps(va_names, ensure_ascii=False), encoding="utf-8")

    meta = {"source": "toy", "grid": grid, "n_train": n_train, "n_val": n_val,
            "n_classes": len(set(tr_names)), "seed": seed}
    (out_dir / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    return meta


def load_pairs(out_dir: Path, split: str = "train"):
    """Charge (images, captions) d'un split."""
    out_dir = Path(out_dir)
    images = np.load(out_dir / f"{split}_images.npy")
    captions = json.loads((out_dir / f"{split}_captions.json").read_text(encoding="utf-8"))
    return images, captions


def main() -> None:
    parser = argparse.ArgumentParser(description="Préparer des paires image-texte (Phase 5).")
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--source", choices=["toy"], default="toy")
    parser.add_argument("--n-train", type=int, default=1800)
    parser.add_argument("--n-val", type=int, default=180)
    parser.add_argument("--grid", type=int, default=32)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()
    meta = prepare_toy(args.out_dir, args.n_train, args.n_val, args.grid, args.seed)
    print(json.dumps(meta, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()