"""
==================================================================================
    Action Recognition - Real-Time Webcam Predictor
==================================================================================
    Project: Action Recognition Dataset
    
    Real-time action recognition using webcam with:
    - Live video feed processing
    - Smooth action predictions
    - Confidence visualization
    - Recording capability
    - Multi-person detection (optional)
    
==================================================================================
    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 deque
import time
from datetime import datetime

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 Config:
    """Real-time predictor configuration"""
    # Webcam
    CAMERA_ID = 0
    FRAME_WIDTH = 640
    FRAME_HEIGHT = 480
    FPS = 30
    
    # Model
    MODEL_PATH = "checkpoints/best_model.pth"
    NUM_FRAMES = 16
    FRAME_SIZE = (112, 112)
    
    # Prediction
    SMOOTHING_WINDOW = 10
    CONFIDENCE_THRESHOLD = 0.5
    
    # Display
    SHOW_FPS = True
    SHOW_CONFIDENCE = True
    SHOW_SKELETON = True
    
    # Recording
    OUTPUT_DIR = "recordings"
    
    # RSK World branding
    LOGO_COLOR = (255, 212, 0)  # Cyan (BGR)
    TEXT_COLOR = (255, 255, 255)


# ==================================================================================
# Action Classes (without model)
# ==================================================================================

ACTION_CLASSES = [
    "walking", "running", "jumping", "waving", "sitting",
    "standing", "dancing", "exercising", "punching", "kicking",
    "yoga", "stretching", "boxing"
]


# ==================================================================================
# Simple Motion-Based Predictor
# ==================================================================================

class MotionBasedPredictor:
    """
    Simple motion-based action predictor using optical flow
    Works without a trained model
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, config):
        self.config = config
        self.prev_frame = None
        self.motion_history = deque(maxlen=30)
        self.position_history = deque(maxlen=30)
        
    def predict(self, frame):
        """Predict action based on motion analysis"""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        gray = cv2.GaussianBlur(gray, (21, 21), 0)
        
        if self.prev_frame is None:
            self.prev_frame = gray
            return "standing", 0.5
        
        # Calculate frame difference
        frame_diff = cv2.absdiff(self.prev_frame, gray)
        thresh = cv2.threshold(frame_diff, 25, 255, cv2.THRESH_BINARY)[1]
        thresh = cv2.dilate(thresh, None, iterations=2)
        
        # Find contours (motion areas)
        contours, _ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Calculate motion metrics
        total_motion = sum(cv2.contourArea(c) for c in contours)
        motion_ratio = total_motion / (frame.shape[0] * frame.shape[1])
        
        self.motion_history.append(motion_ratio)
        
        # Calculate motion statistics
        avg_motion = np.mean(self.motion_history)
        motion_variance = np.var(self.motion_history)
        
        # Determine action based on motion patterns
        action, confidence = self._classify_motion(avg_motion, motion_variance)
        
        self.prev_frame = gray
        return action, confidence
    
    def _classify_motion(self, avg_motion, variance):
        """Classify action based on motion metrics"""
        if avg_motion < 0.001:
            return "standing", 0.9
        elif avg_motion < 0.005:
            return "sitting", 0.7
        elif avg_motion < 0.02:
            return "walking", 0.8
        elif avg_motion < 0.05:
            if variance > 0.0001:
                return "waving", 0.7
            else:
                return "walking", 0.75
        elif avg_motion < 0.1:
            if variance > 0.001:
                return "dancing", 0.7
            else:
                return "running", 0.8
        else:
            if variance > 0.005:
                return "jumping", 0.75
            else:
                return "exercising", 0.7


# ==================================================================================
# Real-Time Predictor
# ==================================================================================

