import cv2
import numpy as np
import gradio as gr
import mediapipe as mp
from mediapipe.tasks.python import vision
from mediapipe.tasks.python import BaseOptions
from mediapipe.tasks.python.vision import PoseLandmarker, PoseLandmarkerOptions, RunningMode

MODEL_PATH = "pose_landmarker.task"

# Color palette (BGR format)
COLORS = {
    "White": (255, 255, 255),
    "Red": (0, 0, 255),
    "Green": (0, 255, 0),
    "Blue": (255, 0, 0),
    "Yellow": (0, 255, 255),
    "Cyan": (255, 255, 0),
    "Magenta": (255, 0, 255),
    "Orange": (0, 165, 255),
    "Purple": (255, 0, 128),
    "Pink": (203, 192, 255),
}

MULTICOLOR_SCHEME = {
    "face": (255, 255, 0),        # Cyan
    "torso": (0, 255, 255),       # Yellow
    "right_arm": (0, 0, 255),     # Red
    "left_arm": (255, 0, 0),      # Blue
    "right_leg": (255, 0, 255),   # Magenta
    "left_leg": (0, 255, 0),      # Green
}

def get_body_part_connections():
    """Define which connections belong to which body part"""
    
    connections = {
        "face": [
            (0, 1), (1, 2), (2, 3), (3, 7),  # Right eye region
            (0, 4), (4, 5), (5, 6), (6, 8),  # Left eye region
            (9, 10),  # Mouth
        ],
        "torso": [
            (11, 12),  # Shoulders
            (11, 23), (12, 24),  # Shoulder to hip
            (23, 24),  # Hips
        ],
        "right_arm": [
            (11, 13), (13, 15),  # Shoulder to elbow to wrist
            (15, 17), (15, 19), (15, 21),  # Wrist connections
            (17, 19),  # Hand
        ],
        "left_arm": [
            (12, 14), (14, 16),  # Shoulder to elbow to wrist
            (16, 18), (16, 20), (16, 22),  # Wrist connections
            (18, 20),  # Hand
        ],
        "right_leg": [
            (23, 25), (25, 27),  # Hip to knee to ankle
            (27, 29), (27, 31),  # Ankle connections
            (29, 31),  # Foot
        ],
        "left_leg": [
            (24, 26), (26, 28),  # Hip to knee to ankle
            (28, 30), (28, 32),  # Ankle connections
            (30, 32),  # Foot
        ],
    }
    return connections

def draw_pose(
    video_path,
    detection_confidence,
    tracking_confidence,
    background_type,
    color_mode,
    line_color,
    joint_color
):
    output_path = "output.mp4"
    
    options = PoseLandmarkerOptions(
        base_options=BaseOptions(model_asset_path=MODEL_PATH),
        running_mode=RunningMode.VIDEO,
        num_poses=1,
        min_pose_detection_confidence=detection_confidence,
        min_tracking_confidence=tracking_confidence,
    )
    landmarker = PoseLandmarker.create_from_options(options)
    
    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(3))
    height = int(cap.get(4))
    fps = cap.get(5) or 24
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    timestamp = 0
    
    line_bgr = COLORS[line_color]
    joint_bgr = COLORS[joint_color]
    

    body_parts = get_body_part_connections()
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        mp_image = mp.Image(
            image_format=mp.ImageFormat.SRGB,
            data=rgb_frame
        )
        
        result = landmarker.detect_for_video(mp_image, timestamp)
        timestamp += int(1000 / fps)
        
        if background_type == "Black Background":
            canvas = np.zeros((height, width, 3), dtype=np.uint8)
        else:  # Original Video
            canvas = frame.copy()
        
        if result.pose_landmarks:
            for pose_landmarks in result.pose_landmarks:
                points = []
                for lm in pose_landmarks:
                    x = int(lm.x * width)
                    y = int(lm.y * height)
                    points.append((x, y))
                
                if color_mode == "Single Color":
                    # Original behavior - single color for all connections
                    connections = mp.solutions.pose.POSE_CONNECTIONS
                    for c in connections:
                        cv2.line(
                            canvas,
                            points[c[0]],
                            points[c[1]],
                            line_bgr,
                            3
                        )
                else:  # Multicolor
                    for part_name, part_connections in body_parts.items():
                        part_color = MULTICOLOR_SCHEME[part_name]
                        for c in part_connections:
                            if c[0] < len(points) and c[1] < len(points):
                                cv2.line(
                                    canvas,
                                    points[c[0]],
                                    points[c[1]],
                                    part_color,
                                    3
                                )
                
                for p in points:
                    cv2.circle(canvas, p, 5, joint_bgr, -1)
        
        out.write(canvas)
    
    cap.release()
    out.release()
    
    return output_path

