"""
==================================================================================
    Action Recognition Dataset - Sample Video Generator
==================================================================================
    Project: Action Recognition Dataset
    Category: Video Data / Data Science
    
    This script generates sample action videos with stick figure animations
    and RSK World logo watermark for demonstration purposes.
    
==================================================================================
    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.
    
==================================================================================

Usage:
    python generate_samples.py
    
This will generate sample videos for each action class in the sample_data directory.
"""

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


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

# Video settings
VIDEO_WIDTH = 640
VIDEO_HEIGHT = 480
FPS = 30
DURATION = 3  # seconds per video
NUM_FRAMES = FPS * DURATION

# Colors (BGR format for OpenCV)
BG_COLOR = (23, 20, 17)  # Dark background #111417
STICK_COLOR = (255, 212, 0)  # Cyan #00D4FF
LABEL_BG = (50, 40, 30)  # Semi-dark for label background
TEXT_COLOR = (255, 255, 255)  # White
LOGO_COLOR = (255, 212, 0)  # Cyan for logo

# Action classes to generate
ACTION_CLASSES = [
    'walking',
    'running', 
    'jumping',
    'waving',
    'sitting',
    'standing',
    'clapping',
    'kicking'
]

# Number of sample videos per class
SAMPLES_PER_CLASS = 3


# ==================================================================================
# RSK World Logo Drawing
# ==================================================================================

def draw_logo(frame, x, y, scale=1.0):
    """
    Draw RSK World logo on frame
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    """
    # Draw video icon (two rectangles)
    icon_size = int(20 * scale)
    cv2.rectangle(frame, 
                  (x, y), 
                  (x + icon_size, y + int(icon_size * 0.7)),
                  LOGO_COLOR, -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, -1)
    
    # Draw text "RSKWorld"
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.6 * scale
    text = "RSKWorld"
    text_x = x + int(icon_size * 1.5)
    cv2.putText(frame, text, (text_x, y + int(icon_size * 0.6)), 
                font, font_scale, TEXT_COLOR, 1, cv2.LINE_AA)
    
    # Draw ".in" in cyan
    text_size = cv2.getTextSize(text, font, font_scale, 1)[0]
    cv2.putText(frame, ".in", (text_x + text_size[0], y + int(icon_size * 0.6)), 
                font, font_scale * 0.8, LOGO_COLOR, 1, cv2.LINE_AA)


def draw_watermark(frame):
    """Draw watermark with copyright info"""
    h, w = frame.shape[:2]
    
    # Draw logo in bottom right
    draw_logo(frame, w - 150, h - 30, scale=0.8)
    
    # Draw copyright text
    font = cv2.FONT_HERSHEY_SIMPLEX
    text = "rskworld.in"
    cv2.putText(frame, text, (10, h - 10), font, 0.4, (100, 100, 100), 1, cv2.LINE_AA)


# ==================================================================================
# Stick Figure Drawing Functions
# ==================================================================================

class StickFigure:
    """
    Animated stick figure for action recognition samples
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Contact: help@rskworld.in
    """
    
    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_angle=0, right_arm_angle=0, 
             left_leg_angle=0, right_leg_angle=0, body_offset_y=0):
        """Draw the stick figure with specified joint angles"""
        
        color = STICK_COLOR
        thickness = max(2, int(3 * self.scale))
        
        # Head position
        head_x = self.x
        head_y = self.y + body_offset_y
        
        # Draw head
        cv2.circle(frame, (head_x, head_y), self.head_radius, color, thickness)
        
        # Body start and end
        body_start = (head_x, head_y + self.head_radius)
        body_end = (head_x, head_y + self.head_radius + self.body_length)
        
        # Draw body
        cv2.line(frame, body_start, body_end, color, thickness)
        
        # Shoulder position (1/4 down the body)
        shoulder_y = head_y + self.head_radius + int(self.body_length * 0.2)
        
        # Draw left arm
        left_arm_end_x = head_x - int(self.limb_length * math.cos(math.radians(left_arm_angle + 90)))
        left_arm_end_y = shoulder_y + int(self.limb_length * math.sin(math.radians(left_arm_angle + 90)))
        cv2.line(frame, (head_x, shoulder_y), (left_arm_end_x, left_arm_end_y), color, thickness)
        
        # Draw right arm
        right_arm_end_x = head_x + int(self.limb_length * math.cos(math.radians(right_arm_angle + 90)))
        right_arm_end_y = shoulder_y + int(self.limb_length * math.sin(math.radians(right_arm_angle + 90)))
        cv2.line(frame, (head_x, shoulder_y), (right_arm_end_x, right_arm_end_y), color, thickness)
        
        # Hip position (end of body)
        hip_y = body_end[1]
        
        # Draw left leg
        left_leg_end_x = head_x - int(self.limb_length * math.sin(math.radians(left_leg_angle)))
        left_leg_end_y = hip_y + int(self.limb_length * math.cos(math.radians(left_leg_angle)))
        cv2.line(frame, (head_x, hip_y), (left_leg_end_x, left_leg_end_y), color, thickness)
        
        # Draw right leg
        right_leg_end_x = head_x + int(self.limb_length * math.sin(math.radians(right_leg_angle)))
        right_leg_end_y = hip_y + int(self.limb_length * math.cos(math.radians(right_leg_angle)))
        cv2.line(frame, (head_x, hip_y), (right_leg_end_x, right_leg_end_y), color, thickness)


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

