"""
==================================================================================
    Action Recognition Dataset - Model Benchmark & Comparison
==================================================================================
    Project: Action Recognition Dataset
    
    Comprehensive benchmarking tool with:
    - Multiple model comparison
    - Performance metrics (accuracy, precision, recall, F1)
    - Speed benchmarks (FPS, inference time)
    - Confusion matrix generation
    - Export results to JSON/CSV
    - Visualization of results
    
==================================================================================
    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
import json
import time
from pathlib import Path
from datetime import datetime
from collections import defaultdict

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 BenchmarkConfig:
    """Benchmark configuration"""
    DATA_DIR = "sample_data"
    OUTPUT_DIR = "benchmark_results"
    NUM_FRAMES = 16
    FRAME_SIZE = (112, 112)
    NUM_RUNS = 3  # Number of runs for averaging
    WARMUP_RUNS = 1


# ==================================================================================
# Metrics Calculator
# ==================================================================================

class MetricsCalculator:
    """
    Calculate evaluation metrics
    
    RSK World (https://rskworld.in)
    """
    
    @staticmethod
    def accuracy(y_true, y_pred):
        """Calculate accuracy"""
        correct = sum(t == p for t, p in zip(y_true, y_pred))
        return correct / len(y_true) if y_true else 0
    
    @staticmethod
    def precision_recall_f1(y_true, y_pred, classes):
        """Calculate precision, recall, F1 per class"""
        metrics = {}
        
        for cls in classes:
            tp = sum(1 for t, p in zip(y_true, y_pred) if t == cls and p == cls)
            fp = sum(1 for t, p in zip(y_true, y_pred) if t != cls and p == cls)
            fn = sum(1 for t, p in zip(y_true, y_pred) if t == cls and p != cls)
            
            precision = tp / (tp + fp) if (tp + fp) > 0 else 0
            recall = tp / (tp + fn) if (tp + fn) > 0 else 0
            f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
            
            metrics[cls] = {
                'precision': round(precision, 4),
                'recall': round(recall, 4),
                'f1': round(f1, 4),
                'support': sum(1 for t in y_true if t == cls)
            }
        
        # Macro averages
        metrics['macro_avg'] = {
            'precision': round(np.mean([m['precision'] for m in metrics.values() if isinstance(m, dict) and 'precision' in m]), 4),
            'recall': round(np.mean([m['recall'] for m in metrics.values() if isinstance(m, dict) and 'recall' in m]), 4),
            'f1': round(np.mean([m['f1'] for m in metrics.values() if isinstance(m, dict) and 'f1' in m]), 4)
        }
        
        return metrics
    
    @staticmethod
    def confusion_matrix(y_true, y_pred, classes):
        """Generate confusion matrix"""
        n_classes = len(classes)
        class_to_idx = {c: i for i, c in enumerate(classes)}
        
        matrix = [[0] * n_classes for _ in range(n_classes)]
        
        for t, p in zip(y_true, y_pred):
            if t in class_to_idx and p in class_to_idx:
                matrix[class_to_idx[t]][class_to_idx[p]] += 1
        
        return matrix


# ==================================================================================
# Simple Motion-Based Model (for benchmark demo)
# ==================================================================================

class MotionModel:
    """
    Motion-based action classifier
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, name="MotionModel"):
        self.name = name
        self.classes = []
        
    def load_video(self, video_path, config):
        """Load and preprocess video"""
        cap = cv2.VideoCapture(video_path)
        frames = []
        
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total_frames == 0:
            cap.release()
            return None
        
        indices = np.linspace(0, total_frames - 1, config.NUM_FRAMES).astype(int)
        
        for i in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ret, frame = cap.read()
            if ret:
                frame = cv2.resize(frame, config.FRAME_SIZE)
                frames.append(frame)
        
        cap.release()
        return frames if len(frames) > 1 else None
    
    def predict(self, frames):
        """Predict action from frames"""
        if frames is None or len(frames) < 2:
            return "unknown", 0.5
        
        # Calculate motion
        motion_scores = []
        prev_gray = None
        
        for frame in frames:
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            if prev_gray is not None:
                diff = cv2.absdiff(prev_gray, gray)
                motion = np.mean(diff)
                motion_scores.append(motion)
            
            prev_gray = gray
        
        avg_motion = np.mean(motion_scores) if motion_scores else 0
        
        # Classify
        if avg_motion < 2:
            return "standing", 0.8
        elif avg_motion < 5:
            return "walking", 0.75
        elif avg_motion < 10:
            return "running", 0.7
        else:
            return "jumping", 0.65