custom_css = """
* {
    font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif;
}

.main-header {
    text-align: center;
    padding: 30px 20px;
    border-bottom: 1px solid #404040;
    margin-bottom: 30px;
}

.main-header h1 {
    font-size: 28px;
    font-weight: 600;
    color: #ffffff;
    margin: 0 0 8px 0;
    letter-spacing: -0.5px;
}

.main-header p {
    font-size: 15px;
    color: #b0b0b0;
    margin: 0;
    font-weight: 400;
}

.section-header {
    font-size: 13px;
    font-weight: 600;
    color: #ffffff;
    text-transform: uppercase;
    letter-spacing: 0.5px;
    margin-bottom: 16px;
    padding-bottom: 8px;
    border-bottom: 2px solid #ffffff;
}

.subsection-title {
    font-size: 12px;
    font-weight: 600;
    color: #b0b0b0;
    text-transform: uppercase;
    letter-spacing: 0.3px;
    margin: 20px 0 12px 0;
}

.warning-box {
    background: #1a1a1a;
    border-left: 3px solid #888;
    padding: 16px 20px;
    margin: 20px 0;
    font-size: 14px;
    color: #d0d0d0;
    line-height: 1.6;
}

.warning-box strong {
    font-weight: 600;
    color: #ffffff;
}

.info-box {
    background: #1a1a1a;
    border: 1px solid #404040;
    padding: 16px;
    border-radius: 4px;
    margin-top: 16px;
}

.info-box p {
    margin: 6px 0;
    font-size: 13px;
    color: #b0b0b0;
    line-height: 1.5;
}

.info-box strong {
    color: #ffffff;
}

.footer {
    text-align: center;
    padding: 30px 20px;
    border-top: 1px solid #404040;
    margin-top: 40px;
    color: #888;
    font-size: 13px;
}

.footer h3 {
    font-size: 14px;
    font-weight: 600;
    color: #ffffff;
    margin-bottom: 12px;
}

.footer p {
    color: #b0b0b0;
}

#submit-btn {
    margin-top: 24px;
    background: #ffffff;
    color: #000000;
    border: none;
    font-weight: 500;
    letter-spacing: 0.3px;
}

#submit-btn:hover {
    background: #e0e0e0;
}

.gr-box {
    border-radius: 4px;
}

.gr-input, .gr-dropdown, .gr-radio {
    border-radius: 4px;
}

.gr-accordion {
    border: 1px solid #404040;
    border-radius: 4px;
}

label {
    color: #d0d0d0 !important;
}

.gr-text-input, .gr-dropdown {
    background: #1a1a1a;
    border: 1px solid #404040;
    color: #ffffff;
}

.color-legend {
    background: #1a1a1a;
    border: 1px solid #404040;
    padding: 12px;
    border-radius: 4px;
    margin-top: 12px;
    font-size: 12px;
}

.color-legend p {
    margin: 4px 0;
    color: #b0b0b0;
}

.color-item {
    display: inline-block;
    width: 12px;
    height: 12px;
    margin-right: 6px;
    border-radius: 2px;
}
"""