def animate_walking(figure, frame_num, total_frames):
    """Walking animation - arms and legs swing"""
    cycle = (frame_num % 30) / 30.0 * 2 * math.pi
    
    left_arm = 30 * math.sin(cycle)
    right_arm = -30 * math.sin(cycle)
    left_leg = 25 * math.sin(cycle)
    right_leg = -25 * math.sin(cycle)
    
    return left_arm, right_arm, left_leg, right_leg, 0


def animate_running(figure, frame_num, total_frames):
    """Running animation - faster arm and leg swing"""
    cycle = (frame_num % 15) / 15.0 * 2 * math.pi
    
    left_arm = 45 * math.sin(cycle)
    right_arm = -45 * math.sin(cycle)
    left_leg = 40 * math.sin(cycle)
    right_leg = -40 * math.sin(cycle)
    
    return left_arm, right_arm, left_leg, right_leg, 0


def animate_jumping(figure, frame_num, total_frames):
    """Jumping animation - body moves up and down"""
    cycle = (frame_num % 45) / 45.0
    
    # Jump arc
    if cycle < 0.3:
        # Crouch
        body_offset = int(20 * cycle / 0.3)
        arm_angle = -30 + 60 * cycle / 0.3
    elif cycle < 0.6:
        # Jump up
        progress = (cycle - 0.3) / 0.3
        body_offset = int(-60 * math.sin(progress * math.pi))
        arm_angle = 30 - 120 * progress
    else:
        # Land
        progress = (cycle - 0.6) / 0.4
        body_offset = int(-60 * math.sin((1 - progress) * math.pi / 2))
        arm_angle = -90 + 60 * progress
    
    return arm_angle, arm_angle, 10, -10, body_offset


def animate_waving(figure, frame_num, total_frames):
    """Waving animation - right arm waves"""
    cycle = (frame_num % 20) / 20.0 * 2 * math.pi
    
    right_arm = -120 + 30 * math.sin(cycle)
    
    return 0, right_arm, 5, -5, 0


def animate_sitting(figure, frame_num, total_frames):
    """Sitting animation - transition to seated"""
    progress = min(1.0, frame_num / (total_frames * 0.4))
    
    # Bend legs
    leg_angle = 90 * progress
    body_offset = int(40 * progress)
    
    return 20, 20, leg_angle, -leg_angle, body_offset


def animate_standing(figure, frame_num, total_frames):
    """Standing animation - stand up from seated"""
    progress = min(1.0, frame_num / (total_frames * 0.4))
    
    # Straighten legs
    leg_angle = 90 * (1 - progress)
    body_offset = int(40 * (1 - progress))
    
    return 20 * (1 - progress), 20 * (1 - progress), leg_angle, -leg_angle, body_offset


def animate_clapping(figure, frame_num, total_frames):
    """Clapping animation - hands come together"""
    cycle = (frame_num % 15) / 15.0
    
    if cycle < 0.5:
        # Hands apart
        arm_spread = 60 * (1 - cycle * 2)
    else:
        # Hands together
        arm_spread = 60 * (cycle - 0.5) * 2
    
    left_arm = -60 + arm_spread
    right_arm = -60 + arm_spread
    
    return left_arm, -right_arm, 5, -5, 0


