# filename: src/data/dataset.py
# purpose:  PyTorch Dataset for SmartVision classification (NUM_CLASSES from config)
# version:  2.0

import logging
import platform
from pathlib import Path
from typing import Optional

from PIL import Image
from torch.utils.data import DataLoader, Dataset

from config import (
    CLASS_TO_IDX,
    CLASSES,
    DATA_PROCESSED_DIR,
    MODEL_CONFIGS,
    RANDOM_STATE,
    TRAIN_SPLIT,
    VAL_SPLIT,
)

# Windows requires num_workers=0 (no fork); Linux/Colab can use 2 workers
_NUM_WORKERS = 0 if platform.system() == "Windows" else 2

logger = logging.getLogger(__name__)

# Session-level cache: guarantees identical splits across multiple get_dataloaders()
# calls. Key = (str(data_root), random_state).
_split_cache: dict[tuple, dict] = {}


def create_stratified_split(
    data_root: Path,
    random_state: int = RANDOM_STATE,
) -> dict[str, list[tuple[Path, int]]]:
    """
    Build a stratified random split from the existing on-disk split directories.

    Reads ALL images from train/ + val/ + test/ subdirs, shuffles per class
    with a fixed seed, then re-splits 70/15/15. Uses round() not int() to avoid
    silently dropping images for class counts not divisible by 20.

    Integrity checks (in order):
      1. Completeness: len(assigned) == len(all_samples)  — catches round() edge cases
      2. Uniqueness:   no path appears twice              — structural guarantee
      3. Cross-split assertions (redundant by construction; kept as documentation)

    Caches by (data_root, random_state): identical splits across get_dataloaders() calls.

    Returns:
        {'train': [(Path, int), ...], 'val': [...], 'test': [...]}
    """
    cache_key = (str(data_root), random_state)
    if cache_key in _split_cache:
        logger.debug("create_stratified_split: returning cached split")
        return _split_cache[cache_key]

    import random as _random

    all_samples: list[tuple[Path, int]] = []
    for split_name in ("train", "val", "test"):
        for class_name in sorted(CLASS_TO_IDX.keys()):
            class_dir = data_root / split_name / class_name
            if not class_dir.exists():
                continue
            class_idx = CLASS_TO_IDX[class_name]
            for img_path in sorted(class_dir.glob("*.jpg")):
                all_samples.append((img_path, class_idx))

    rng = _random.Random(random_state)
    splits: dict[str, list] = {"train": [], "val": [], "test": []}
    per_class: dict[int, list] = {}
    for item in all_samples:
        per_class.setdefault(item[1], []).append(item)

    for class_idx, items in per_class.items():
        rng.shuffle(items)
        n = len(items)
        train_end = round(n * TRAIN_SPLIT)
        val_end = round(n * (TRAIN_SPLIT + VAL_SPLIT))
        splits["train"].extend(items[:train_end])
        splits["val"].extend(items[train_end:val_end])
        splits["test"].extend(items[val_end:])

    all_assigned = splits["train"] + splits["val"] + splits["test"]
    assigned_paths = [str(p) for p, _ in all_assigned]
    assert len(assigned_paths) == len(all_samples), (
        f"Split dropped {len(all_samples) - len(assigned_paths)} images. "
        f"total={len(all_samples)} assigned={len(assigned_paths)}"
    )
    assert len(set(assigned_paths)) == len(assigned_paths), (
        f"Duplicate paths: {len(assigned_paths) - len(set(assigned_paths))} duplicates"
    )
    train_set = {str(p) for p, _ in splits["train"]}
    val_set = {str(p) for p, _ in splits["val"]}
    test_set = {str(p) for p, _ in splits["test"]}
    assert not (train_set & val_set), f"Train/val overlap: {len(train_set & val_set)}"
    assert not (train_set & test_set), (
        f"Train/test overlap: {len(train_set & test_set)}"
    )
    assert not (val_set & test_set), f"Val/test overlap: {len(val_set & test_set)}"

    logger.info(
        f"create_stratified_split: train={len(splits['train'])} "
        f"val={len(splits['val'])} test={len(splits['test'])}"
    )
    _split_cache[cache_key] = splits
    return splits


class SmartVisionDataset(Dataset):
    """
    Classification dataset: one cropped 224×224 image per sample, one class label.
    Folder structure: {root}/{split}/{class_name}/*.jpg
    """

    def __init__(
        self,
        split: str,
        transform=None,
        root: Optional[Path] = None,
        samples: Optional[list] = None,
    ) -> None:
        if split not in ("train", "val", "test"):
            raise ValueError(f"split must be 'train', 'val', or 'test'; got {split!r}")
        self.split = split
        self.transform = transform
        self.root = Path(root) if root else DATA_PROCESSED_DIR / "classification"

        if samples is not None:
            self.samples = list(
                samples
            )  # shallow copy — (Path, int) tuples are immutable
            logger.info(
                f"SmartVisionDataset [{self.split}]: {len(self.samples)} samples (pre-split)"
            )
        else:
            self.samples: list[tuple[Path, int]] = []
            self._load_samples()

    def _load_samples(self) -> None:
        split_dir = self.root / self.split
        if not split_dir.exists():
            logger.warning(f"Split directory not found: {split_dir}")
            return

        for class_name in sorted(CLASS_TO_IDX.keys()):
            class_dir = split_dir / class_name
            if not class_dir.exists():
                continue
            class_idx = CLASS_TO_IDX[class_name]
            for img_path in sorted(class_dir.glob("*.jpg")):
                self.samples.append((img_path, class_idx))

        logger.info(f"SmartVisionDataset [{self.split}]: {len(self.samples)} samples")

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int):
        img_path, class_idx = self.samples[idx]
        img = Image.open(img_path).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, class_idx

    def class_counts(self) -> dict[str, int]:
        counts: dict[str, int] = {cls: 0 for cls in CLASSES}
        for _, class_idx in self.samples:
            counts[CLASSES[class_idx]] += 1
        return counts


def get_dataloaders(
    model_name: str,
    train_transform,
    eval_transform,
    fast_mode: bool = True,
    use_random_split: bool = True,
) -> tuple[DataLoader, DataLoader, DataLoader]:
    """Returns (train_loader, val_loader, test_loader) for the given model."""
    batch_size = MODEL_CONFIGS[model_name]["batch"]

    if use_random_split:
        root = DATA_PROCESSED_DIR / "classification"
        all_splits = create_stratified_split(root)
        train_ds = SmartVisionDataset(
            "train", transform=train_transform, samples=all_splits["train"]
        )
        val_ds = SmartVisionDataset(
            "val", transform=eval_transform, samples=all_splits["val"]
        )
        test_ds = SmartVisionDataset(
            "test", transform=eval_transform, samples=all_splits["test"]
        )
    else:
        train_ds = SmartVisionDataset("train", transform=train_transform)
        val_ds = SmartVisionDataset("val", transform=eval_transform)
        test_ds = SmartVisionDataset("test", transform=eval_transform)

    if fast_mode and len(train_ds) < 50:
        logger.warning(
            f"Only {len(train_ds)} training samples found. "
            "Run 01_data_acquisition.py with FAST_MODE=True first."
        )

    train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=_NUM_WORKERS,
        pin_memory=True,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=_NUM_WORKERS,
        pin_memory=True,
    )
    test_loader = DataLoader(
        test_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=_NUM_WORKERS,
        pin_memory=True,
    )

    logger.info(
        f"DataLoaders [{model_name}]: "
        f"train={len(train_ds)} val={len(val_ds)} test={len(test_ds)} batch={batch_size}"
    )
    return train_loader, val_loader, test_loader