"""Create a stratified train / val / test split manifest for the Augmented Dataset.

Outputs a single JSON manifest so the same split is reused by every model run
(training, k-fold CV, and the final independent-test evaluation).

The independent test set is held out FIRST and is never used during k-fold CV.
The k-fold CV runs on the remaining train+val pool (the script also stores
five stratified train/val folds so they can be reproduced exactly).
"""

from __future__ import annotations

import argparse
import json
import random
from collections import Counter, defaultdict
from pathlib import Path

from sklearn.model_selection import StratifiedKFold, train_test_split


def collect_samples(data_dir: Path) -> tuple[list[tuple[str, str]], list[str]]:
    classes = sorted([p.name for p in data_dir.iterdir() if p.is_dir()])
    samples: list[tuple[str, str]] = []
    valid_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
    for cls in classes:
        for image_path in sorted((data_dir / cls).iterdir()):
            if image_path.suffix.lower() in valid_exts:
                samples.append((str(image_path.relative_to(data_dir)), cls))
    return samples, classes


def stratified_split(samples, test_size, val_size, seed):
    paths = [s[0] for s in samples]
    labels = [s[1] for s in samples]
    paths_pool, paths_test, labels_pool, labels_test = train_test_split(
        paths, labels, test_size=test_size, stratify=labels, random_state=seed
    )
    relative_val = val_size / (1.0 - test_size)
    paths_train, paths_val, labels_train, labels_val = train_test_split(
        paths_pool, labels_pool, test_size=relative_val, stratify=labels_pool, random_state=seed
    )
    return (paths_train, labels_train), (paths_val, labels_val), (paths_test, labels_test)


def kfold_indices(paths, labels, folds, seed):
    skf = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed)
    out = []
    for k, (train_idx, val_idx) in enumerate(skf.split(paths, labels), start=1):
        out.append({
            "fold": k,
            "train": [int(i) for i in train_idx],
            "val": [int(i) for i in val_idx],
        })
    return out


def class_distribution(labels):
    return dict(Counter(labels))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="Database/Augmented_Dataset")
    parser.add_argument("--output", default="holdout_split.json")
    parser.add_argument("--test-size", type=float, default=0.15)
    parser.add_argument("--val-size", type=float, default=0.15)
    parser.add_argument("--folds", type=int, default=5)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    data_dir = Path(args.data_dir).resolve()
    samples, classes = collect_samples(data_dir)
    print(f"Dataset: {data_dir}")
    print(f"  total images: {len(samples)}")
    print(f"  classes: {classes}")

    (train, val, test) = stratified_split(samples, args.test_size, args.val_size, args.seed)
    folds = kfold_indices(
        train[0] + val[0],
        train[1] + val[1],
        args.folds,
        args.seed,
    )

    manifest = {
        "data_dir": str(data_dir),
        "classes": classes,
        "seed": args.seed,
        "test_size": args.test_size,
        "val_size": args.val_size,
        "splits": {
            "train": list(zip(train[0], train[1])),
            "val": list(zip(val[0], val[1])),
            "test": list(zip(test[0], test[1])),
        },
        "kfold": {
            "folds": args.folds,
            "pool_paths": train[0] + val[0],
            "pool_labels": train[1] + val[1],
            "indices": folds,
        },
        "class_distribution": {
            "train": class_distribution(train[1]),
            "val": class_distribution(val[1]),
            "test": class_distribution(test[1]),
        },
    }

    out_path = Path(args.output).resolve()
    out_path.write_text(json.dumps(manifest, indent=2))
    print(f"Manifest written: {out_path}")
    for split_name in ("train", "val", "test"):
        print(f"  {split_name}: {len(manifest['splits'][split_name])} images")


if __name__ == "__main__":
    main()