def animate_kicking(figure, frame_num, total_frames):
    """Kicking animation - leg kicks forward"""
    cycle = (frame_num % 30) / 30.0
    
    if cycle < 0.4:
        # Wind up
        right_leg = -30 * cycle / 0.4
    elif cycle < 0.6:
        # Kick
        progress = (cycle - 0.4) / 0.2
        right_leg = -30 + 120 * progress
    else:
        # Return
        progress = (cycle - 0.6) / 0.4
        right_leg = 90 * (1 - progress)
    
    return 20, -20, 0, right_leg, 0


# Animation function mapping
ANIMATION_FUNCTIONS = {
    'walking': animate_walking,
    'running': animate_running,
    'jumping': animate_jumping,
    'waving': animate_waving,
    'sitting': animate_sitting,
    'standing': animate_standing,
    'clapping': animate_clapping,
    'kicking': animate_kicking
}


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

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


def draw_grid_background(frame):
    """Draw subtle grid pattern on background"""
    h, w = frame.shape[:2]
    grid_color = (30, 28, 25)  # Slightly lighter than background
    
    # Vertical lines
    for x in range(0, w, 40):
        cv2.line(frame, (x, 0), (x, h), grid_color, 1)
    
    # Horizontal lines
    for y in range(0, h, 40):
        cv2.line(frame, (0, y), (w, y), grid_color, 1)


