from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import numpy as np
from sklearn.model_selection import GroupShuffleSplit


@dataclass(frozen=True)
class DatasetSplits:
    train: np.ndarray
    val: np.ndarray
    test: np.ndarray


def load_pose_dataset(path: str | Path) -> dict[str, np.ndarray]:
    """Load the compressed dataset produced by prepare_dataset.py."""
    with np.load(path, allow_pickle=False) as data:
        required = {"poses", "labels", "groups", "sources"}
        missing = required.difference(data.files)
        if missing:
            raise ValueError(f"Dataset is missing arrays: {sorted(missing)}")
        return {key: data[key] for key in required}


def group_train_val_test_split(
    labels: np.ndarray,
    groups: np.ndarray,
    test_size: float = 0.20,
    val_size: float = 0.20,
    seed: int = 42,
) -> DatasetSplits:
    """Split without putting windows from the same source video in two sets."""
    indices = np.arange(len(labels))
    relative_val_size = val_size / (1.0 - test_size)
    selected = None
    for attempt in range(100):
        outer = GroupShuffleSplit(
            n_splits=1, test_size=test_size, random_state=seed + attempt * 2
        )
        train_val_idx, test_idx = next(outer.split(indices, labels, groups))
        inner = GroupShuffleSplit(
            n_splits=1,
            test_size=relative_val_size,
            random_state=seed + attempt * 2 + 1,
        )
        inner_train, inner_val = next(
            inner.split(train_val_idx, labels[train_val_idx], groups[train_val_idx])
        )
        train_idx = train_val_idx[inner_train]
        val_idx = train_val_idx[inner_val]
        if all(len(np.unique(labels[idx])) == 2 for idx in (train_idx, val_idx, test_idx)):
            selected = (train_idx, val_idx, test_idx)
            break
    if selected is None:
        raise ValueError("Could not create group splits containing both classes")
    train_idx, val_idx, test_idx = selected

    split_groups = [set(groups[idx].tolist()) for idx in (train_idx, val_idx, test_idx)]
    if any(split_groups[i] & split_groups[j] for i in range(3) for j in range(i + 1, 3)):
        raise RuntimeError("Group leakage detected")
    return DatasetSplits(train=train_idx, val=val_idx, test=test_idx)


def save_splits(path: str | Path, splits: DatasetSplits) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, train=splits.train, val=splits.val, test=splits.test)


def load_splits(path: str | Path) -> DatasetSplits:
    with np.load(path) as data:
        return DatasetSplits(train=data["train"], val=data["val"], test=data["test"])