import cv2
import tkinter as tk
import mediapipe as mp
import numpy as np
import os
import math
from rembg import remove 
from PIL import Image

import dobot
import vector_draw

# Load images with transparency
mario_hat_image_path = "Filters/MArio.png"
sunglasses_image_path = "Filters/glasses.png"
moustache_image_path = "Filters/MoustacheMario.png"

# Load images
mario_hat = cv2.imread(mario_hat_image_path, cv2.IMREAD_UNCHANGED)
sunglasses = cv2.imread(sunglasses_image_path, cv2.IMREAD_UNCHANGED)
moustache = cv2.imread(moustache_image_path, cv2.IMREAD_UNCHANGED)

# Check if images were loaded correctly
if mario_hat is None:
    print("Error: Mario hat image not found.")
    exit()
if sunglasses is None:
    print("Error: Sunglasses image not found.")
    exit()
if moustache is None:
    print("Error: Moustache image not found.")
    exit()

# Initialize MediaPipe FaceMesh
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(min_detection_confidence=0.5, min_tracking_confidence=0.5)

# Variables for toggling filters
mario_hat_active = False
sunglasses_active = False
moustache_active = False
show_angles = False

# Open webcam for capturing live feed
cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("Error: The webcam cannot be opened")
    exit()

# Variable to hold the contour frame
contour_frame = None
resized_edges = None

def calculate_angles(landmarks):
    left_eye = np.array(landmarks[33])
    right_eye = np.array(landmarks[263])
    nose_tip = np.array(landmarks[1])
    chin = np.array(landmarks[152])
    yaw = math.degrees(math.atan2(right_eye[1] - left_eye[1], right_eye[0] - left_eye[0]))
    pitch = math.degrees(math.atan2(chin[1] - nose_tip[1], chin[0] - nose_tip[0]))
    return yaw, pitch

def apply_mario_hat(frame, landmarks):
    global mario_hat
    if mario_hat_active and mario_hat is not None:
        forehead = landmarks[10]
        chin = landmarks[152]
        left_side = landmarks[234]
        right_side = landmarks[454]
        face_width = int(np.linalg.norm(np.array(left_side) - np.array(right_side)))
        hat_width = int(face_width * 4.0)
        hat_height = int(hat_width * mario_hat.shape[0] / mario_hat.shape[1])
        mario_hat_resized = cv2.resize(mario_hat, (hat_width, hat_height))
        x = int(forehead[0] - hat_width / 2)
        y = int(forehead[1] - hat_height * 0.7)
        alpha_channel = mario_hat_resized[:, :, 3] / 255.0
        hat_rgb = mario_hat_resized[:, :, :3]
        for i in range(hat_height):
            for j in range(hat_width):
                if 0 <= y + i < frame.shape[0] and 0 <= x + j < frame.shape[1]:
                    alpha = alpha_channel[i, j]
                    if alpha > 0:
                        for c in range(3):
                            frame[y + i, x + j, c] = (1 - alpha) * frame[y + i, x + j, c] + alpha * hat_rgb[i, j, c]
    return frame

def apply_sunglasses(frame, landmarks):
    global sunglasses
    if sunglasses_active and sunglasses is not None:
        left_eye = landmarks[33]
        right_eye = landmarks[263]
        eye_dist = np.linalg.norm(np.array(left_eye) - np.array(right_eye))
        scaling_factor = 1.75
        sunglasses_width = int(eye_dist * scaling_factor)
        sunglasses_height = int(sunglasses_width * sunglasses.shape[0] / sunglasses.shape[1])
        sunglasses_resized = cv2.resize(sunglasses, (sunglasses_width, sunglasses_height))
        center_x = int((left_eye[0] + right_eye[0]) / 2)
        center_y = int((left_eye[1] + right_eye[1]) / 2)
        x = int(center_x - sunglasses_resized.shape[1] / 2)
        y = int(center_y - sunglasses_resized.shape[0] / 2)
        alpha_channel = sunglasses_resized[:, :, 3] / 255.0
        sunglasses_rgb = sunglasses_resized[:, :, :3]
        for i in range(sunglasses_resized.shape[0]):
            for j in range(sunglasses_resized.shape[1]):
                if alpha_channel[i, j] > 0:
                    for c in range(3):
                        frame[y + i, x + j, c] = (1 - alpha_channel[i, j]) * frame[y + i, x + j, c] + alpha_channel[i, j] * sunglasses_rgb[i, j, c]
    return frame

def apply_moustache(frame, landmarks):
    global moustache
    if moustache_active and moustache is not None:
        nose_base = landmarks[1]
        mouth_left = landmarks[61]
        mouth_right = landmarks[291]
        mouth_width = int(np.linalg.norm(np.array(mouth_left) - np.array(mouth_right)))
        moustache_width = int(mouth_width * 1.5)
        moustache_height = int(moustache_width * moustache.shape[0] / moustache.shape[1])
        moustache_resized = cv2.resize(moustache, (moustache_width, moustache_height))
        x = int(nose_base[0] - moustache_width / 2)
        y = int(nose_base[1])
        alpha_channel = moustache_resized[:, :, 3] / 255.0
        moustache_rgb = moustache_resized[:, :, :3]
        for i in range(moustache_height):
            for j in range(moustache_width):
                if 0 <= y + i < frame.shape[0] and 0 <= x + j < frame.shape[1]:
                    alpha = alpha_channel[i, j]
                    if alpha > 0:
                        for c in range(3):
                            frame[y + i, x + j, c] = (1 - alpha) * frame[y + i, x + j, c] + alpha * moustache_rgb[i, j, c]
    return frame

