import random
from typing import Callable, List, Tuple

from torch.utils.data import Dataset, Subset

from src.data.dataset import get_source_name
from src.preprocessing import get_transforms
from src.utils import create_group_kfold_splits, get_basename


# Defined at module level so DataLoader workers (num_workers > 0) can serialize it safely
class TransformSubset(Dataset):
    def __init__(self, subset, transform):
        self.subset, self.transform = subset, transform

    def __len__(self):
        return len(self.subset)

    def __getitem__(self, idx):
        img, label = self.subset[idx]
        return self.transform(img), label


# ── Splitting ──────────────────────────────────────────────────────────────

# Builds grouped CV fold indices from config; this is the only split strategy used.
def get_splits(
    raw_dataset,
    cfg,
) -> List[Tuple[List[int], List[int], List[int]]]:
    splits = create_group_kfold_splits(
        raw_dataset.samples,
        n_splits=cfg.get("cv_folds", 5),
        seed=cfg.get("seed", 42),
    )

    # Optional source holdout: train/val from train_sources, test from eval_sources.
    train_sources = cfg.get("train_sources")
    eval_sources = cfg.get("eval_sources")
    if train_sources or eval_sources:
        all_sources = {get_source_name(path) for path, _ in raw_dataset.samples}
        ts = set(train_sources or all_sources)
        es = set(eval_sources or all_sources)
        unknown = (ts | es) - all_sources
        if unknown:
            raise ValueError(f"Unknown sources requested: {sorted(unknown)}")
        splits = [
            (
                [i for i in tr if get_source_name(raw_dataset.samples[i][0]) in ts],
                [i for i in val if get_source_name(raw_dataset.samples[i][0]) in ts],
                [i for i in te if get_source_name(raw_dataset.samples[i][0]) in es],
            )
            for tr, val, te in splits
        ]
    return splits


# Deterministic subsampling shared by training and reevaluation.
def apply_subsample(raw_dataset, cfg) -> tuple[int, int] | None:
    subsample = cfg.get("subsample", 1.0)
    if subsample >= 1.0:
        return None

    total = len(raw_dataset.samples)
    if total == 0:
        return 0, 0

    # Subsample at basename-group level to preserve identity grouping guarantees.
    group_to_indices = {}
    for idx, (path, _) in enumerate(raw_dataset.samples):
        group = get_basename(str(path))
        group_to_indices.setdefault(group, []).append(idx)

    groups = list(group_to_indices.keys())
    n_groups = len(groups)
    target_groups = max(1, int(n_groups * subsample))
    rng = random.Random(cfg.get("seed", 42))
    rng.shuffle(groups)
    keep_groups = set(groups[:target_groups])

    keep_indices = [
        idx
        for group in keep_groups
        for idx in group_to_indices[group]
    ]
    keep_indices.sort()
    raw_dataset.samples = [raw_dataset.samples[i] for i in keep_indices]
    return len(raw_dataset.samples), total


# Controls stochastic augmentations (flip, jitter, etc.)
# augment=False -> NO_AUGMENT preset (square-crop + resize + normalize still run)
# augment=None, pipeline defaults,  augment=dict -> override specific params
# Face cropping is handled upstream via data_dir swap, not here
def build_transforms(raw_dataset, cfg, augment=None) -> Callable:
    image_size = cfg["image_size"]

    if augment is False:
        from src.preprocessing.pipeline import DFFImagePipeline

        augment = DFFImagePipeline.NO_AUGMENT

    def transform_builder(indices, train=True, normalize_mean=None, normalize_std=None):
        subset = Subset(raw_dataset, indices)
        return TransformSubset(
            subset,
            get_transforms(
                train=train,
                image_size=image_size,
                augment=augment,
                normalize_mean=normalize_mean,
                normalize_std=normalize_std,
            ),
        )
    return transform_builder

