"""
==================================================================================
    Action Recognition Dataset - Browser-Compatible Video Generator
==================================================================================
    Project: Action Recognition Dataset
    Category: Video Data / Data Science
    
    This script generates sample action videos with stick figure animations
    in browser-compatible H.264 format.
    
==================================================================================
    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
==================================================================================
    (C) 2026 RSK World. All Rights Reserved.
    
==================================================================================

Usage:
    python generate_browser_videos.py
"""

import cv2
import numpy as np
import os
import math
from pathlib import Path
import imageio

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

VIDEO_WIDTH = 640
VIDEO_HEIGHT = 480
FPS = 30
DURATION = 3
NUM_FRAMES = FPS * DURATION

# Colors (BGR for OpenCV, RGB for imageio)
BG_COLOR_BGR = (23, 20, 17)
BG_COLOR_RGB = (17, 20, 23)
STICK_COLOR_BGR = (255, 212, 0)
STICK_COLOR_RGB = (0, 212, 255)  # Cyan
LABEL_BG_BGR = (50, 40, 30)
TEXT_COLOR = (255, 255, 255)
LOGO_COLOR_BGR = (255, 212, 0)

# Action classes
ACTION_CLASSES = ['dancing', 'running', 'jumping', 'walking', 'stretching', 'yoga', 'exercising']


# ==================================================================================
# Drawing Functions
# ==================================================================================

def draw_grid_background(frame):
    """Draw subtle grid pattern"""
    h, w = frame.shape[:2]
    grid_color = (30, 28, 25)
    for x in range(0, w, 40):
        cv2.line(frame, (x, 0), (x, h), grid_color, 1)
    for y in range(0, h, 40):
        cv2.line(frame, (0, y), (w, y), grid_color, 1)


def draw_logo(frame, x, y, scale=1.0):
    """Draw RSK World logo"""
    icon_size = int(20 * scale)
    cv2.rectangle(frame, (x, y), (x + icon_size, y + int(icon_size * 0.7)), LOGO_COLOR_BGR, -1)
    cv2.rectangle(frame, (x + int(icon_size * 0.3), y + int(icon_size * 0.15)),
                  (x + int(icon_size * 1.3), y + int(icon_size * 0.85)), LOGO_COLOR_BGR, -1)
    
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame, "RSKWorld", (x + int(icon_size * 1.5), y + int(icon_size * 0.6)),
                font, 0.6 * scale, TEXT_COLOR, 1, cv2.LINE_AA)
    text_w = cv2.getTextSize("RSKWorld", font, 0.6 * scale, 1)[0][0]
    cv2.putText(frame, ".in", (x + int(icon_size * 1.5) + text_w, y + int(icon_size * 0.6)),
                font, 0.5 * scale, LOGO_COLOR_BGR, 1, cv2.LINE_AA)