def generate_video(action_name, output_path, sample_num=1):
    """
    Generate a sample video for the given action
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    """
    print(f"  Generating {action_name} sample {sample_num}...")
    
    # Create video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(str(output_path), fourcc, FPS, (VIDEO_WIDTH, VIDEO_HEIGHT))
    
    # Create stick figure
    figure = StickFigure(VIDEO_WIDTH // 2, VIDEO_HEIGHT // 2 - 60, scale=1.5)
    
    # Get animation function
    animate_func = ANIMATION_FUNCTIONS.get(action_name, animate_walking)
    
    # Add variation based on sample number
    np.random.seed(sample_num * 42)
    x_offset = np.random.randint(-50, 50)
    figure.x += x_offset
    
    # Generate frames
    for frame_num in range(NUM_FRAMES):
        # Create frame with background
        frame = np.full((VIDEO_HEIGHT, VIDEO_WIDTH, 3), BG_COLOR, dtype=np.uint8)
        
        # Draw grid background
        draw_grid_background(frame)
        
        # Get animation parameters
        left_arm, right_arm, left_leg, right_leg, body_offset = animate_func(
            figure, frame_num, NUM_FRAMES
        )
        
        # Draw stick figure
        figure.draw(frame, left_arm, right_arm, left_leg, right_leg, body_offset)
        
        # Draw action label with random confidence
        confidence = 95.0 + np.random.uniform(-3, 3)
        draw_action_label(frame, action_name, confidence)
        
        # Draw watermark
        draw_watermark(frame)
        
        # Write frame
        out.write(frame)
    
    out.release()
    print(f"    [OK] Saved: {output_path}")


def create_directory_structure():
    """Create the sample data directory structure"""
    base_dir = Path('sample_data')
    
    for split in ['train', 'val', 'test']:
        for action in ACTION_CLASSES:
            dir_path = base_dir / split / action
            dir_path.mkdir(parents=True, exist_ok=True)
            print(f"  Created: {dir_path}")


def generate_all_samples():
    """
    Generate all sample videos
    
    RSK World (https://rskworld.in)
    Founder: Molla Samser
    Designer: Rima Khatun
    Contact: help@rskworld.in | +91 93305 39277
    """
    print("=" * 60)
    print("Action Recognition Dataset - Sample Video Generator")
    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()
    
    # Create directory structure
    print("Creating directory structure...")
    create_directory_structure()
    print()
    
    # Generate samples for each split
    splits = {
        'train': 3,  # 3 samples per class for training
        'val': 1,    # 1 sample per class for validation
        'test': 1    # 1 sample per class for testing
    }
    
    total_videos = 0
    
    for split, num_samples in splits.items():
        print(f"\nGenerating {split} samples...")
        print("-" * 40)
        
        for action in ACTION_CLASSES:
            print(f"\n[{action.upper()}]")
            
            for i in range(1, num_samples + 1):
                output_path = Path('sample_data') / split / action / f"{action}_{i:03d}.mp4"
                generate_video(action, output_path, i)
                total_videos += 1
    
    print()
    print("=" * 60)
    print(f"Generation complete! Total videos: {total_videos}")
    print("=" * 60)
    print()
    print("Dataset structure:")
    print("  sample_data/")
    print("  ├── train/     (24 videos)")
    print("  ├── val/       (8 videos)")
    print("  └── test/      (8 videos)")
    print()
    print("Thank you for using RSK World datasets!")
    print("Visit: https://rskworld.in")


# ==================================================================================
# Generate Preview Image
# ==================================================================================

def generate_preview_image():
    """Generate a preview image for the dataset"""
    print("\nGenerating preview image...")
    
    # Create image
    img_width = 1200
    img_height = 630
    img = np.full((img_height, img_width, 3), BG_COLOR, dtype=np.uint8)
    
    # Draw gradient background
    for y in range(img_height):
        alpha = y / img_height
        color = tuple(int(c * (1 - alpha * 0.3)) for c in BG_COLOR)
        cv2.line(img, (0, y), (img_width, y), color, 1)
    
    # Draw grid
    draw_grid_background(img)
    
    # Draw title
    font = cv2.FONT_HERSHEY_SIMPLEX
    title = "Action Recognition Dataset"
    title_size = cv2.getTextSize(title, font, 2.0, 3)[0]
    title_x = (img_width - title_size[0]) // 2
    cv2.putText(img, title, (title_x, 100), font, 2.0, TEXT_COLOR, 3, cv2.LINE_AA)
    
    # Draw subtitle
    subtitle = "Video Dataset for AI/ML Training"
    sub_size = cv2.getTextSize(subtitle, font, 1.0, 2)[0]
    sub_x = (img_width - sub_size[0]) // 2
    cv2.putText(img, subtitle, (sub_x, 150), font, 1.0, LOGO_COLOR, 2, cv2.LINE_AA)
    
    # Draw multiple stick figures with different actions
    actions = ['walking', 'running', 'jumping', 'waving']
    figure_positions = [
        (200, 300), (450, 300), (750, 300), (1000, 300)
    ]
    
    for i, (action, (x, y)) in enumerate(zip(actions, figure_positions)):
        figure = StickFigure(x, y, scale=1.2)
        animate_func = ANIMATION_FUNCTIONS[action]
        left_arm, right_arm, left_leg, right_leg, body_offset = animate_func(figure, 15, 30)
        figure.draw(img, left_arm, right_arm, left_leg, right_leg, body_offset)
        
        # Draw action label
        label_size = cv2.getTextSize(action.upper(), font, 0.6, 1)[0]
        label_x = x - label_size[0] // 2
        cv2.putText(img, action.upper(), (label_x, y + 120), font, 0.6, LOGO_COLOR, 1, cv2.LINE_AA)
    
    # Draw stats
    stats = ["1,500+ Videos", "12 Action Classes", "50+ Hours"]
    stat_y = 520
    for i, stat in enumerate(stats):
        stat_x = 200 + i * 350
        cv2.putText(img, stat, (stat_x, stat_y), font, 0.8, TEXT_COLOR, 2, cv2.LINE_AA)
    
    # Draw logo
    draw_logo(img, img_width - 200, img_height - 50, scale=1.2)
    
    # Draw copyright
    copyright_text = "© 2026 RSK World | rskworld.in"
    cv2.putText(img, copyright_text, (20, img_height - 20), font, 0.5, (100, 100, 100), 1, cv2.LINE_AA)
    
    # Save image
    cv2.imwrite('action-recognition.png', img)
    print("  [OK] Saved: action-recognition.png")


# ==================================================================================
# Main Entry Point
# ==================================================================================

if __name__ == '__main__':
    try:
        # Generate preview image
        generate_preview_image()
        
        # Generate sample videos
        generate_all_samples()
        
    except Exception as e:
        print(f"\nError: {e}")
        print("\nMake sure you have OpenCV installed:")
        print("  pip install opencv-python numpy")
        print("\nContact: help@rskworld.in for support")