with gr.Blocks(title="MediaPipe Pose Estimation", theme=gr.themes.Default(), css=custom_css) as demo:
    
    # Header
    gr.HTML(
        """
        <div class="main-header">
            <h1>MediaPipe Pose Estimation</h1>
            <p>Advanced pose detection and visualization system</p>
        </div>
        """
    )
    
    with gr.Accordion("Instructions", open=False):
        gr.Markdown(
            """
            1. Upload a video file containing human subjects
            2. Adjust detection and tracking confidence parameters as needed
            3. Select the desired background type for the output
            4. Choose between single color or multicolor visualization
            5. If single color is selected, customize the skeleton colors
            6. Click Process Video to generate the result
            """
        )
    
    with gr.Accordion("Parameter Guidelines", open=False):
        gr.Markdown(
            """
            **Detection Confidence**
            
            Controls the minimum confidence threshold for initial pose detection. Higher values (0.7-0.9) reduce false positives but may miss some poses. Lower values (0.3-0.5) detect more poses but may include incorrect detections. Default: 0.5
            
            **Tracking Confidence**
            
            Determines the reliability threshold for tracking poses across consecutive frames. Higher values provide more stable tracking but may lose poses more easily. Lower values maintain tracking longer but may be less stable. Default: 0.5
            
            **Color Modes**
            
            Single Color: All skeleton lines use the same color. Multicolor: Different body parts are colored differently (face, torso, arms, legs).
            """
        )
    
    # Warning message
    gr.HTML(
        """
        <div class="warning-box">
            <strong>Processing Time Notice:</strong> Processing duration is proportional to video length and resolution. Videos exceeding 2 minutes or high-resolution files may require several minutes to process. Please wait while the system completes the analysis.
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            gr.HTML('<div class="section-header">Input</div>')
            
            video_input = gr.Video(label="Video File")
            
            gr.HTML('<div class="subsection-title">Confidence Parameters</div>')
            detection_conf = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.5,
                step=0.05,
                label="Detection Confidence"
            )
            tracking_conf = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.5,
                step=0.05,
                label="Tracking Confidence"
            )
            
            gr.HTML('<div class="subsection-title">Visualization Options</div>')
            background_type = gr.Radio(
                choices=["Black Background", "Original Video"],
                value="Black Background",
                label="Background Type"
            )
            
            color_mode = gr.Radio(
                choices=["Single Color", "Multicolor"],
                value="Single Color",
                label="Color Mode"
            )
            
            with gr.Row():
                line_color = gr.Dropdown(
                    choices=list(COLORS.keys()),
                    value="White",
                    label="Line Color",
                    info="Only applies in Single Color mode"
                )
                joint_color = gr.Dropdown(
                    choices=list(COLORS.keys()),
                    value="Red",
                    label="Joint Color",
                    info='Single & Multicolor'
                )
            
            gr.HTML(
                """
                <div class="color-legend">
                    <p><strong>Multicolor Legend:</strong></p>
                    <p><span class="color-item" style="background: rgb(255, 255, 0);"></span>Torso: Yellow</p>
                    <p><span class="color-item" style="background: rgb(255, 0, 0);"></span>Right Arm: Red</p>
                    <p><span class="color-item" style="background: rgb(0, 0, 255);"></span>Left Arm: Blue</p>
                    <p><span class="color-item" style="background: rgb(255, 0, 255);"></span>Right Leg: Magenta</p>
                    <p><span class="color-item" style="background: rgb(0, 255, 0);"></span>Left Leg: Green</p>
                </div>
                """
            )
            
            submit_btn = gr.Button("Process Video", variant="primary", elem_id="submit-btn")
        
        with gr.Column():
            gr.HTML('<div class="section-header">Output</div>')
            video_output = gr.Video(label="Processed Video")
            
            gr.HTML(
                """
                <div class="info-box">
                    <p><strong>Output Specifications:</strong></p>
                    <p>Format: MP4 (H.264 encoding)</p>
                    <p>Resolution: Matches input resolution</p>
                    <p>Frame Rate: Matches input frame rate</p>
                    <p>Keypoints: 33 body landmarks tracked per frame</p>
                </div>
                """
            )
    
    # Footer
    gr.HTML(
        """
        <div class="footer">
            <h3>Technical Information</h3>
            <p>This application utilizes MediaPipe Pose Landmarker for real-time pose detection and tracking.</p>
            <p>The system identifies 33 anatomical keypoints and visualizes skeletal structure with customizable styling.</p>
            <p style="margin-top: 16px;">Supported formats: MP4, AVI, MOV, WebM</p>
            <p style="margin-top: 16px; color: #666;">Powered by Google MediaPipe</p>
        </div>
        """
    )
    
    submit_btn.click(
        fn=draw_pose,
        inputs=[
            video_input,
            detection_conf,
            tracking_conf,
            background_type,
            color_mode,
            line_color,
            joint_color
        ],
        outputs=video_output
    )

demo.launch(server_name="0.0.0.0", server_port=7860)