import gradio as gr
import cv2
import os
import mediapipe as mp
import numpy as np
import tempfile
from maps import *
from functions import *

class CosmeticInjectionVisualizer:
    def __init__(self, muscles_map, tasks_map):
        self.mp_face_mesh = mp.solutions.face_mesh
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles
        self.muscles_map = muscles_map
        self.tasks_map = tasks_map

    def process_image(self, image, task_name):
        frame_shape = image.shape
        with self.mp_face_mesh.FaceMesh(
            static_image_mode=True,
            refine_landmarks=True,
            max_num_faces=1,
            min_detection_confidence=0.5) as face_mesh:

            results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            if results.multi_face_landmarks:
                for face_landmarks in results.multi_face_landmarks:
                    points = self.get_muscle_points(task_name, face_landmarks, frame_shape)
                    self.draw_muscle_points(image, points, face_landmarks, task_name)
        return image

    def draw_rounded_rectangle(self,image, start_point, end_point, color, thickness, radius):
      top_left = start_point
      bottom_right = end_point

      if thickness < 0:  # Filled rectangle
          thickness = cv2.FILLED

      # Draw filled rectangle with rounded corners
      if thickness == cv2.FILLED:
          # Top-left corner
          cv2.ellipse(image, (top_left[0] + radius, top_left[1] + radius), (radius, radius), 180, 0, 90, color, -1)
          # Top-right corner
          cv2.ellipse(image, (bottom_right[0] - radius, top_left[1] + radius), (radius, radius), 270, 0, 90, color, -1)
          # Bottom-left corner
          cv2.ellipse(image, (top_left[0] + radius, bottom_right[1] - radius), (radius, radius), 90, 0, 90, color, -1)
          # Bottom-right corner
          cv2.ellipse(image, (bottom_right[0] - radius, bottom_right[1] - radius), (radius, radius), 0, 0, 90, color, -1)

          # Top and bottom border
          cv2.rectangle(image, (top_left[0] + radius, top_left[1]), (bottom_right[0] - radius, top_left[1] + radius), color, -1)
          cv2.rectangle(image, (top_left[0] + radius, bottom_right[1] - radius), (bottom_right[0] - radius, bottom_right[1]), color, -1)
          # Left and right border
          cv2.rectangle(image, (top_left[0], top_left[1] + radius), (top_left[0] + radius, bottom_right[1] - radius), color, -1)
          cv2.rectangle(image, (bottom_right[0] - radius, top_left[1] + radius), (bottom_right[0], bottom_right[1] - radius), color, -1)
          # Center rectangle
          cv2.rectangle(image, (top_left[0] + radius, top_left[1] + radius), (bottom_right[0] - radius, bottom_right[1] - radius), color, -1)
      else:
          # Top-left corner
          cv2.ellipse(image, (top_left[0] + radius, top_left[1] + radius), (radius, radius), 180, 0, 90, color, thickness)
          # Top-right corner
          cv2.ellipse(image, (bottom_right[0] - radius, top_left[1] + radius), (radius, radius), 270, 0, 90, color, thickness)
          # Bottom-left corner
          cv2.ellipse(image, (top_left[0] + radius, bottom_right[1] - radius), (radius, radius), 90, 0, 90, color, thickness)
          # Bottom-right corner
          cv2.ellipse(image, (bottom_right[0] - radius, bottom_right[1] - radius), (radius, radius), 0, 0, 90, color, thickness)

          # Top border
          cv2.line(image, (top_left[0] + radius, top_left[1]), (bottom_right[0] - radius, top_left[1]), color, thickness)
          # Bottom border
          cv2.line(image, (top_left[0] + radius, bottom_right[1]), (bottom_right[0] - radius, bottom_right[1]), color, thickness)
          # Left border
          cv2.line(image, (top_left[0], top_left[1] + radius), (top_left[0], bottom_right[1] - radius), color, thickness)
          # Right border
          cv2.line(image, (bottom_right[0], top_left[1] + radius), (bottom_right[0], bottom_right[1] - radius), color, thickness)

    def draw_muscle_points(self, image, points, face_landmarks, task, draw_background=False, verbose=False):
        # Calculate bounding box of the face landmarks
        x_min = min([landmark.x for landmark in face_landmarks.landmark]) * image.shape[1]
        y_min = min([landmark.y for landmark in face_landmarks.landmark]) * image.shape[0]
        x_max = max([landmark.x for landmark in face_landmarks.landmark]) * image.shape[1]
        y_max = max([landmark.y for landmark in face_landmarks.landmark]) * image.shape[0]
        face_width = x_max - x_min
        face_height = y_max - y_min

        # Determine text size and circle size based on face size relative to the image
        scale_factor = 0.0005 * (face_width + face_height)
        text_scale = scale_factor
        thickness = int(2 * scale_factor)
        margin = int(10 * scale_factor)
        circle_radius = int(5 * scale_factor)
        radius = int(10 * scale_factor)  # For rounded corners

        muscle_names = set()
        for (x, y, muscle) in points:
            # Draw the circle for the muscle point
            cv2.circle(image, (x, y), circle_radius, (0, 0, 255), -1)

            # Determine text size and background size
            text = "3 U"
            (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, text_scale, thickness)
            text_x = x
            text_y = y - 10
            text_x_end = text_x + text_width + 2 * margin
            text_y_end = text_y - text_height - 2 * margin

            if draw_background:
                # Draw background rectangle with margins and rounded corners
                self.draw_rounded_rectangle(image, (text_x - margin, text_y + margin), (text_x_end, text_y_end), (0, 0, 0), cv2.FILLED, radius)

            if verbose:
              # Draw the text on top of the rectangle
              cv2.putText(image, text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, text_scale, (0, 255, 0), thickness, cv2.LINE_AA)

            muscle_names.add(muscle)
        muscles = ','.join(muscle_names)
        # Draw legend in the bottom-left corner with background
        legend_texts = ['Total dose : 42','Name of patient: Julia Juila',f'Muscle : {muscles}', f'Task : {task}','Cosmetic App']
        legend_x = 10
        legend_y = image.shape[0] - 10
        legend_margin = 5
        legend_scale_factor = 1.5 * text_scale  # Make legend text larger
        legend_thickness = int(2 * legend_scale_factor)
        legend_radius = int(10 * legend_scale_factor)
        max_text_width = 0
        total_text_height = 0

        for text in legend_texts:
            (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, legend_scale_factor, legend_thickness)
            max_text_width = max(max_text_width, text_width)
            total_text_height += text_height + 2 * legend_margin

        legend_start_point = (legend_x - legend_margin, legend_y - total_text_height - legend_margin)
        legend_end_point = (legend_x + max_text_width + legend_margin, legend_y + legend_margin)

        if True:
            self.draw_rounded_rectangle(image, legend_start_point, legend_end_point, (0, 0, 0), cv2.FILLED, legend_radius)

        for text in legend_texts:
            (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, legend_scale_factor, legend_thickness)
            cv2.putText(image, text, (legend_x, legend_y), cv2.FONT_HERSHEY_SIMPLEX, legend_scale_factor, (255, 255, 255), legend_thickness, cv2.LINE_AA)
            legend_y -= text_height + 2 * legend_margin


    def get_muscle_points(self, task_name, face_landmarks, frame_shape):
        if task_name not in self.tasks_map:
            raise ValueError(f"Task '{task_name}' not found in tasks map.")
        muscles_names = self.tasks_map[task_name]['muscles']
        points = []
        for muscle in muscles_names:
            for region in self.muscles_map[muscle]:
                if 'points' in self.muscles_map[muscle][region]:
                    for point_idx in self.muscles_map[muscle][region]['points']:
                        landmark = face_landmarks.landmark[point_idx]
                        x = int(landmark.x * frame_shape[1])
                        y = int(landmark.y * frame_shape[0])
                        points.append((x, y, muscle))
                else:
                    for subregion in self.muscles_map[muscle][region]:
                        for point_idx in self.muscles_map[muscle][region][subregion]['points']:
                            landmark = face_landmarks.landmark[point_idx]
                            x = int(landmark.x * frame_shape[1])
                            y = int(landmark.y * frame_shape[0])
                            points.append((x, y, muscle))
        return points

    def process_video(self, video_path, task_name):
        cap = cv2.VideoCapture(video_path)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        output_path = temp_file.name

        # Get the width and height of the frames
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(cap.get(cv2.CAP_PROP_FPS))

        # Define the codec and create VideoWriter object
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Use appropriate codec for .mp4 files
        out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))

        with self.mp_face_mesh.FaceMesh(
            static_image_mode=False,
            refine_landmarks=True,
            max_num_faces=1,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5) as face_mesh:

            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                frame_shape = frame.shape
                results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                if results.multi_face_landmarks:
                    for face_landmarks in results.multi_face_landmarks:
                        points = self.get_muscle_points(task_name, face_landmarks, frame_shape)
                        self.draw_muscle_points(frame, points, face_landmarks, task_name)
                
                # Write the processed frame to the output video file
                out.write(frame)

            cap.release()
            out.release()
            cv2.destroyAllWindows()

        return output_path

    def process_webcam(self, task_name):
        cap = cv2.VideoCapture(0)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        output_path = temp_file.name

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, 20.0, (640, 480))
        
        with self.mp_face_mesh.FaceMesh(
            static_image_mode=False,
            refine_landmarks=True,
            max_num_faces=1,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5) as face_mesh:

            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                frame_shape = frame.shape
                results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                if results.multi_face_landmarks:
                    for face_landmarks in results.multi_face_landmarks:
                        points = self.get_muscle_points(task_name, face_landmarks, frame_shape)
                        self.draw_muscle_points(frame, points, face_landmarks, task_name)
                out.write(frame)

            cap.release()
            out.release()
            cv2.destroyAllWindows()

        return output_path