class RealtimePredictor:
    """
    Real-time action recognition from webcam
    
    RSK World (https://rskworld.in)
    Founder: Molla Samser
    """
    
    def __init__(self, config=None):
        self.config = config or Config()
        self.predictor = MotionBasedPredictor(self.config)
        self.fps_history = deque(maxlen=30)
        self.prediction_history = deque(maxlen=self.config.SMOOTHING_WINDOW)
        self.is_recording = False
        self.video_writer = None
        
        # Create output directory
        Path(self.config.OUTPUT_DIR).mkdir(exist_ok=True)
        
    def _add_overlay(self, frame, action, confidence, fps):
        """Add RSK World branding and information overlay"""
        h, w = frame.shape[:2]
        
        # Create semi-transparent overlay
        overlay = frame.copy()
        
        # Top bar
        cv2.rectangle(overlay, (0, 0), (w, 70), (20, 20, 30), -1)
        
        # Bottom bar
        cv2.rectangle(overlay, (0, h - 80), (w, h), (20, 20, 30), -1)
        
        # Blend overlay
        cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame)
        
        # Font
        font = cv2.FONT_HERSHEY_SIMPLEX
        
        # RSK World logo (top left)
        cv2.putText(frame, "RSK", (15, 35), font, 0.9, self.config.TEXT_COLOR, 2, cv2.LINE_AA)
        cv2.putText(frame, "World", (75, 35), font, 0.9, self.config.LOGO_COLOR, 2, cv2.LINE_AA)
        cv2.putText(frame, ".in", (165, 35), font, 0.6, (150, 150, 150), 1, cv2.LINE_AA)
        
        # Title (top center)
        title = "Action Recognition"
        title_size = cv2.getTextSize(title, font, 0.7, 2)[0]
        cv2.putText(frame, title, ((w - title_size[0]) // 2, 45), font, 0.7, 
                   self.config.TEXT_COLOR, 2, cv2.LINE_AA)
        
        # FPS (top right)
        if self.config.SHOW_FPS:
            fps_text = f"FPS: {fps:.1f}"
            cv2.putText(frame, fps_text, (w - 120, 35), font, 0.6, 
                       (100, 255, 100), 2, cv2.LINE_AA)
        
        # Recording indicator
        if self.is_recording:
            cv2.circle(frame, (w - 30, 25), 10, (0, 0, 255), -1)
            cv2.putText(frame, "REC", (w - 80, 32), font, 0.5, (0, 0, 255), 2, cv2.LINE_AA)
        
        # Action label (bottom center)
        action_text = action.upper()
        action_size = cv2.getTextSize(action_text, font, 1.2, 3)[0]
        cv2.putText(frame, action_text, ((w - action_size[0]) // 2, h - 45), 
                   font, 1.2, self.config.LOGO_COLOR, 3, cv2.LINE_AA)
        
        # Confidence bar (bottom)
        if self.config.SHOW_CONFIDENCE:
            bar_width = int(300 * confidence)
            bar_x = (w - 300) // 2
            
            # Background bar
            cv2.rectangle(frame, (bar_x, h - 25), (bar_x + 300, h - 15), (50, 50, 50), -1)
            
            # Confidence bar
            color = (0, 255, 0) if confidence > 0.7 else ((0, 255, 255) if confidence > 0.5 else (0, 0, 255))
            cv2.rectangle(frame, (bar_x, h - 25), (bar_x + bar_width, h - 15), color, -1)
            
            # Confidence text
            conf_text = f"{confidence * 100:.1f}%"
            cv2.putText(frame, conf_text, (bar_x + 310, h - 12), font, 0.5, 
                       self.config.TEXT_COLOR, 1, cv2.LINE_AA)
        
        # Instructions (bottom left)
        instructions = "Q: Quit | R: Record | S: Screenshot"
        cv2.putText(frame, instructions, (10, h - 12), font, 0.4, 
                   (150, 150, 150), 1, cv2.LINE_AA)
        
        return frame
    
    def _smooth_prediction(self, action, confidence):
        """Smooth predictions over time"""
        self.prediction_history.append((action, confidence))
        
        # Count action occurrences
        action_counts = {}
        confidence_sums = {}
        
        for a, c in self.prediction_history:
            action_counts[a] = action_counts.get(a, 0) + 1
            confidence_sums[a] = confidence_sums.get(a, 0) + c
        
        # Get most common action
        best_action = max(action_counts, key=action_counts.get)
        avg_confidence = confidence_sums[best_action] / action_counts[best_action]
        
        return best_action, avg_confidence
    
    def _toggle_recording(self, frame):
        """Toggle video recording"""
        if self.is_recording:
            self.is_recording = False
            if self.video_writer:
                self.video_writer.release()
                self.video_writer = None
                print("[RSK World] Recording stopped")
        else:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_path = Path(self.config.OUTPUT_DIR) / f"recording_{timestamp}.mp4"
            
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            self.video_writer = cv2.VideoWriter(
                str(output_path), fourcc, self.config.FPS,
                (self.config.FRAME_WIDTH, self.config.FRAME_HEIGHT)
            )
            self.is_recording = True
            print(f"[RSK World] Recording started: {output_path}")
    
    def _save_screenshot(self, frame):
        """Save current frame as screenshot"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = Path(self.config.OUTPUT_DIR) / f"screenshot_{timestamp}.png"
        cv2.imwrite(str(output_path), frame)
        print(f"[RSK World] Screenshot saved: {output_path}")
    
    def run(self):
        """Run real-time prediction"""
        print("=" * 60)
        print("Action Recognition - Real-Time Predictor")
        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)
        print()
        print("Controls:")
        print("  Q - Quit")
        print("  R - Start/Stop Recording")
        print("  S - Save Screenshot")
        print()
        print("Starting webcam...")
        
        # Open webcam
        cap = cv2.VideoCapture(self.config.CAMERA_ID)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.config.FRAME_WIDTH)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.config.FRAME_HEIGHT)
        cap.set(cv2.CAP_PROP_FPS, self.config.FPS)
        
        if not cap.isOpened():
            print("[ERROR] Cannot open webcam!")
            print("Make sure your webcam is connected and not in use by another application.")
            return
        
        print("[RSK World] Webcam opened successfully!")
        print("Press 'Q' to quit\n")
        
        prev_time = time.time()
        
        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    print("[ERROR] Cannot read frame from webcam")
                    break
                
                # Calculate FPS
                current_time = time.time()
                fps = 1.0 / (current_time - prev_time)
                prev_time = current_time
                self.fps_history.append(fps)
                avg_fps = np.mean(self.fps_history)
                
                # Predict action
                action, confidence = self.predictor.predict(frame)
                
                # Smooth prediction
                action, confidence = self._smooth_prediction(action, confidence)
                
                # Add overlay
                frame = self._add_overlay(frame, action, confidence, avg_fps)
                
                # Record if enabled
                if self.is_recording and self.video_writer:
                    self.video_writer.write(frame)
                
                # Display frame
                cv2.imshow("RSK World - Action Recognition", frame)
                
                # Handle key presses
                key = cv2.waitKey(1) & 0xFF
                
                if key == ord('q') or key == ord('Q'):
                    print("\n[RSK World] Exiting...")
                    break
                elif key == ord('r') or key == ord('R'):
                    self._toggle_recording(frame)
                elif key == ord('s') or key == ord('S'):
                    self._save_screenshot(frame)
        
        finally:
            # Cleanup
            cap.release()
            if self.video_writer:
                self.video_writer.release()
            cv2.destroyAllWindows()
        
        print()
        print("=" * 60)
        print("Thank you for using RSK World!")
        print("Visit: https://rskworld.in")
        print("=" * 60)


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

def main():
    """Main function"""
    predictor = RealtimePredictor()
    predictor.run()


if __name__ == "__main__":
    main()