"""
Video Dataset for Training.

Uses shared preprocessing from core.preprocessing for consistency with inference.
Supports both binary classification (numeric labels 0/1) and multi-class
classification (string labels like 'collision_front', 'normal_driving', etc.).
"""

import os
import random
from typing import List, Optional

import cv2
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

try:
    import decord
    decord.bridge.set_bridge("native")
    _DECORD_AVAILABLE = True
except ImportError:
    _DECORD_AVAILABLE = False

from badas.training.config import Config
from badas.training.utils import get_frames_ffmpeg
from badas.core.preprocessing import (
    calculate_frame_indices,
    sample_frames_uniform,
    preprocess_frames_batch,
)


class VideoDataset(Dataset):
    """Video dataset with shared preprocessing pipeline.

    Uses CSV metadata for loading clips from flat directory structure.
    All clips are stored in a single 'clips' folder with labels in CSV.

    Supports both binary (numeric 0/1) and multi-class (string) labels.
    """

    def __init__(
        self,
        config: Config,
        csv_path: str,
        data_root: str,
        split_filter: Optional[str] = None,
        resample: bool = False,
        resample_ratio: float = 1.0,
        transform=None,
        class_names: Optional[List[str]] = None,
        group_boost_column: Optional[str] = None,
        return_path: bool = False,
    ):
        """Initialize VideoDataset.

        Args:
            config: Training configuration
            csv_path: Path to CSV file with columns: path, label, split
            data_root: Root directory for relative paths in CSV
            split_filter: Filter CSV by split column (e.g., 'train', 'val', 'test')
            resample: Whether to resample for class balance
            resample_ratio: Target ratio for resampling (binary: neg:pos ratio,
                           multi-class: oversample all to match largest class)
            transform: Augmentation transforms
            class_names: Optional explicit list of class names for consistent ordering.
                        If None, classes are inferred from CSV labels.
            group_boost_column: Optional boolean CSV column name. Paths where this
                        column is truthy are tracked so the dataloader can boost
                        their sampling probability. None = no group boosting.
        """
        self.config = config
        self.transform = transform
        self.frame_count = config.data.frame_count
        self.img_size = config.data.img_size

        # Class mapping will be built from CSV or provided class_names
        self.classes: List[str] = []
        self.class_to_idx: dict = {}
        self.samples: List[tuple] = []

        # Paths belonging to the boosted group (populated in _load_from_csv)
        self._boost_paths: set = set()

        self.return_path = return_path
        self._load_from_csv(csv_path, data_root, split_filter, class_names, group_boost_column)

        if resample:
            self.resample(ratio=resample_ratio)

    def _load_from_csv(
        self,
        csv_path: str,
        data_root: str,
        split_filter: Optional[str] = None,
        class_names: Optional[List[str]] = None,
        group_boost_column: Optional[str] = None,
    ):
        """Load samples from CSV file.

        CSV must have columns:
            - path: relative path to clip file (e.g., 'clips/video_win001_l1.mp4')
            - label: numeric (0/1) for binary, or string for multi-class
            - split: 'train', 'val', or 'test'

        Args:
            csv_path: Path to CSV metadata file
            data_root: Root directory for relative paths
            split_filter: Filter by split column value
            class_names: Optional explicit class ordering
            group_boost_column: Optional boolean column name for group boosting
        """
        df = pd.read_csv(csv_path)

        # Validate required columns
        required_cols = ['path', 'label']
        missing = [c for c in required_cols if c not in df.columns]
        if missing:
            raise ValueError(f"CSV missing required columns: {missing}")

        # Filter by split if specified
        if split_filter is not None:
            if 'split' not in df.columns:
                raise ValueError(f"CSV has no 'split' column but split_filter='{split_filter}' was requested")
            df = df[df['split'] == split_filter]
            print(f"Filtered to split '{split_filter}': {len(df)} samples")

        # Build class mapping from labels
        self._build_class_mapping(df['label'], class_names)

        # Warn if boost column is missing instead of silently ignoring
        if group_boost_column is not None and group_boost_column not in df.columns:
            raise ValueError(
                f"group_boost_column='{group_boost_column}' not found in CSV. "
                f"Available columns: {list(df.columns)}"
            )

        for _, row in df.iterrows():
            path = row['path']
            label = row['label']

            # Convert label to index
            label_idx = self._label_to_index(label)

            # Make path absolute using data_root
            if not os.path.isabs(path):
                path = os.path.join(data_root, path)

            self.samples.append((path, label_idx))

            # Track paths belonging to the boosted group
            if group_boost_column is not None and row[group_boost_column]:
                self._boost_paths.add(path)

        print(f"Loaded {len(self.samples)} samples from CSV ({len(self.classes)} classes)")
        if group_boost_column is not None:
            print(f"  Group boost column '{group_boost_column}': {len(self._boost_paths)} samples in boosted group")

    def _build_class_mapping(
        self,
        labels: pd.Series,
        class_names: Optional[List[str]] = None,
    ):
        """Build class name to index mapping.

        Args:
            labels: Series of label values from CSV
            class_names: Optional explicit class ordering
        """
        # Check if labels are numeric (binary classification)
        sample_label = labels.iloc[0]
        is_numeric = isinstance(sample_label, (int, np.integer, float, np.floating))

        if is_numeric:
            # Binary classification with numeric labels
            unique_labels = sorted(labels.unique())
            if set(unique_labels).issubset({0, 1, 0.0, 1.0}):
                self.classes = ["negative", "positive"]
                self.class_to_idx = {"negative": 0, "positive": 1}
            else:
                # Numeric but not 0/1 - treat as class indices
                self.classes = [str(int(l)) for l in unique_labels]
                self.class_to_idx = {name: idx for idx, name in enumerate(self.classes)}
        else:
            # String labels (multi-class classification)
            if class_names is not None:
                # Use provided class ordering
                self.classes = list(class_names)
            else:
                # Infer from data, sorted alphabetically for consistency
                self.classes = sorted(labels.unique().tolist())

            self.class_to_idx = {name: idx for idx, name in enumerate(self.classes)}

        print(f"Classes ({len(self.classes)}): {self.classes}")

    def _label_to_index(self, label) -> int:
        """Convert a label (string or numeric) to class index."""
        if isinstance(label, (int, np.integer, float, np.floating)):
            # Numeric label
            if len(self.classes) == 2 and self.classes == ["negative", "positive"]:
                # Binary: 0 -> 0, 1 -> 1
                return int(label)
            else:
                # Numeric class index
                return int(label)
        else:
            # String label
            if label not in self.class_to_idx:
                raise ValueError(f"Unknown label '{label}'. Known classes: {self.classes}")
            return self.class_to_idx[label]

    def resample(self, ratio: float = 1.0):
        """Resample the dataset by oversampling minority classes.

        For binary classification:
            ratio: Target ratio of negative:positive samples.
                   1.0 = balanced (1:1)
                   2.0 = 2:1 negative:positive

        For multi-class classification:
            ratio: Ignored. All classes are oversampled to match the largest class.
        """
        grouped_by_class = {}
        for path, class_idx in self.samples:
            if class_idx not in grouped_by_class:
                grouped_by_class[class_idx] = []
            grouped_by_class[class_idx].append((path, class_idx))

        print("Current classes:")
        for k in sorted(grouped_by_class.keys()):
            v = grouped_by_class[k]
            print(f"  {self.classes[k]}: {len(v)} samples")

        num_classes = len(self.classes)

        if num_classes == 2:
            # Binary classification - use ratio-based resampling
            targets = self._compute_binary_targets(grouped_by_class, ratio)
        else:
            # Multi-class - oversample all to match largest class
            targets = self._compute_multiclass_targets(grouped_by_class)

        # Apply oversampling
        for class_idx, samples in grouped_by_class.items():
            target = targets.get(class_idx, len(samples))
            if len(samples) >= target:
                continue
            num_extra = target - len(samples)
            repeats = num_extra // len(samples)
            remainder = num_extra - len(samples) * repeats
            extras = random.sample(samples, remainder) if remainder > 0 else []
            print(f"  Adding {num_extra} samples to '{self.classes[class_idx]}'")
            for _ in range(repeats):
                self.samples.extend(samples)
            self.samples.extend(extras)

        total_samples = sum(targets.values())
        print(f"Total samples after resampling: {total_samples}")

    def _compute_binary_targets(self, grouped_by_class: dict, ratio: float) -> dict:
        """Compute target counts for binary classification."""
        neg_idx = self.class_to_idx.get("negative", 0)
        pos_idx = self.class_to_idx.get("positive", 1)

        neg_count = len(grouped_by_class.get(neg_idx, []))
        pos_count = len(grouped_by_class.get(pos_idx, []))

        if ratio == 1.0:
            target_count = max(neg_count, pos_count)
            target_neg = target_count
            target_pos = target_count
        elif ratio > 1.0:
            target_neg = neg_count
            target_pos = int(neg_count / ratio)
            if target_pos < pos_count:
                target_pos = pos_count
                target_neg = int(pos_count * ratio)
        else:
            target_pos = pos_count
            target_neg = int(pos_count * ratio)
            if target_neg < neg_count:
                target_neg = neg_count
                target_pos = int(neg_count / ratio)

        print(f"Binary resampling: neg={target_neg}, pos={target_pos} (ratio={target_neg/max(target_pos,1):.2f}:1)")
        return {neg_idx: target_neg, pos_idx: target_pos}

    def _compute_multiclass_targets(self, grouped_by_class: dict) -> dict:
        """Compute target counts for multi-class classification (balance all classes)."""
        max_count = max(len(samples) for samples in grouped_by_class.values())
        targets = {class_idx: max_count for class_idx in grouped_by_class.keys()}
        print(f"Multi-class resampling: all classes to {max_count} samples")
        return targets

    def get_sample_weights(self, boost_factor: float) -> List[float]:
        """Return a per-sample weight list for use with WeightedRandomSampler.

        Samples whose path is in the boosted group get weight=boost_factor;
        all others get weight=1.0. Call this after resample() so repeated
        samples are included.

        Args:
            boost_factor: Weight multiplier for the boosted group. Must be > 0.
                         1.0 = all weights equal (equivalent to uniform sampling).

        Returns:
            List of float weights, one per entry in self.samples.
        """
        if not self._boost_paths or boost_factor == 1.0:
            return [1.0] * len(self.samples)
        return [
            boost_factor if path in self._boost_paths else 1.0
            for path, _ in self.samples
        ]

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

    def __getitem__(self, idx: int) -> tuple:
        video_path, label = self.samples[idx]

        try:
            if self.config.system.use_decord and _DECORD_AVAILABLE:
                frames = self._load_frames_decord(video_path)
            elif self.config.system.use_ffmpeg:
                frames = self._load_frames_ffmpeg(video_path)
            else:
                frames = self._load_frames_cv2(video_path)
        except Exception:
            import numpy as np
            frames = [np.zeros((self.img_size, self.img_size, 3), dtype=np.uint8)]

        # Pad if needed
        while len(frames) < self.frame_count:
            frames.append(frames[-1].copy())

        # Apply augmentations (before preprocessing)
        if self.transform is not None:
            frames = self.transform(images=frames)["images"]

        # Preprocess using shared pipeline (returns CPU tensor for DataLoader)
        video_tensor = preprocess_frames_batch(frames, self.img_size)

        if self.return_path:
            return video_tensor, torch.tensor(label, dtype=torch.long), video_path
        return video_tensor, torch.tensor(label, dtype=torch.long)

    def _load_frames_decord(self, video_path: str) -> list:
        """Load frames using decord — fastest option, no subprocess."""
        num_threads = getattr(self.config.system, 'decord_num_threads', 4)
        vr = decord.VideoReader(video_path, ctx=decord.cpu(0), num_threads=num_threads)
        total = len(vr)
        if total == self.frame_count:
            indices = list(range(total))
        else:
            fps = self.config.data.dataset_fps
            indices = self._get_frame_indices(total, fps)
        frames = vr.get_batch(indices).asnumpy()  # (T, H, W, 3) RGB uint8
        return [frames[i] for i in range(len(indices))]

    def _load_frames_cv2(self, video_path: str) -> list:
        """Load frames using OpenCV."""
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Failed to open video file: {video_path}")

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS) or self.config.data.dataset_fps
        indices = self._get_frame_indices(total_frames, fps)

        frames = []
        for i in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ret, frame = cap.read()
            if ret:
                # Convert BGR to RGB
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            elif frames:
                frames.append(frames[-1].copy())
        cap.release()
        return frames

    def _load_frames_ffmpeg(self, video_path: str) -> list:
        """Load frames using ffmpeg."""
        frames = list(get_frames_ffmpeg(video_path))
        total_frames = len(frames)

        if total_frames == self.frame_count:
            return frames

        # Use float-based sampling
        fps = self.config.data.dataset_fps
        indices = self._get_frame_indices(total_frames, fps)
        return [frames[i] for i in indices]

    def _get_frame_indices(self, total_frames: int, source_fps: float) -> list:
        """Get frame indices using float-based sampling for accurate timing."""
        if self.config.data.frame_sampling_method == "original_fps":
            # Sample at model's original FPS
            target_fps = self.config.model.original_fps
            indices = calculate_frame_indices(total_frames, source_fps, target_fps)

            # Take the last frame_count frames (for temporal consistency)
            if len(indices) >= self.frame_count:
                return indices[-self.frame_count:]
            return indices
        else:
            # Uniform sampling across the video
            return sample_frames_uniform(total_frames, self.frame_count)
