"""
Lightning DataModule for Video Collision Anticipation
====================================================

DataModule implementation that wraps the VideoDataset class
and provides standardized data loading for PyTorch Lightning.

Uses shared preprocessing from core.preprocessing - no HuggingFace dependency.
"""

import os
import random
from typing import Optional

import albumentations as A
import numpy as np
import pytorch_lightning as pl
import torch
from torch.utils.data import DataLoader, WeightedRandomSampler

from badas.training.config import Config
from badas.training.dataset import VideoDataset


class _VideoAugmentation:
    """Two-stage augmentation for video clips.

    clip_transform: applied with images= so all frames share the same random
        parameters (consistent lighting, flip, crop across the clip).
    frame_transforms: list of (transform, probability) pairs. Each transform's
        on/off decision is made ONCE per clip, then applied independently to
        every frame so weather effects (rain, snow, shadow) move naturally
        without switching type mid-clip.
    """

    def __init__(self, clip_transform, frame_transforms=None):
        self.clip_transform = clip_transform
        self.frame_transforms = frame_transforms or []

    def __call__(self, images, **kwargs):
        images = self.clip_transform(images=images)["images"]
        for transforms_list, p in self.frame_transforms:
            if random.random() < p:
                # Pick ONCE per clip so all frames get the same effect type
                chosen = random.choice(transforms_list)
                images = [chosen(image=f)["image"] for f in images]
        return {"images": images}


class VideoDataModule(pl.LightningDataModule):
    """Lightning DataModule for video classification datasets."""

    def __init__(self, config: Config, return_path: bool = False):
        super().__init__()
        self.config = config
        self.return_path = return_path

        # Save hyperparameters for checkpointing
        from dataclasses import asdict
        self.save_hyperparameters(asdict(config))

        # Datasets
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None

        # Setup transforms
        self._setup_transforms()

    def _setup_transforms(self):
        """Setup data transforms based on configuration."""
        if not self.config.data.use_augmentations:
            self.train_transform = None
            self.val_transform = None
            return

        # Clip-level transforms: same random params applied to ALL frames.
        # Probabilities kept moderate so clips frequently pass through clean.
        clip_transform = A.Compose(
            [
                # Spatial
                A.HorizontalFlip(p=0.5),
                A.Affine(scale=(0.9, 1.1), translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)}, rotate=(-5, 5), p=0.3),
                A.Perspective(scale=(0.02, 0.05), p=0.2),
                # Photometric
                A.RandomBrightnessContrast(p=0.4),
                A.HueSaturationValue(p=0.4),
                A.GaussianBlur(blur_limit=(3, 5), p=0.2),
                A.CLAHE(p=0.2),
                A.RandomGamma(p=0.2),
                A.RGBShift(p=0.2),
                A.MotionBlur(blur_limit=5, p=0.2),
                # Camera / lens artifacts
                A.Spatter(p=0.1),
                # Visibility — at most one dramatic effect per clip
                A.OneOf([
                    A.RandomFog(fog_coef_range=(0.1, 0.3), p=1.0),
                    A.RandomSunFlare(flare_roi=(0, 0, 1, 0.5), num_flare_circles_range=(3, 6),
                                     src_radius=100, p=1.0),
                ], p=0.15),
            ]
        )

        # Frame-level transforms: decided ONCE per clip, applied per-frame so
        # the effect moves naturally. Probabilities kept low to avoid overloading.
        # Each entry is ([option1, option2, ...], probability). One option is chosen
        # per clip so all frames get the same effect type (e.g. all rain or all snow).
        frame_transforms = [
            ([A.RandomRain(p=1.0), A.RandomSnow(p=1.0)], 0.15),
            ([A.RandomShadow(p=1.0)], 0.15),
            ([A.ISONoise(color_shift=(0.01, 0.03), intensity=(0.1, 0.3), p=1.0)], 0.15),
            ([A.ImageCompression(quality_range=(75, 95), p=1.0)], 0.2),
        ]

        self.train_transform = _VideoAugmentation(clip_transform, frame_transforms)

        # Validation/test transforms without augmentation
        self.val_transform = None

    def setup(self, stage: Optional[str] = None):
        """Setup datasets for different stages.

        Uses CSV-based loading with flat directory structure.
        All clips are in a 'clips' folder with labels managed via CSV.
        """
        if stage == "fit" or stage is None:
            csv_path = self.config.data.csv_path
            data_root = self.config.data.data_root
            class_names = getattr(self.config.data, 'class_names', None)

            if csv_path is None:
                raise ValueError(
                    "csv_path is required in config. "
                    "Set data.csv_path to your clips_metadata.csv file."
                )

            # Training dataset from CSV
            self.train_dataset = VideoDataset(
                config=self.config,
                csv_path=csv_path,
                data_root=data_root,
                split_filter="train",
                resample=self.config.data.oversample,
                resample_ratio=self.config.data.oversample_ratio,
                transform=self.train_transform,
                class_names=class_names,
                group_boost_column=self.config.data.group_boost_column,
                return_path=self.return_path,
            )
            print(f"Training from CSV with {len(self.train_dataset)} samples.")

            # Validation dataset - use same class mapping as training for consistency
            self.val_dataset = VideoDataset(
                config=self.config,
                csv_path=csv_path,
                data_root=data_root,
                split_filter="val",
                resample=False,
                transform=self.val_transform,
                class_names=self.train_dataset.classes,  # Use train's class ordering
                return_path=self.return_path,
            )
            print(f"Validation from CSV with {len(self.val_dataset)} samples.")

    def _worker_init_fn(self, worker_id):
        """Initialize worker with reproducible random seed."""
        worker_seed = self.config.system.seed + worker_id
        random.seed(worker_seed)
        np.random.seed(worker_seed)
        torch.manual_seed(worker_seed)

    def train_dataloader(self) -> DataLoader:
        """Return training dataloader."""
        boost_col = self.config.data.group_boost_column
        boost_factor = self.config.data.group_boost_factor

        if boost_col is not None and boost_factor != 1.0:
            # Group boosting: use WeightedRandomSampler so the boosted group
            # (both positives and negatives) is drawn more frequently.
            # This replaces shuffle=True — sampling is still random, just non-uniform.
            weights = self.train_dataset.get_sample_weights(boost_factor)
            sampler = WeightedRandomSampler(
                weights=weights,
                num_samples=len(weights),
                replacement=True,
            )
            print(
                f"Group boost active: column='{boost_col}', factor={boost_factor}. "
                f"Using WeightedRandomSampler."
            )
            return DataLoader(
                self.train_dataset,
                batch_size=self.config.training.batch_size,
                sampler=sampler,
                num_workers=self.config.data.num_workers,
                pin_memory=False,
                persistent_workers=self.config.data.num_workers > 0,
                worker_init_fn=self._worker_init_fn,
            )

        # Default: uniform shuffle — existing behavior unchanged
        return DataLoader(
            self.train_dataset,
            batch_size=self.config.training.batch_size,
            shuffle=True,
            num_workers=self.config.data.num_workers,
            pin_memory=False,
            persistent_workers=self.config.data.num_workers > 0,
            worker_init_fn=self._worker_init_fn,
        )

    def val_dataloader(self) -> DataLoader:
        """Return validation dataloader."""
        if self.config.training.val_batch_size is not None:
            batch_size = self.config.training.val_batch_size
        else:
            batch_size = self.config.training.batch_size
        return DataLoader(
            self.val_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=self.config.data.num_workers,
            pin_memory=False,
            persistent_workers=self.config.data.num_workers > 0,
            worker_init_fn=self._worker_init_fn,
        )