class EnhancedMotionModel(MotionModel):
    """
    Enhanced motion model with optical flow
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self):
        super().__init__("EnhancedMotionModel")
    
    def predict(self, frames):
        """Predict with optical flow analysis"""
        if frames is None or len(frames) < 2:
            return "unknown", 0.5
        
        # Motion analysis
        motion_scores = []
        flow_magnitudes = []
        
        prev_gray = cv2.cvtColor(frames[0], cv2.COLOR_BGR2GRAY)
        
        for frame in frames[1:]:
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # Frame difference
            diff = cv2.absdiff(prev_gray, gray)
            motion_scores.append(np.mean(diff))
            
            # Simplified flow (edge detection as proxy)
            edges = cv2.Canny(gray, 50, 150)
            flow_magnitudes.append(np.mean(edges))
            
            prev_gray = gray
        
        avg_motion = np.mean(motion_scores)
        avg_flow = np.mean(flow_magnitudes)
        motion_var = np.var(motion_scores)
        
        # Enhanced classification
        if avg_motion < 2:
            if avg_flow < 10:
                return "standing", 0.85
            else:
                return "sitting", 0.7
        elif avg_motion < 5:
            if motion_var > 1:
                return "waving", 0.7
            else:
                return "walking", 0.8
        elif avg_motion < 10:
            if motion_var > 3:
                return "dancing", 0.7
            else:
                return "running", 0.75
        elif avg_motion < 20:
            return "exercising", 0.7
        else:
            return "jumping", 0.65


# ==================================================================================
# Benchmark Runner
# ==================================================================================

class BenchmarkRunner:
    """
    Run benchmarks on multiple models
    
    RSK World (https://rskworld.in)
    Founder: Molla Samser
    """
    
    def __init__(self, config=None):
        self.config = config or BenchmarkConfig()
        self.metrics = MetricsCalculator()
        self.results = {}
        
        # Create output directory
        Path(self.config.OUTPUT_DIR).mkdir(exist_ok=True)
    
    def load_dataset(self, split='test'):
        """Load dataset for evaluation"""
        data_dir = Path(self.config.DATA_DIR) / split
        
        if not data_dir.exists():
            print(f"[RSK World] Dataset not found: {data_dir}")
            return [], []
        
        samples = []
        labels = []
        
        for class_dir in sorted(data_dir.iterdir()):
            if class_dir.is_dir():
                for video_file in class_dir.iterdir():
                    if video_file.suffix.lower() in ['.mp4', '.avi', '.mov']:
                        samples.append(str(video_file))
                        labels.append(class_dir.name)
        
        return samples, labels
    
    def run_benchmark(self, model, samples, true_labels):
        """Run benchmark for a single model"""
        print(f"\n[RSK World] Benchmarking: {model.name}")
        print("-" * 50)
        
        predictions = []
        confidences = []
        inference_times = []
        
        # Warmup
        for _ in range(self.config.WARMUP_RUNS):
            if samples:
                frames = model.load_video(samples[0], self.config)
                model.predict(frames)
        
        # Benchmark runs
        for run in range(self.config.NUM_RUNS):
            print(f"  Run {run + 1}/{self.config.NUM_RUNS}...")
            
            run_predictions = []
            run_times = []
            
            for video_path in samples:
                # Load video
                frames = model.load_video(video_path, self.config)
                
                # Time prediction
                start_time = time.time()
                pred, conf = model.predict(frames)
                end_time = time.time()
                
                run_predictions.append(pred)
                run_times.append(end_time - start_time)
            
            predictions.append(run_predictions)
            inference_times.extend(run_times)
        
        # Use last run predictions for accuracy
        final_predictions = predictions[-1] if predictions else []
        
        # Calculate metrics
        classes = sorted(set(true_labels))
        
        accuracy = self.metrics.accuracy(true_labels, final_predictions)
        class_metrics = self.metrics.precision_recall_f1(true_labels, final_predictions, classes)
        confusion = self.metrics.confusion_matrix(true_labels, final_predictions, classes)
        
        # Speed metrics
        avg_time = np.mean(inference_times) if inference_times else 0
        fps = 1.0 / avg_time if avg_time > 0 else 0
        
        results = {
            'model_name': model.name,
            'accuracy': round(accuracy, 4),
            'class_metrics': class_metrics,
            'confusion_matrix': confusion,
            'classes': classes,
            'speed': {
                'avg_inference_time_ms': round(avg_time * 1000, 2),
                'fps': round(fps, 2),
                'total_samples': len(samples)
            },
            'config': {
                'num_frames': self.config.NUM_FRAMES,
                'frame_size': self.config.FRAME_SIZE,
                'num_runs': self.config.NUM_RUNS
            }
        }
        
        # Print summary
        print(f"\n  Results for {model.name}:")
        print(f"    Accuracy: {accuracy * 100:.2f}%")
        print(f"    F1 (macro): {class_metrics.get('macro_avg', {}).get('f1', 0) * 100:.2f}%")
        print(f"    Avg Inference: {avg_time * 1000:.2f}ms")
        print(f"    FPS: {fps:.2f}")
        
        return results
    
    def compare_models(self, models):
        """Compare multiple models"""
        print("\n" + "=" * 60)
        print("Action Recognition - Model Benchmark")
        print("=" * 60)
        print("Website: https://rskworld.in")
        print("Founder: Molla Samser")
        print("Designer: Rima Khatun")
        print("=" * 60)
        
        # Load dataset
        print("\n[RSK World] Loading test dataset...")
        samples, labels = self.load_dataset('test')
        
        if not samples:
            print("[ERROR] No test samples found!")
            print("Please run: python download_youtube_videos.py")
            print("Then run: python process_downloaded.py")
            return {}
        
        print(f"[RSK World] Loaded {len(samples)} samples")
        print(f"[RSK World] Classes: {sorted(set(labels))}")
        
        # Benchmark each model
        all_results = {}
        
        for model in models:
            results = self.run_benchmark(model, samples, labels)
            all_results[model.name] = results
            self.results[model.name] = results
        
        # Generate comparison
        self._generate_comparison(all_results)
        
        # Save results
        self._save_results(all_results)
        
        return all_results
    
    def _generate_comparison(self, results):
        """Generate comparison summary"""
        print("\n" + "=" * 60)
        print("BENCHMARK COMPARISON SUMMARY")
        print("=" * 60)
        
        # Table header
        print(f"\n{'Model':<25} {'Accuracy':>10} {'F1':>10} {'FPS':>10} {'Time(ms)':>10}")
        print("-" * 65)
        
        for name, res in results.items():
            acc = res.get('accuracy', 0) * 100
            f1 = res.get('class_metrics', {}).get('macro_avg', {}).get('f1', 0) * 100
            fps = res.get('speed', {}).get('fps', 0)
            time_ms = res.get('speed', {}).get('avg_inference_time_ms', 0)
            
            print(f"{name:<25} {acc:>9.2f}% {f1:>9.2f}% {fps:>10.2f} {time_ms:>10.2f}")
        
        print("-" * 65)
        
        # Best model
        if results:
            best_acc = max(results.items(), key=lambda x: x[1].get('accuracy', 0))
            best_fps = max(results.items(), key=lambda x: x[1].get('speed', {}).get('fps', 0))
            
            print(f"\nBest Accuracy: {best_acc[0]} ({best_acc[1].get('accuracy', 0) * 100:.2f}%)")
            print(f"Best Speed: {best_fps[0]} ({best_fps[1].get('speed', {}).get('fps', 0):.2f} FPS)")
    
    def _save_results(self, results):
        """Save benchmark results"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Save JSON
        json_path = Path(self.config.OUTPUT_DIR) / f"benchmark_{timestamp}.json"
        with open(json_path, 'w') as f:
            json.dump({
                'timestamp': datetime.now().isoformat(),
                'developer': 'RSK World (rskworld.in)',
                'results': results
            }, f, indent=4)
        
        print(f"\n[RSK World] Results saved to: {json_path}")
        
        # Save CSV summary
        csv_path = Path(self.config.OUTPUT_DIR) / f"benchmark_{timestamp}.csv"
        with open(csv_path, 'w') as f:
            f.write("Model,Accuracy,F1,Precision,Recall,FPS,Inference_ms\n")
            for name, res in results.items():
                acc = res.get('accuracy', 0)
                macro = res.get('class_metrics', {}).get('macro_avg', {})
                f1 = macro.get('f1', 0)
                prec = macro.get('precision', 0)
                rec = macro.get('recall', 0)
                fps = res.get('speed', {}).get('fps', 0)
                time_ms = res.get('speed', {}).get('avg_inference_time_ms', 0)
                f.write(f"{name},{acc},{f1},{prec},{rec},{fps},{time_ms}\n")
        
        print(f"[RSK World] CSV saved to: {csv_path}")


# ==================================================================================
# Generate Visualization
# ==================================================================================

def generate_html_report(results, output_path="benchmark_results/report.html"):
    """Generate HTML benchmark report"""
    
    html = f"""<!DOCTYPE html>
<html>
<head>
    <title>Benchmark Report - RSK World</title>
    <style>
        body {{ font-family: 'Segoe UI', sans-serif; margin: 40px; background: #0a0a0f; color: #fff; }}
        .header {{ text-align: center; margin-bottom: 40px; }}
        .header h1 {{ color: #00d4ff; }}
        .header .brand {{ color: #ffd400; }}
        table {{ width: 100%; border-collapse: collapse; margin: 20px 0; }}
        th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #333; }}
        th {{ background: #1a1a25; color: #00d4ff; }}
        tr:hover {{ background: #1a1a25; }}
        .best {{ color: #00ff88; font-weight: bold; }}
        .metric {{ text-align: right; }}
        .footer {{ text-align: center; margin-top: 40px; color: #666; }}
    </style>
</head>
<body>
    <div class="header">
        <h1>Action Recognition <span class="brand">Benchmark Report</span></h1>
        <p>Generated by RSK World (rskworld.in)</p>
        <p>Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
    </div>
    
    <h2>Model Comparison</h2>
    <table>
        <tr>
            <th>Model</th>
            <th class="metric">Accuracy</th>
            <th class="metric">F1 Score</th>
            <th class="metric">Precision</th>
            <th class="metric">Recall</th>
            <th class="metric">FPS</th>
            <th class="metric">Inference (ms)</th>
        </tr>
"""
    
    # Find best values
    best_acc = max(r.get('accuracy', 0) for r in results.values())
    best_fps = max(r.get('speed', {}).get('fps', 0) for r in results.values())
    
    for name, res in results.items():
        acc = res.get('accuracy', 0) * 100
        macro = res.get('class_metrics', {}).get('macro_avg', {})
        f1 = macro.get('f1', 0) * 100
        prec = macro.get('precision', 0) * 100
        rec = macro.get('recall', 0) * 100
        fps = res.get('speed', {}).get('fps', 0)
        time_ms = res.get('speed', {}).get('avg_inference_time_ms', 0)
        
        acc_class = 'best' if res.get('accuracy', 0) == best_acc else ''
        fps_class = 'best' if fps == best_fps else ''
        
        html += f"""        <tr>
            <td>{name}</td>
            <td class="metric {acc_class}">{acc:.2f}%</td>
            <td class="metric">{f1:.2f}%</td>
            <td class="metric">{prec:.2f}%</td>
            <td class="metric">{rec:.2f}%</td>
            <td class="metric {fps_class}">{fps:.2f}</td>
            <td class="metric">{time_ms:.2f}</td>
        </tr>
"""
    
    html += """    </table>
    
    <div class="footer">
        <p>&copy; 2026 RSK World. All Rights Reserved.</p>
        <p>Website: rskworld.in | Email: help@rskworld.in | Phone: +91 93305 39277</p>
        <p>Founder: Molla Samser | Designer: Rima Khatun</p>
    </div>
</body>
</html>
"""
    
    Path(output_path).parent.mkdir(exist_ok=True)
    with open(output_path, 'w') as f:
        f.write(html)
    
    print(f"[RSK World] HTML report saved to: {output_path}")


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

def main():
    """Run benchmark"""
    print("=" * 60)
    print("Action Recognition Dataset - Model Benchmark")
    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)
    
    # Create models to benchmark
    models = [
        MotionModel("BasicMotion"),
        EnhancedMotionModel(),
    ]
    
    # Run benchmark
    runner = BenchmarkRunner()
    results = runner.compare_models(models)
    
    # Generate HTML report
    if results:
        generate_html_report(results)
    
    print("\n" + "=" * 60)
    print("Benchmark complete!")
    print("Thank you for using RSK World!")
    print("Visit: https://rskworld.in")
    print("=" * 60)


if __name__ == "__main__":
    main()