import gradio as gr
import cv2
import tempfile
import numpy as np
import urllib.request
import os
import subprocess
import shutil
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

model_url="https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/1/pose_landmarker_lite.task"
model_path="/tmp/pose_landmarker_lite.task"
if not os.path.exists(model_path):
    urllib.request.urlretrieve(model_url,model_path)

class PoseDetector:
    def __init__(self,confidence=0.25):
        self.confidence=confidence
        self.prev_landmarks=None
        self.setup_detector(confidence)
    def setup_detector(self,confidence):
        base_options=python.BaseOptions(model_asset_path=model_path)
        options=vision.PoseLandmarkerOptions(
            base_options=base_options,
            running_mode=vision.RunningMode.VIDEO,
            num_poses=1,
            min_pose_detection_confidence=confidence,
            min_pose_presence_confidence=confidence,
            min_tracking_confidence=confidence,
            output_segmentation_masks=True)
        self.pose_landmarker=vision.PoseLandmarker.create_from_options(options)
    def enhance_image(self,frame):
        lab=cv2.cvtColor(frame,cv2.COLOR_BGR2LAB)
        l,a,b=cv2.split(lab)
        clahe=cv2.createCLAHE(clipLimit=2.0,tileGridSize=(8,8))
        l=clahe.apply(l)
        enhanced=cv2.merge([l,a,b])
        enhanced=cv2.cvtColor(enhanced,cv2.COLOR_LAB2BGR)
        kernel=np.array([[-0.5,-0.5,-0.5],[-0.5,5,-0.5],[-0.5,-0.5,-0.5]])
        enhanced=cv2.filter2D(enhanced,-1,kernel)
        enhanced=cv2.convertScaleAbs(enhanced,alpha=1.1,beta=10)
        return enhanced
    def remove_background(self,frame,segmentation_mask):
        if segmentation_mask is None:
            return frame
        try:
            mask=segmentation_mask.numpy_view()
            if mask.shape[:2]!=(frame.shape[0],frame.shape[1]):
                mask=cv2.resize(mask,(frame.shape[1],frame.shape[0]))
            binary_mask=(mask>0.5).astype(np.uint8)*255
            kernel=np.ones((5,5),np.uint8)
            binary_mask=cv2.morphologyEx(binary_mask,cv2.MORPH_CLOSE,kernel)
            binary_mask=cv2.morphologyEx(binary_mask,cv2.MORPH_OPEN,kernel)
            background=np.zeros_like(frame)
            mask_3ch=cv2.cvtColor(binary_mask,cv2.COLOR_GRAY2BGR)/255.0
            result=(frame*mask_3ch+background*(1-mask_3ch)).astype(np.uint8)
            return result
        except:
            return frame
    def temporal_smoothing(self,current_landmarks,alpha=0.7):
        if self.prev_landmarks is None or len(self.prev_landmarks)!=len(current_landmarks):
            self.prev_landmarks=current_landmarks
            return current_landmarks
        smoothed=[]
        for i in range(len(current_landmarks)):
            class SmoothLandmark:pass
            landmark=SmoothLandmark()
            landmark.x=alpha*current_landmarks[i].x+(1-alpha)*self.prev_landmarks[i].x
            landmark.y=alpha*current_landmarks[i].y+(1-alpha)*self.prev_landmarks[i].y
            landmark.z=alpha*current_landmarks[i].z+(1-alpha)*self.prev_landmarks[i].z
            landmark.visibility=current_landmarks[i].visibility
            smoothed.append(landmark)
        self.prev_landmarks=smoothed
        return smoothed
    def detect_pose_multi_pass(self,frame,timestamp_ms):
        rgb_frame=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
        mp_image=mp.Image(image_format=mp.ImageFormat.SRGB,data=rgb_frame)
        result=self.pose_landmarker.detect_for_video(mp_image,timestamp_ms)
        if result.pose_landmarks:
            return result,frame
        enhanced_frame=self.enhance_image(frame)
        rgb_enhanced=cv2.cvtColor(enhanced_frame,cv2.COLOR_BGR2RGB)
        mp_image_enhanced=mp.Image(image_format=mp.ImageFormat.SRGB,data=rgb_enhanced)
        result=self.pose_landmarker.detect_for_video(mp_image_enhanced,timestamp_ms)
        if result.pose_landmarks:
            return result,enhanced_frame
        return None,frame

SIMPLER_CONNECTIONS=[(11,12),(11,23),(12,24),(23,24),(11,13),(13,15),(12,14),(14,16),(23,25),(25,27),(24,26),(26,28)]

