"""
==================================================================================
    Action Recognition Dataset - Advanced Visualization Tools
==================================================================================
    Project: Action Recognition Dataset
    
    Visualization tools including:
    - Dataset statistics dashboard
    - Video frame visualization
    - Class distribution charts
    - Motion heatmaps
    - Dataset explorer
    
==================================================================================
    DEVELOPER INFORMATION
==================================================================================
    Website: RSK World (https://rskworld.in)
    Founded by: Molla Samser
    Designer & Tester: Rima Khatun
    Contact: help@rskworld.in | +91 93305 39277
    
    (c) 2026 RSK World. All Rights Reserved.
==================================================================================
"""

import sys
import os
from pathlib import Path
from collections import defaultdict
import json

try:
    import cv2
    import numpy as np
except ImportError:
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python", "numpy"])
    import cv2
    import numpy as np


# ==================================================================================
# Configuration
# ==================================================================================

class VisConfig:
    """Visualization configuration"""
    DATA_DIR = "sample_data"
    OUTPUT_DIR = "visualizations"
    THUMBNAIL_SIZE = (160, 120)
    GRID_COLS = 4
    COLORS = {
        'primary': (255, 212, 0),      # Cyan (BGR)
        'secondary': (0, 212, 255),    # Yellow (BGR)
        'background': (15, 15, 20),
        'text': (255, 255, 255),
        'accent': (0, 255, 136)        # Green
    }


# ==================================================================================
# Dataset Statistics
# ==================================================================================

