from __future__ import annotations

from pathlib import Path

import numpy as np

from fall_detection.config import load_config
from fall_detection.data import (
    DatasetSplits,
    group_train_val_test_split,
    load_pose_dataset,
    load_splits,
    save_splits,
)
from fall_detection.features import featurize_dataset


def prepare_experiment_data(
    dataset_path: str | Path,
    config_path: str | Path,
    output_root: str | Path,
    seed: int | None = None,
) -> tuple[dict, dict[str, np.ndarray], np.ndarray, DatasetSplits]:
    config = load_config(config_path)
    if seed is not None:
        config["seed"] = int(seed)
    dataset = load_pose_dataset(dataset_path)
    features = featurize_dataset(
        dataset["poses"], config["data"]["visibility_threshold"]
    )
    split_path = Path(output_root) / "splits.npz"
    if split_path.exists():
        splits = load_splits(split_path)
        all_indices = np.concatenate([splits.train, splits.val, splits.test])
        if len(all_indices) != len(dataset["labels"]) or all_indices.max() >= len(dataset["labels"]):
            raise ValueError(
                f"Existing {split_path} does not match this dataset; remove it or use another output directory"
            )
    else:
        splits = group_train_val_test_split(
            dataset["labels"],
            dataset["groups"],
            test_size=config["data"]["test_size"],
            val_size=config["data"]["val_size"],
            seed=config["seed"],
        )
        save_splits(split_path, splits)
    return config, dataset, features, splits


def split_summary(
    labels: np.ndarray, groups: np.ndarray, splits: DatasetSplits
) -> dict[str, dict[str, int]]:
    result: dict[str, dict[str, int]] = {}
    for name, indices in (
        ("train", splits.train),
        ("validation", splits.val),
        ("test", splits.test),
    ):
        result[name] = {
            "samples": int(len(indices)),
            "groups": int(len(np.unique(groups[indices]))),
            "normal": int(np.sum(labels[indices] == 0)),
            "fall": int(np.sum(labels[indices] == 1)),
        }
    return result