"""
==================================================================================
    Action Recognition Dataset - Python Loader
==================================================================================
    Project: Action Recognition Dataset
    Category: Video Data / Data Science
    
==================================================================================
    DEVELOPER INFORMATION
==================================================================================
    Website: RSK World (https://rskworld.in)
    Founded by: Molla Samser
    Designer & Tester: Rima Khatun
    
    Contact Information:
    - Email: help@rskworld.in
    - Email: support@rskworld.in
    - Phone: +91 93305 39277
    
==================================================================================
    COPYRIGHT
==================================================================================
    © 2026 RSK World. All Rights Reserved.
    
==================================================================================
"""

import cv2
import os
import json
import numpy as np
from pathlib import Path
from typing import List, Tuple, Dict, Optional


class ActionRecognitionDataset:
    """
    Action Recognition Dataset Loader
    
    A comprehensive video dataset loader for action recognition tasks.
    Supports loading videos, extracting frames, and preprocessing for
    various deep learning frameworks.
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Contact: help@rskworld.in
    """
    
    def __init__(
        self,
        root_dir: str,
        split: str = 'train',
        num_frames: int = 16,
        frame_size: Tuple[int, int] = (224, 224),
        normalize: bool = True
    ):
        """
        Initialize the Action Recognition Dataset loader.
        
        Args:
            root_dir: Path to the dataset root directory
            split: Dataset split ('train', 'val', or 'test')
            num_frames: Number of frames to sample from each video
            frame_size: Target frame size (height, width)
            normalize: Whether to normalize pixel values to [0, 1]
        """
        self.root_dir = Path(root_dir)
        self.split = split
        self.num_frames = num_frames
        self.frame_size = frame_size
        self.normalize = normalize
        
        # Load class labels
        self.class_labels = self._load_class_labels()
        self.num_classes = len(self.class_labels)
        
        # Build sample index
        self.samples = self._build_sample_index()
        
        print(f"[RSK World] Loaded {len(self.samples)} samples from {split} split")
        print(f"[RSK World] Number of classes: {self.num_classes}")
    
    def _load_class_labels(self) -> Dict[int, str]:
        """Load class labels from annotation file."""
        label_file = self.root_dir / 'annotations' / 'class_labels.json'
        
        if label_file.exists():
            with open(label_file, 'r') as f:
                data = json.load(f)
                return {int(k): v for k, v in data['classes'].items()}
        else:
            # Infer from directory structure
            split_dir = self.root_dir / self.split
            if split_dir.exists():
                classes = sorted([d.name for d in split_dir.iterdir() if d.is_dir()])
                return {i: c for i, c in enumerate(classes)}
            return {}
    
    def _build_sample_index(self) -> List[Tuple[str, int]]:
        """Build index of all video samples with labels."""
        samples = []
        split_dir = self.root_dir / self.split
        
        if not split_dir.exists():
            print(f"[RSK World] Warning: Split directory not found: {split_dir}")
            return samples
        
        for class_idx, class_name in self.class_labels.items():
            class_dir = split_dir / class_name
            
            if not class_dir.exists():
                continue
            
            for video_file in class_dir.iterdir():
                if video_file.suffix.lower() in ['.mp4', '.avi', '.mov', '.mkv']:
                    samples.append((str(video_file), class_idx))
        
        return samples
    
    def load_video(self, video_path: str) -> np.ndarray:
        """
        Load and preprocess a video file.
        
        Args:
            video_path: Path to the video file
            
        Returns:
            numpy array of shape (num_frames, height, width, channels)
        """
        cap = cv2.VideoCapture(video_path)
        frames = []
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Resize frame
            frame = cv2.resize(frame, self.frame_size)
            
            # Convert BGR to RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            frames.append(frame)
        
        cap.release()
        
        if len(frames) == 0:
            # Return zeros if video is empty
            return np.zeros(
                (self.num_frames, self.frame_size[0], self.frame_size[1], 3),
                dtype=np.float32 if self.normalize else np.uint8
            )
        
        # Sample frames uniformly
        indices = np.linspace(0, len(frames) - 1, self.num_frames).astype(int)
        sampled_frames = np.array([frames[i] for i in indices])
        
        # Normalize if required
        if self.normalize:
            sampled_frames = sampled_frames.astype(np.float32) / 255.0
        
        return sampled_frames
    
    def __len__(self) -> int:
        """Return the total number of samples."""
        return len(self.samples)
    
    def __getitem__(self, idx: int) -> Tuple[np.ndarray, int]:
        """
        Get a sample by index.
        
        Args:
            idx: Sample index
            
        Returns:
            Tuple of (frames, label)
        """
        video_path, label = self.samples[idx]
        frames = self.load_video(video_path)
        return frames, label
    
    def get_class_name(self, class_idx: int) -> str:
        """Get class name by index."""
        return self.class_labels.get(class_idx, f"unknown_{class_idx}")
    
    def get_class_distribution(self) -> Dict[str, int]:
        """Get distribution of samples across classes."""
        distribution = {name: 0 for name in self.class_labels.values()}
        
        for _, label in self.samples:
            class_name = self.class_labels[label]
            distribution[class_name] += 1
        
        return distribution


def create_data_generator(
    root_dir: str,
    split: str = 'train',
    batch_size: int = 8,
    num_frames: int = 16,
    shuffle: bool = True
):
    """
    Create a data generator for batch processing.
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    
    Args:
        root_dir: Path to the dataset root directory
        split: Dataset split ('train', 'val', or 'test')
        batch_size: Number of samples per batch
        num_frames: Number of frames to sample from each video
        shuffle: Whether to shuffle samples
        
    Yields:
        Tuple of (batch_frames, batch_labels)
    """
    dataset = ActionRecognitionDataset(
        root_dir=root_dir,
        split=split,
        num_frames=num_frames
    )
    
    indices = np.arange(len(dataset))
    
    while True:
        if shuffle:
            np.random.shuffle(indices)
        
        for start_idx in range(0, len(indices), batch_size):
            batch_indices = indices[start_idx:start_idx + batch_size]
            
            batch_frames = []
            batch_labels = []
            
            for idx in batch_indices:
                frames, label = dataset[idx]
                batch_frames.append(frames)
                batch_labels.append(label)
            
            yield np.array(batch_frames), np.array(batch_labels)


# Example usage
if __name__ == '__main__':
    print("=" * 60)
    print("Action Recognition Dataset Loader")
    print("RSK World (https://rskworld.in)")
    print("Founder: Molla Samser | Designer: Rima Khatun")
    print("Contact: help@rskworld.in | +91 93305 39277")
    print("=" * 60)
    
    # Example: Load training dataset
    # dataset = ActionRecognitionDataset(
    #     root_dir='./action-recognition',
    #     split='train',
    #     num_frames=16
    # )
    
    # # Get a sample
    # frames, label = dataset[0]
    # print(f"Frames shape: {frames.shape}")
    # print(f"Label: {label} ({dataset.get_class_name(label)})")
    
    # # Get class distribution
    # distribution = dataset.get_class_distribution()
    # print(f"Class distribution: {distribution}")
    
    print("\nDataset loader ready!")
    print("For usage examples, see README.md")