def draw_watermark(frame):
    """Draw watermark"""
    h, w = frame.shape[:2]
    draw_logo(frame, w - 150, h - 30, scale=0.8)
    cv2.putText(frame, "rskworld.in", (10, h - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (100, 100, 100), 1, cv2.LINE_AA)


def draw_action_label(frame, action_name, confidence=None):
    """Draw action label"""
    h, w = frame.shape[:2]
    label_h = 60
    overlay = frame.copy()
    cv2.rectangle(overlay, (0, h - label_h), (w, h), LABEL_BG_BGR, -1)
    cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame)
    
    font = cv2.FONT_HERSHEY_SIMPLEX
    text = action_name.upper()
    text_size = cv2.getTextSize(text, font, 1.0, 2)[0]
    cv2.putText(frame, text, ((w - text_size[0]) // 2, h - 30), font, 1.0, LOGO_COLOR_BGR, 2, cv2.LINE_AA)
    
    if confidence:
        conf_text = f"Confidence: {confidence:.1f}%"
        conf_size = cv2.getTextSize(conf_text, font, 0.5, 1)[0]
        cv2.putText(frame, conf_text, ((w - conf_size[0]) // 2, h - 10), font, 0.5, TEXT_COLOR, 1, cv2.LINE_AA)


# ==================================================================================
# Stick Figure Class
# ==================================================================================

class StickFigure:
    def __init__(self, x, y, scale=1.0):
        self.x = x
        self.y = y
        self.scale = scale
        self.head_radius = int(20 * scale)
        self.body_length = int(60 * scale)
        self.limb_length = int(40 * scale)
    
    def draw(self, frame, left_arm=0, right_arm=0, left_leg=0, right_leg=0, body_offset=0):
        color = STICK_COLOR_BGR
        thickness = max(2, int(3 * self.scale))
        
        head_x, head_y = self.x, self.y + body_offset
        cv2.circle(frame, (head_x, head_y), self.head_radius, color, thickness)
        
        body_start = (head_x, head_y + self.head_radius)
        body_end = (head_x, head_y + self.head_radius + self.body_length)
        cv2.line(frame, body_start, body_end, color, thickness)
        
        shoulder_y = head_y + self.head_radius + int(self.body_length * 0.2)
        
        # Arms
        for arm_angle, direction in [(left_arm, -1), (right_arm, 1)]:
            end_x = head_x + direction * int(self.limb_length * math.cos(math.radians(arm_angle + 90)))
            end_y = shoulder_y + int(self.limb_length * math.sin(math.radians(arm_angle + 90)))
            cv2.line(frame, (head_x, shoulder_y), (end_x, end_y), color, thickness)
        
        # Legs
        hip_y = body_end[1]
        for leg_angle, direction in [(left_leg, -1), (right_leg, 1)]:
            end_x = head_x + direction * int(self.limb_length * math.sin(math.radians(leg_angle)))
            end_y = hip_y + int(self.limb_length * math.cos(math.radians(leg_angle)))
            cv2.line(frame, (head_x, hip_y), (end_x, end_y), color, thickness)


# ==================================================================================
# Animation Functions
# ==================================================================================

def animate_dancing(frame_num):
    cycle = (frame_num % 20) / 20.0 * 2 * math.pi
    return (-60 + 30 * math.sin(cycle), -60 - 30 * math.sin(cycle),
            20 * math.sin(cycle), -20 * math.sin(cycle), int(5 * math.sin(cycle * 2)))

def animate_running(frame_num):
    cycle = (frame_num % 15) / 15.0 * 2 * math.pi
    return (45 * math.sin(cycle), -45 * math.sin(cycle),
            40 * math.sin(cycle), -40 * math.sin(cycle), 0)

def animate_jumping(frame_num):
    cycle = (frame_num % 45) / 45.0
    if cycle < 0.3:
        return (-30, -30, 10, -10, int(20 * cycle / 0.3))
    elif cycle < 0.6:
        progress = (cycle - 0.3) / 0.3
        return (-90, -90, 20, -20, int(-60 * math.sin(progress * math.pi)))
    else:
        progress = (cycle - 0.6) / 0.4
        return (-30, -30, 10, -10, int(-30 * (1 - progress)))

def animate_walking(frame_num):
    cycle = (frame_num % 30) / 30.0 * 2 * math.pi
    return (30 * math.sin(cycle), -30 * math.sin(cycle),
            25 * math.sin(cycle), -25 * math.sin(cycle), 0)

def animate_stretching(frame_num):
    cycle = (frame_num % 60) / 60.0 * 2 * math.pi
    return (-150 + 30 * math.sin(cycle), -150 - 30 * math.sin(cycle),
            0, 0, int(10 * math.sin(cycle)))

def animate_yoga(frame_num):
    cycle = (frame_num % 90) / 90.0
    if cycle < 0.5:
        angle = -90 * (cycle * 2)
        return (angle, angle, 0, 0, 0)
    else:
        angle = -90
        leg = 90 * ((cycle - 0.5) * 2)
        return (angle, angle, 0, leg, 0)

def animate_exercising(frame_num):
    cycle = (frame_num % 30) / 30.0 * 2 * math.pi
    squat = int(30 * (1 - math.cos(cycle)) / 2)
    return (-90 + 20 * math.sin(cycle), -90 - 20 * math.sin(cycle),
            30 * (1 - math.cos(cycle)) / 2, -30 * (1 - math.cos(cycle)) / 2, squat)

ANIMATIONS = {
    'dancing': animate_dancing,
    'running': animate_running,
    'jumping': animate_jumping,
    'walking': animate_walking,
    'stretching': animate_stretching,
    'yoga': animate_yoga,
    'exercising': animate_exercising
}


# ==================================================================================
# Video Generation
# ==================================================================================

def generate_video(action_name, output_path):
    """Generate browser-compatible video using imageio"""
    print(f"  Generating {action_name}...")
    
    figure = StickFigure(VIDEO_WIDTH // 2, VIDEO_HEIGHT // 2 - 60, scale=1.5)
    animate_func = ANIMATIONS.get(action_name, animate_walking)
    
    frames = []
    
    for frame_num in range(NUM_FRAMES):
        # Create frame
        frame = np.full((VIDEO_HEIGHT, VIDEO_WIDTH, 3), BG_COLOR_BGR, dtype=np.uint8)
        draw_grid_background(frame)
        
        # Animate
        left_arm, right_arm, left_leg, right_leg, body_offset = animate_func(frame_num)
        figure.draw(frame, left_arm, right_arm, left_leg, right_leg, body_offset)
        
        # Labels
        confidence = 95.0 + np.random.uniform(-3, 3)
        draw_action_label(frame, action_name, confidence)
        draw_watermark(frame)
        
        # Convert BGR to RGB for imageio
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame_rgb)
    
    # Write video with H.264 codec (browser compatible)
    imageio.mimwrite(
        str(output_path),
        frames,
        fps=FPS,
        codec='libx264',
        quality=8,
        output_params=['-pix_fmt', 'yuv420p', '-movflags', '+faststart']
    )
    
    print(f"    [OK] Saved: {output_path}")


def main():
    """Generate all sample videos"""
    print("=" * 60)
    print("Action Recognition - Browser-Compatible Video Generator")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("=" * 60)
    print()
    
    # Create directories
    base_dir = Path('sample_data')
    for split in ['train']:
        for action in ACTION_CLASSES:
            (base_dir / split / action).mkdir(parents=True, exist_ok=True)
    
    # Generate videos
    for action in ACTION_CLASSES:
        output_path = base_dir / 'train' / action / f'{action}_001.mp4'
        generate_video(action, output_path)
    
    print()
    print("=" * 60)
    print("Generation complete! Videos are now browser-compatible.")
    print("=" * 60)


if __name__ == '__main__':
    main()