import gradio as gr
from matplotlib.pyplot import draw
import mediapipe as mp
import numpy as np
import tempfile
import mediapy as media
import log_utils

logger = log_utils.get_logger()

mp_hands = mp.solutions.hands
mp_hands_connections = mp.solutions.hands_connections
hands = mp_hands.Hands()
mp_draw = mp.solutions.drawing_utils 

connections = {
    'HAND_CONNECTIONS': mp_hands_connections.HAND_CONNECTIONS,
    'HAND_PALM_CONNECTIONS': mp_hands_connections.HAND_PALM_CONNECTIONS,
    'HAND_THUMB_CONNECTIONS': mp_hands_connections.HAND_THUMB_CONNECTIONS,
    'HAND_INDEX_FINGER_CONNECTIONS': mp_hands_connections.HAND_INDEX_FINGER_CONNECTIONS,
    'HAND_MIDDLE_FINGER_CONNECTIONS': mp_hands_connections.HAND_MIDDLE_FINGER_CONNECTIONS,
    'HAND_RING_FINGER_CONNECTIONS': mp_hands_connections.HAND_RING_FINGER_CONNECTIONS,
    'HAND_PINKY_FINGER_CONNECTIONS': mp_hands_connections.HAND_PINKY_FINGER_CONNECTIONS,
}

def draw_landmarks(img, selected_connections, draw_background):
  results = hands.process(img)
  output_img = img if draw_background else np.zeros_like(img)  
  if results.multi_hand_landmarks:
    for hand_landmarks in results.multi_hand_landmarks:
      mp_draw.draw_landmarks(output_img, hand_landmarks, connections[selected_connections])    
  return output_img

def process_image(img, selected_connections, draw_background):
  logger.info(f"Processing image with connections: {selected_connections}, draw background: {draw_background}")
  return draw_landmarks(img, selected_connections, draw_background)

def process_video(video_path, selected_connections, draw_background):
  logger.info(f"Processing video with connections: {selected_connections}, draw background: {draw_background}")
  with tempfile.NamedTemporaryFile() as f:
    out_path = f"{f.name}.{video_path.split('.')[-1]}"
    with media.VideoReader(video_path) as r:
      with media.VideoWriter(
          out_path, shape=r.shape, fps=r.fps, bps=r.bps) as w:
        for image in r:
          w.add_image(draw_landmarks(image, selected_connections, draw_background))
    return out_path


demo = gr.Blocks()

with demo:
  gr.Markdown(
    """
    # Hand & Finger Tracking
    This is a demo of hand and finger tracking using [Google's MediaPipe](https://google.github.io/mediapipe/solutions/hands.html).
    """)

  with gr.Column():  
    draw_background = gr.Checkbox(value=True, label="Draw background?")
    connection_keys = list(connections.keys())
    selected_connections = gr.Dropdown(
      label="Select connections to draw", 
      choices=connection_keys,
      value=connection_keys[0],
      )
    with gr.Tabs():
      with gr.TabItem(label="Upload an image"):
        uploaded_image = gr.Image(type="numpy")
        submit_uploaded_image = gr.Button(value="Process Image")
      with gr.TabItem(label="Take a picture"):
        camera_picture = gr.Image(source="webcam", type="numpy")
        submit_camera_picture = gr.Button(value="Process Image")  
      with gr.TabItem(label="Record a video"):
        recorded_video = gr.Video(source="webcam", format="mp4")
        submit_recorded_video = gr.Button(value="Process Video")
      with gr.TabItem(label="Upload a video"):
        uploaded_video = gr.Video(format="mp4")
        submit_uploaded_video = gr.Button(value="Process Video")  

  with gr.Column():
    processed_video = gr.Video()
    processed_image = gr.Image()

  gr.Markdown('<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=kristyc.mediapipe-hands" />')
  setting_inputs = [selected_connections, draw_background]
  submit_uploaded_image.click(fn=process_image, inputs=[uploaded_image, *setting_inputs], outputs=[processed_image])
  submit_camera_picture.click(fn=process_image, inputs=[camera_picture, *setting_inputs], outputs=[processed_image])
  submit_recorded_video.click(fn=process_video, inputs=[recorded_video, *setting_inputs], outputs=[processed_video])
  submit_uploaded_video.click(fn=process_video, inputs=[recorded_video, *setting_inputs], outputs=[processed_video])

demo.launch()