import cv2
import gradio as gr
from ultralytics import YOLO
import numpy as np
import math
import tempfile
import os

# Load the model
# Load the model
# Using custom model provided by user
model = YOLO('yolo26x-pose.pt')

def get_angle(p1, p2, p3):
    """
    Calculate the angle between three points (p1-p2-p3) at joint p2.
    """
    rad = math.atan2(p3[1]-p2[1], p3[0]-p2[0]) - math.atan2(p1[1]-p2[1], p1[0]-p2[0])
    deg = abs(rad * 180.0 / math.pi)
    return 360 - deg if deg > 180 else deg

def draw_pose(img, kps):
    """
    Draw pose landmarks and classify posture based on knee angle.
    """
    legs = [(11, 13, 15), (12, 14, 16)] # hip-knee-ankle indices
    status = "Unknown"

    for h_idx, k_idx, a_idx in legs:
        hip, knee, ankle = kps[h_idx], kps[k_idx], kps[a_idx]
        
        # Check confidence scores (index 2)
        if hip[2] > 0.5 and knee[2] > 0.5 and ankle[2] > 0.5:
            ang = get_angle(hip[:2], knee[:2], ankle[:2])
            
            if ang > 160: 
                posture, color = "STANDING", (0, 255, 0)
            elif ang < 140: 
                posture, color = "SITTING", (255, 0, 0)
            else: 
                posture, color = "BENDING", (0, 165, 255)
            
            # Draw lines and text
            pt = lambda p: (int(p[0]), int(p[1]))
            cv2.line(img, pt(hip), pt(knee), (255,0,255), 2)
            cv2.line(img, pt(knee), pt(ankle), (255,0,255), 2)
            
            label = f"{posture} {int(ang)}"
            (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)
            cv2.rectangle(img, (int(hip[0]), int(hip[1]) - 30), (int(hip[0]) + w, int(hip[1])), color, -1)
            cv2.putText(img, label, (int(hip[0]), int(hip[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
            
            return posture
            
    return status

def process_img(img):
    """
    Process a single image for posture detection.
    """
    if img is None: 
        return None, "No Image"
    
    # Run inference
    res = model(img)
    out = img.copy()
    msg = "No Person Detected"
    
    if res and res[0].keypoints is not None:
        for k in res[0].keypoints.data.cpu().numpy():
             msg = f"POSTURE: {draw_pose(out, k)}"
             
    return out, msg

def process_vid(vid_path):
    """
    Process a video file for posture detection.
    """
    if not vid_path: 
        return None
        
    cap = cv2.VideoCapture(vid_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # Create temp file for output
    fd, out_path = tempfile.mkstemp(suffix='.mp4')
    os.close(fd)
    
    # Initialize video writer
    writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret: 
            break
        
        # Run inference
        res = model(frame, verbose=False)
        if res and res[0].keypoints is not None:
            for k in res[0].keypoints.data.cpu().numpy():
                draw_pose(frame, k)
                
        writer.write(frame)
        
    cap.release()
    writer.release()
    return out_path

# Define Gradio Interface
with gr.Blocks(title="Postures") as app:
    gr.Markdown("# Posture Detection\nSimple angle-based classification using YOLOv8: Standing (>160), Sitting (<140)")
    
    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column():
                img_inp = gr.Image(type="numpy", label="Input Image")
                img_btn = gr.Button("Detect Posture")
            with gr.Column():
                img_out = gr.Image(label="Result")
                img_stat = gr.Textbox(label="Status")
        img_btn.click(process_img, inputs=img_inp, outputs=[img_out, img_stat])
        
    with gr.Tab("Video"):
        with gr.Row():
            with gr.Column():
                vid_inp = gr.Video(label="Input Video")
                vid_btn = gr.Button("Process Video")
            with gr.Column():
                vid_out = gr.Video(label="Processed Video")
        vid_btn.click(process_vid, inputs=vid_inp, outputs=vid_out)

if __name__ == "__main__":
    app.launch()