visualizer = CosmeticInjectionVisualizer(muscles_map, tasks_map)

def inference_image(image, task_name):
    result_image = visualizer.process_image(image, task_name)
    return result_image

def inference_video(video_path, task_name):
    result_video_path = visualizer.process_video(video_path, task_name)
    return result_video_path

def inference_webcam(task_name):
    result_video_path = visualizer.process_webcam(task_name)
    return result_video_path

task_names = list(tasks_map.keys())
base_path=os.getcwd()
default_image_path = os.path.join(base_path,'image.jpg')
default_video_path = os.path.join(base_path,'video.mp4')

with gr.Blocks() as demo:
    gr.Markdown("# Cosmetic Injection Visualizer")
    
    with gr.Tabs():
        with gr.TabItem("Image"):
            image_input = gr.Image(type="numpy", label="Input Image", value=default_image_path)
            task_input_image = gr.Dropdown(choices=task_names, label="Task Name")
            image_output = gr.Image(type="numpy", label="Output Image")
            gr.Button("Process Image").click(inference_image, inputs=[image_input, task_input_image], outputs=image_output)

        with gr.TabItem("Video"):
            video_input = gr.Video(label="Input Video", value=default_video_path)
            task_input_video = gr.Dropdown(choices=task_names, label="Task Name")
            video_output = gr.Video(label="Output Video")
            gr.Button("Process Video").click(inference_video, inputs=[video_input, task_input_video], outputs=video_output)

        with gr.TabItem("Webcam"):
            task_input_webcam = gr.Dropdown(choices=task_names, label="Task Name")
            webcam_output = gr.Video(label="Output Video")
            gr.Button("Process Webcam").click(inference_webcam, inputs=[task_input_webcam], outputs=webcam_output)

demo.launch()