def update_frame():
    global mario_hat_active, sunglasses_active, show_angles, contour_frame, moustache_active
    ret, frame = cap.read()
    if ret:
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = face_mesh.process(rgb_frame)
        if results.multi_face_landmarks:
            for face_landmarks in results.multi_face_landmarks:
                landmarks = [(lm.x * frame.shape[1], lm.y * frame.shape[0]) for lm in face_landmarks.landmark]
                yaw, pitch = calculate_angles(landmarks)
                if mario_hat_active:
                    frame = apply_mario_hat(frame, landmarks)
                if sunglasses_active:
                    frame = apply_sunglasses(frame, landmarks)
                if moustache_active:
                    frame = apply_moustache(frame, landmarks)
                if show_angles:
                    cv2.putText(frame, f"Yaw: {yaw:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                    cv2.putText(frame, f"Pitch: {pitch:.2f}", (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.imshow("Webcam Feed", frame)
        contour_frame = frame
    root.after(100, update_frame)

def toggle_mario_hat():
    global mario_hat_active
    mario_hat_active = not mario_hat_active
    status = "activated" if mario_hat_active else "deactivated"
    print(f"Mario hat filter {status}")

def toggle_sunglasses():
    global sunglasses_active
    sunglasses_active = not sunglasses_active
    status = "activated" if sunglasses_active else "deactivated"
    print(f"Sunglasses filter {status}")

def toggle_moustache():
    global moustache_active
    moustache_active = not moustache_active
    status = "activated" if moustache_active else "deactivated"
    print(f"Moustache filter {status}")

def toggle_angles():
    global show_angles
    show_angles = not show_angles
    status = "shown" if show_angles else "hidden"
    print(f"Angles display {status}")

def show_contour_frame():
    if contour_frame is not None:
        # Display the result
        cv2.imshow('Edges', resized_edges)


def save_image():
    global contour_frame, resized_edges
    if contour_frame is not None:
        save_path = "Tmp/captured_face.png"
        cv2.imwrite(save_path, contour_frame)
        print(f"Image saved to {save_path}")

        # Store path of the image in the variable input_path 
        input_path =  'Tmp/captured_face.png' 
        
        # Store path of the output image in the variable output_path 
        output_path = 'Tmp/captured_face_nobg.png' 
        
        # Processing the image 
        input = Image.open(input_path) 
        
        # Removing the background from the given Image 
        output = remove(input) 
        
        #Saving the image in the given path 
        output.save(output_path)
        image = cv2.imread(output_path, cv2.IMREAD_GRAYSCALE)
        mask = (image > 1) & (image < 254)
        blurred_image = cv2.GaussianBlur(image, (9, 9), 0)
        median_val = np.median(blurred_image[mask])
        lower_threshold = int(max(0, 0.5 * median_val))
        upper_threshold = int(min(255, 1.2 * median_val))
        print(f"Automatic lower threshold: {lower_threshold}")
        print(f"Automatic upper threshold: {upper_threshold}")

        # Apply Canny edge detection using the calculated thresholds
        edges = cv2.Canny(blurred_image, lower_threshold, upper_threshold)

        # Resize the output image to a smaller size (e.g., 50% of the original size)
        output_height, output_width = edges.shape[:2]
        resized_edges = cv2.resize(edges, (output_width // 2, output_height // 2), interpolation=cv2.INTER_AREA)

        # Save the resized result to a file
        cv2.imwrite('Tmp/final_output_image.png', resized_edges)

def start_dobot():
    vector_draw.vector_draw()

# Tkinter GUI setup
root = tk.Tk()
root.title("Control Tab")
root.geometry("300x370")
root.configure(bg="#004346")

# Buttons on the control window with updated font and colors
mario_hat_button = tk.Button(root, text="Add Mario Hat", font=("Arial", 12, "bold"), command=toggle_mario_hat, bg="#4C8577", fg="white", padx=10, pady=5, height=1, width=20)
mario_hat_button.pack(pady=10)

sunglasses_button = tk.Button(root, text="Add Glasses", font=("Arial", 12, "bold"), command=toggle_sunglasses, bg="#4C8577", fg="white", padx=10, pady=5, height=1, width=20)
sunglasses_button.pack(pady=10)

moustache_button = tk.Button(root, text="Add Mario Moustache", font=("Arial", 12, "bold"), command=toggle_moustache, bg="#4C8577", fg="white", padx=10, pady=5,height=1, width=20)
moustache_button.pack(pady=10)

save_image_button = tk.Button(root, text="Save/Retake Image", font=("Arial", 12, "bold"), command=save_image, bg="#49A078", fg="white", padx=10, pady=5,height=1, width=20)
save_image_button.pack(pady=10)

contour_frame_button = tk.Button(root, text="Show Contour Image", font=("Arial", 12, "bold"), command=show_contour_frame, bg="#216869", fg="white", padx=10, pady=5,height=1, width=20)
contour_frame_button.pack(pady=10)

contour_frame_button = tk.Button(root, text="Start Dobot Drawing", font=("Arial", 12, "bold"), command=start_dobot, bg="#49A078", fg="white", padx=10, pady=5,height=1, width=20)
contour_frame_button.pack(pady=10)

# Graceful exit
def on_closing():
    cap.release()
    cv2.destroyAllWindows()
    root.destroy()

root.protocol("WM_DELETE_WINDOW", on_closing)

show_contour_frame()

# Start Tkinter event loop and OpenCV frame updates
update_frame()
root.mainloop()