def draw_landmarks(frame,landmarks):
    h,w=frame.shape[:2]
    for connection in SIMPLER_CONNECTIONS:
        idx1,idx2=connection
        if idx1<len(landmarks)and idx2<len(landmarks):
            if landmarks[idx1].visibility>0.5 and landmarks[idx2].visibility>0.5:
                pt1=(int(landmarks[idx1].x*w),int(landmarks[idx1].y*h))
                pt2=(int(landmarks[idx2].x*w),int(landmarks[idx2].y*h))
                if 0<=pt1[0]<w and 0<=pt1[1]<h and 0<=pt2[0]<w and 0<=pt2[1]<h:
                    cv2.line(frame,pt1,pt2,(0,255,0),3)
    for i,landmark in enumerate(landmarks):
        if landmark.visibility>0.5:
            x=int(landmark.x*w)
            y=int(landmark.y*h)
            if 0<=x<w and 0<=y<h:
                if i in[11,12]:color,radius=(0,165,255),6
                elif i in[23,24]:color,radius=(0,255,255),6
                elif i in[13,14,15,16]:color,radius=(255,0,0),5
                elif i in[25,26,27,28]:color,radius=(0,255,0),5
                else:color,radius=(0,0,255),4
                cv2.circle(frame,(x,y),radius,color,-1)
                cv2.circle(frame,(x,y),radius,(255,255,255),1)

def process_video(input_video,remove_background=False,progress=gr.Progress()):
    if input_video is None:
        return None
    temp_dir=tempfile.mkdtemp()
    output_path=tempfile.NamedTemporaryFile(delete=False,suffix='.mp4')
    output_path.close()
    cap=cv2.VideoCapture(input_video)
    if not cap.isOpened():
        raise ValueError("Could not open video")
    fps=int(cap.get(cv2.CAP_PROP_FPS))or 30
    total_frames=int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    detector=PoseDetector(confidence=0.25)
    frame_count=0
    poses_detected=0
    while True:
        ret,frame=cap.read()
        if not ret:
            break
        progress(frame_count/total_frames,desc=f"Frame {frame_count}/{total_frames}")
        timestamp_ms=int(frame_count*(1000/fps))
        result,processed_frame=detector.detect_pose_multi_pass(frame,timestamp_ms)
        if result and result.pose_landmarks:
            poses_detected+=1
            for landmarks in result.pose_landmarks:
                smoothed_landmarks=detector.temporal_smoothing(landmarks)
                draw_landmarks(processed_frame,smoothed_landmarks)
                if remove_background and result.segmentation_masks:
                    processed_frame=detector.remove_background(processed_frame,result.segmentation_masks[0])
        frame_path=os.path.join(temp_dir,f"frame_{frame_count:06d}.jpg")
        cv2.imwrite(frame_path,processed_frame)
        frame_count+=1
    cap.release()
    progress(0.95,desc="Encoding video...")
    cmd=['ffmpeg','-y','-framerate',str(fps),'-pattern_type','glob','-i',f'{temp_dir}/frame_*.jpg','-c:v','libx264','-profile:v','main','-level','3.0','-pix_fmt','yuv420p','-movflags','+faststart','-crf','23','-preset','medium',output_path.name]
    try:
        result=subprocess.run(cmd,capture_output=True,text=True,timeout=300)
        if result.returncode==0 and os.path.exists(output_path.name):
            shutil.rmtree(temp_dir)
            return output_path.name
        else:
            return input_video
    except:
        shutil.rmtree(temp_dir)
        return input_video

def get_default_video():
    video_list=['obj_01.mp4','obj_04.mp4','obj_05.mp4','obj_07.mp4']
    for video_path in video_list:
        if os.path.exists(video_path):
            return video_path
    return None

with gr.Blocks(title="FormAI - Body Movement Tracker",theme=gr.themes.Soft())as demo:
    gr.Markdown("# 🏋️ FormAI - Your Virtual Movement Coach\n\n**Choose tracking mode:**\n- **Standard Tracking**: Faster processing, keeps original background\n- **Background Removal**: Better for complex backgrounds, creates clean output\n")
    with gr.Row():
        with gr.Column(scale=2):
            video_input=gr.Video(label="Select Video",interactive=True,sources=["upload"],height=350,value=get_default_video())
            gr.Markdown("### Sample Videos")
            with gr.Row():
                videos=[('01','obj_01.mp4'),('03','obj_04.mp4'),('05','obj_05.mp4'),('07','obj_07.mp4')]
                for label,path in videos:
                    if os.path.exists(path):
                        btn=gr.Button(f"🎬 Video {label}",variant="secondary",size="sm")
                        btn.click(fn=lambda p=path:p,outputs=video_input)
            gr.Markdown("### Tracking Options")
            with gr.Row():
                track_btn=gr.Button("🎯 Standard Tracking",variant="primary",size="lg")
                track_bg_btn=gr.Button("🎯 Track with Background Removal",variant="secondary",size="lg")
        with gr.Column(scale=3):
            video_output=gr.Video(label="Tracked Output",height=500,autoplay=True)
    if get_default_video():
        demo.load(fn=lambda:process_video(get_default_video(),False),outputs=video_output)
    track_btn.click(fn=lambda vid:process_video(vid,False),inputs=video_input,outputs=video_output)
    track_bg_btn.click(fn=lambda vid:process_video(vid,True),inputs=video_input,outputs=video_output)

if __name__=="__main__":
    demo.launch(server_name="0.0.0.0",server_port=7860)