class DatasetStats:
    """
    Calculate and display dataset statistics
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, data_dir):
        self.data_dir = Path(data_dir)
        self.stats = {}
        
    def calculate_stats(self):
        """Calculate comprehensive dataset statistics"""
        stats = {
            'splits': {},
            'total_videos': 0,
            'total_frames': 0,
            'total_duration_seconds': 0,
            'classes': set(),
            'class_distribution': defaultdict(lambda: defaultdict(int)),
            'video_properties': []
        }
        
        for split in ['train', 'val', 'test']:
            split_dir = self.data_dir / split
            
            if not split_dir.exists():
                continue
            
            split_stats = {
                'videos': 0,
                'classes': 0,
                'class_counts': {}
            }
            
            for class_dir in sorted(split_dir.iterdir()):
                if not class_dir.is_dir():
                    continue
                
                class_name = class_dir.name
                stats['classes'].add(class_name)
                
                video_count = 0
                for video_file in class_dir.iterdir():
                    if video_file.suffix.lower() in ['.mp4', '.avi', '.mov', '.mkv']:
                        video_count += 1
                        stats['total_videos'] += 1
                        
                        # Get video properties
                        props = self._get_video_properties(str(video_file))
                        if props:
                            stats['video_properties'].append(props)
                            stats['total_frames'] += props['frame_count']
                            stats['total_duration_seconds'] += props['duration']
                
                split_stats['class_counts'][class_name] = video_count
                stats['class_distribution'][split][class_name] = video_count
                split_stats['videos'] += video_count
            
            split_stats['classes'] = len(split_stats['class_counts'])
            stats['splits'][split] = split_stats
        
        stats['classes'] = sorted(stats['classes'])
        self.stats = stats
        return stats
    
    def _get_video_properties(self, video_path):
        """Get video file properties"""
        cap = cv2.VideoCapture(video_path)
        
        if not cap.isOpened():
            return None
        
        props = {
            'path': video_path,
            'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
            'fps': cap.get(cv2.CAP_PROP_FPS),
            'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
            'duration': 0
        }
        
        if props['fps'] > 0:
            props['duration'] = props['frame_count'] / props['fps']
        
        cap.release()
        return props
    
    def print_stats(self):
        """Print statistics summary"""
        if not self.stats:
            self.calculate_stats()
        
        print("\n" + "=" * 60)
        print("DATASET STATISTICS")
        print("=" * 60)
        print(f"Website: https://rskworld.in")
        print(f"Founder: Molla Samser")
        print("=" * 60)
        
        print(f"\n{'Total Videos:':<25} {self.stats['total_videos']}")
        print(f"{'Total Frames:':<25} {self.stats['total_frames']}")
        print(f"{'Total Duration:':<25} {self.stats['total_duration_seconds']:.1f} seconds")
        print(f"{'Number of Classes:':<25} {len(self.stats['classes'])}")
        
        print(f"\n{'Split Statistics:'}")
        print("-" * 50)
        
        for split, split_stats in self.stats.get('splits', {}).items():
            print(f"\n  {split.upper()}")
            print(f"    Videos: {split_stats['videos']}")
            print(f"    Classes: {split_stats['classes']}")
            
            print(f"    Distribution:")
            for cls, count in sorted(split_stats.get('class_counts', {}).items()):
                bar = '#' * min(count * 2, 20)
                print(f"      {cls:<15} {count:>4} {bar}")
        
        if self.stats.get('video_properties'):
            props = self.stats['video_properties']
            resolutions = [(p['width'], p['height']) for p in props]
            unique_res = list(set(resolutions))
            
            print(f"\n{'Video Properties:'}")
            print("-" * 50)
            print(f"  {'Unique Resolutions:':<20} {len(unique_res)}")
            print(f"  {'Most Common:':<20} {max(set(resolutions), key=resolutions.count)}")
            
            fps_values = [p['fps'] for p in props if p['fps'] > 0]
            if fps_values:
                print(f"  {'Average FPS:':<20} {np.mean(fps_values):.1f}")
            
            durations = [p['duration'] for p in props]
            if durations:
                print(f"  {'Avg Duration:':<20} {np.mean(durations):.2f}s")
                print(f"  {'Min Duration:':<20} {min(durations):.2f}s")
                print(f"  {'Max Duration:':<20} {max(durations):.2f}s")
        
        print("\n" + "=" * 60)
        
        return self.stats


# ==================================================================================
# Visualization Generator
# ==================================================================================

class DatasetVisualizer:
    """
    Generate visualizations for the dataset
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, config=None):
        self.config = config or VisConfig()
        Path(self.config.OUTPUT_DIR).mkdir(exist_ok=True)
        
    def create_sample_grid(self, data_dir, split='train', output_path=None):
        """Create a grid of sample frames from each class"""
        data_dir = Path(data_dir) / split
        
        if not data_dir.exists():
            print(f"[RSK World] Directory not found: {data_dir}")
            return
        
        classes = sorted([d.name for d in data_dir.iterdir() if d.is_dir()])
        
        if not classes:
            print("[RSK World] No classes found!")
            return
        
        # Calculate grid size
        n_classes = len(classes)
        cols = min(self.config.GRID_COLS, n_classes)
        rows = (n_classes + cols - 1) // cols
        
        tw, th = self.config.THUMBNAIL_SIZE
        padding = 10
        header_height = 40
        
        # Create canvas
        canvas_width = cols * (tw + padding) + padding
        canvas_height = rows * (th + padding + 25) + padding + header_height
        
        canvas = np.full((canvas_height, canvas_width, 3), self.config.COLORS['background'], dtype=np.uint8)
        
        # Add header
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(canvas, "RSK", (padding, 28), font, 0.8, self.config.COLORS['text'], 2, cv2.LINE_AA)
        cv2.putText(canvas, "World", (padding + 55, 28), font, 0.8, self.config.COLORS['secondary'], 2, cv2.LINE_AA)
        cv2.putText(canvas, f" - {split.upper()} Samples", (padding + 130, 28), font, 0.6, 
                   self.config.COLORS['text'], 1, cv2.LINE_AA)
        
        # Add samples
        for idx, class_name in enumerate(classes):
            class_dir = data_dir / class_name
            
            # Get first video
            videos = list(class_dir.glob('*.mp4')) + list(class_dir.glob('*.avi'))
            
            if videos:
                # Extract middle frame
                cap = cv2.VideoCapture(str(videos[0]))
                total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                cap.set(cv2.CAP_PROP_POS_FRAMES, total // 2)
                ret, frame = cap.read()
                cap.release()
                
                if ret:
                    # Resize frame
                    frame = cv2.resize(frame, self.config.THUMBNAIL_SIZE)
                else:
                    frame = np.zeros((th, tw, 3), dtype=np.uint8)
            else:
                frame = np.zeros((th, tw, 3), dtype=np.uint8)
            
            # Calculate position
            col = idx % cols
            row = idx // cols
            x = padding + col * (tw + padding)
            y = header_height + padding + row * (th + padding + 25)
            
            # Place frame
            canvas[y:y+th, x:x+tw] = frame
            
            # Add border
            cv2.rectangle(canvas, (x-1, y-1), (x+tw, y+th), self.config.COLORS['primary'], 1)
            
            # Add class label
            label_y = y + th + 18
            cv2.putText(canvas, class_name.capitalize(), (x, label_y), font, 0.45,
                       self.config.COLORS['text'], 1, cv2.LINE_AA)
        
        # Save
        if output_path is None:
            output_path = Path(self.config.OUTPUT_DIR) / f"sample_grid_{split}.png"
        
        cv2.imwrite(str(output_path), canvas)
        print(f"[RSK World] Sample grid saved: {output_path}")
        
        return canvas
    
    def create_motion_heatmap(self, video_path, output_path=None):
        """Create motion heatmap from video"""
        cap = cv2.VideoCapture(video_path)
        
        if not cap.isOpened():
            print(f"[RSK World] Cannot open video: {video_path}")
            return
        
        # Read first frame
        ret, prev_frame = cap.read()
        if not ret:
            cap.release()
            return
        
        prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
        
        # Accumulate motion
        h, w = prev_gray.shape
        motion_accumulator = np.zeros((h, w), dtype=np.float32)
        frame_count = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            diff = cv2.absdiff(prev_gray, gray)
            motion_accumulator += diff.astype(np.float32)
            
            prev_gray = gray
            frame_count += 1
        
        cap.release()
        
        if frame_count == 0:
            return
        
        # Normalize
        motion_accumulator /= frame_count
        motion_normalized = (motion_accumulator / motion_accumulator.max() * 255).astype(np.uint8)
        
        # Apply colormap
        heatmap = cv2.applyColorMap(motion_normalized, cv2.COLORMAP_JET)
        
        # Add RSK World branding
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(heatmap, "RSKWorld.in", (10, 25), font, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
        cv2.putText(heatmap, "Motion Heatmap", (10, h - 10), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
        
        # Save
        if output_path is None:
            video_name = Path(video_path).stem
            output_path = Path(self.config.OUTPUT_DIR) / f"heatmap_{video_name}.png"
        
        cv2.imwrite(str(output_path), heatmap)
        print(f"[RSK World] Motion heatmap saved: {output_path}")
        
        return heatmap
    
    def create_class_distribution_chart(self, stats, output_path=None):
        """Create class distribution visualization"""
        if not stats.get('classes'):
            print("[RSK World] No classes found in stats!")
            return
        
        classes = stats['classes']
        n_classes = len(classes)
        
        # Chart dimensions
        chart_width = 800
        chart_height = 400
        margin = 60
        bar_width = (chart_width - 2 * margin) // n_classes - 5
        
        # Create canvas
        canvas = np.full((chart_height + 100, chart_width, 3), self.config.COLORS['background'], dtype=np.uint8)
        
        # Get max count for scaling
        max_count = 0
        for split_data in stats.get('class_distribution', {}).values():
            for count in split_data.values():
                max_count = max(max_count, count)
        
        if max_count == 0:
            max_count = 1
        
        font = cv2.FONT_HERSHEY_SIMPLEX
        
        # Title
        cv2.putText(canvas, "RSKWorld.in - Class Distribution", (margin, 35), font, 0.8, 
                   self.config.COLORS['primary'], 2, cv2.LINE_AA)
        
        # Draw bars for each class
        splits = list(stats.get('class_distribution', {}).keys())
        colors = [(255, 100, 100), (100, 255, 100), (100, 100, 255)]  # BGR
        
        for i, class_name in enumerate(classes):
            x = margin + i * (bar_width + 5) + bar_width // 2
            
            for j, split in enumerate(splits):
                count = stats['class_distribution'].get(split, {}).get(class_name, 0)
                bar_height = int((count / max_count) * (chart_height - 2 * margin - 50))
                
                bar_x = x - bar_width // 2 + j * (bar_width // len(splits))
                bar_w = bar_width // len(splits) - 2
                bar_y = chart_height - margin - bar_height
                
                color = colors[j % len(colors)]
                cv2.rectangle(canvas, (bar_x, bar_y), (bar_x + bar_w, chart_height - margin), color, -1)
            
            # Class label
            label = class_name[:8] + ".." if len(class_name) > 10 else class_name
            cv2.putText(canvas, label, (x - 30, chart_height - margin + 20), font, 0.35,
                       self.config.COLORS['text'], 1, cv2.LINE_AA)
        
        # Legend
        legend_y = chart_height + 50
        for j, split in enumerate(splits):
            cv2.rectangle(canvas, (margin + j * 100, legend_y), (margin + j * 100 + 20, legend_y + 15), 
                         colors[j % len(colors)], -1)
            cv2.putText(canvas, split.capitalize(), (margin + j * 100 + 25, legend_y + 12), font, 0.4,
                       self.config.COLORS['text'], 1, cv2.LINE_AA)
        
        # Save
        if output_path is None:
            output_path = Path(self.config.OUTPUT_DIR) / "class_distribution.png"
        
        cv2.imwrite(str(output_path), canvas)
        print(f"[RSK World] Distribution chart saved: {output_path}")
        
        return canvas


# ==================================================================================
# Main
# ==================================================================================

def main():
    """Generate dataset visualizations"""
    print("=" * 60)
    print("Action Recognition Dataset - Visualization Tools")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("Designer: Rima Khatun")
    print("Contact: help@rskworld.in | +91 93305 39277")
    print("=" * 60)
    
    config = VisConfig()
    
    # Calculate statistics
    stats_calc = DatasetStats(config.DATA_DIR)
    stats = stats_calc.calculate_stats()
    stats_calc.print_stats()
    
    # Create visualizations
    visualizer = DatasetVisualizer(config)
    
    # Sample grids
    for split in ['train', 'val', 'test']:
        visualizer.create_sample_grid(config.DATA_DIR, split)
    
    # Class distribution chart
    visualizer.create_class_distribution_chart(stats)
    
    # Motion heatmaps for first video in each class
    train_dir = Path(config.DATA_DIR) / 'train'
    if train_dir.exists():
        for class_dir in list(train_dir.iterdir())[:3]:  # First 3 classes
            if class_dir.is_dir():
                videos = list(class_dir.glob('*.mp4'))
                if videos:
                    visualizer.create_motion_heatmap(str(videos[0]))
    
    print("\n" + "=" * 60)
    print("Visualizations complete!")
    print(f"Output directory: {config.OUTPUT_DIR}")
    print("\nThank you for using RSK World!")
    print("Visit: https://rskworld.in")
    print("=" * 60)


if __name__ == "__main__":
    main()