from cmath import atan, pi
from email.contentmanager import raw_data_manager
import cv2
from cv2 import sqrt
import mediapipe as mp
import numpy as np

from scipy.misc import face #delete
import scipy.ndimage
import scipy.interpolate as interp

# Takes in coordinates of a trapezoid ordered counter clockwise
# and creates a square cutout from the frame
def create_cut_out(frame, coords, size):

    p1, p2, p3, p4 = coords

    r_channel = []
    b_channel = []
    g_channel = []

    length_x = np.sqrt(((p1[0]-p4[0]) * (p1[0]-p4[0])) + ((p1[1]-p4[1]) * (p1[1]-p4[1])))
    length_y = np.sqrt(((p1[0]-p2[0]) * (p1[0]-p2[0])) + ((p1[1]-p2[1]) * (p1[1]-p2[1])))

    # Assign base points
    base_point1 = p1
    base_point2 = p4

    # Get the angles for the base offset and between the base and the two arms
    alpha = np.arctan2(base_point2[1]-base_point1[1], base_point2[0]-base_point1[0])    # Base offset from horizontal position
    beta = np.arctan2(p2[1]-base_point1[1], p2[0]-base_point1[0])                       # Angle between the x axis and the left arm
    gama = np.arctan2(p3[1]-base_point2[1], p3[0]-base_point2[0])                       # Angle between the x axis and the right arm

    # Calculate the height of the trapezoid
    height = length_y*np.sin(np.pi-abs(alpha-beta))

    # Generic rotation matrices
    alpha_matrix = [[np.cos(alpha), -np.sin(alpha)], [np.sin(alpha), np.cos(alpha)]]
    beta_matrix = [[np.cos(beta), -np.sin(beta)], [np.sin(beta), np.cos(beta)]]
    gama_matrix = [[np.cos(gama), -np.sin(gama)], [np.sin(gama), np.cos(gama)]]

    # Now we're going to take the pixels of each row (offset by angle alpha) of the trapezoid up to its height
    for ycoord in range(0, int(height)):
        left_arm_point = np.dot(beta_matrix, [ycoord, 0]) + base_point1
        right_arm_point = np.dot(gama_matrix, [ycoord, 0]) + base_point2 # Revolves around 4th corner
        length_x = np.sqrt((left_arm_point[0]-right_arm_point[0]) * (left_arm_point[0]-right_arm_point[0]) + (left_arm_point[1]-right_arm_point[1]) * (left_arm_point[1]-right_arm_point[1]))
        
        # First pass we prepare temporal variables that then end up in the different video channels
        if ycoord == 0:
            r_tmp = []
            b_tmp = []
            g_tmp = []
        else: # After first pass, squish the extracted row to desired size
            # Convert to numpy arrays for interpolation
            r_tmp = np.array(r_tmp)
            b_tmp = np.array(b_tmp)
            g_tmp = np.array(g_tmp)

            # Squish the R channel and append
            resized = interp.interp1d(np.arange(r_tmp.size), r_tmp)
            r_tmp = resized(np.linspace(0,r_tmp.size-1, size))
            r_channel.append(r_tmp.astype(int))

            # Squish the G channel and append
            resized = interp.interp1d(np.arange(g_tmp.size), g_tmp)
            g_tmp = resized(np.linspace(0,g_tmp.size-1, size))
            g_channel.append(g_tmp.astype(int))

            # Squish the B channel and append
            resized = interp.interp1d(np.arange(b_tmp.size), b_tmp)
            b_tmp = resized(np.linspace(0,b_tmp.size-1, size))
            b_channel.append(b_tmp.astype(int))

            # Re-initialize temporary arrays
            r_tmp = []
            b_tmp = []
            g_tmp = []

        # Now extract the appropriate row of the trapezoid
        for xcoord in range(0, int(length_x)):
            pixel = np.dot(alpha_matrix, [xcoord, 0]) + left_arm_point
            r_tmp.append(frame[int(pixel[1]), int(pixel[0]), 0])
            frame[int(pixel[1]), int(pixel[0]), 0] = 0
            g_tmp.append(frame[int(pixel[1]), int(pixel[0]), 1])
            frame[int(pixel[1]), int(pixel[0]), 1] = 0
            b_tmp.append(frame[int(pixel[1]), int(pixel[0]), 2])
            frame[int(pixel[1]), int(pixel[0]), 2] = 0

    # The above code guarantees that the cut out has the correct amount of columns
    # Now we stretch it to fit the row requirement

    # Stretching R channel
    size_difference = size/len(r_channel)
    r_channel = scipy.ndimage.zoom(r_channel, (size_difference,1), order=1)

    # Stretching G channel
    size_difference = size/len(g_channel)
    g_channel = scipy.ndimage.zoom(g_channel, (size_difference,1), order=1)

    # Stretching B channel
    size_difference = size/len(b_channel)
    b_channel = scipy.ndimage.zoom(b_channel, (size_difference,1), order=1)

    #rgb_out = np.dstack((r_channel, g_channel, b_channel))
    rgb_out = (np.dstack((r_channel,g_channel,b_channel))).astype(np.uint8)
    return rgb_out, frame

def detect_face(video_name, video_out_name, fps, cutout_size = 32):
    mp_face_mesh = mp.solutions.face_mesh
    face_mesh = mp_face_mesh.FaceMesh(min_detection_confidence = 0.5, min_tracking_confidence = 0.5)

    mp_drawing = mp.solutions.drawing_utils
    #mp_drawing_styles = mp.solutions.drawing_styles delete
    drawing_spec = mp_drawing.DrawingSpec(thickness = 1, circle_radius = 1)

    video = cv2.VideoCapture(video_name)

    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    writer = cv2.VideoWriter(video_out_name, fourcc, fps, (cutout_size,cutout_size))
    #video = cv2.VideoCapture(0)

    while video.isOpened():
        
        ret, image = video.read()
        original = image

        if ret == True:
            # Maybe Delete this line later
            # Convert the image from BGR to RGB and flip it
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            # Improve performace
            image.flags.writeable = False

            results = face_mesh.process(image)

            # Improve performace
            image.flags.writeable = True

            # Maybe Delete this line later
            # Convert the image back to BRG
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

            img_h, img_w, img_c = image.shape
            face_3d = []
            face_2d = []


            if results.multi_face_landmarks:
                for face_landmarks in results.multi_face_landmarks:
                    for idx, lm in enumerate(face_landmarks.landmark):
                        if idx == 33 or idx == 263 or idx == 1 or idx == 61 or idx == 291 or idx == 199 or idx == 107 or idx == 109 or idx == 336 or idx == 338 or idx == 151:
                            if idx == 1:
                                nose_2d = (lm.x * img_w, lm.y * img_h)
                                nose_3d = (lm.x * img_w, lm.y * img_h, lm.z * 3000)
                            if idx == 151:
                                middle_forehead = (lm.x * img_w, lm.y * img_h)
                            if idx == 107:
                                point1 = (lm.x * img_w, lm.y * img_h)
                            if idx == 109:
                                point2 = (lm.x * img_w, lm.y * img_h)
                            if idx == 338:
                                point3 = (lm.x * img_w, lm.y * img_h)
                            if idx == 336:
                                point4 = (lm.x * img_w, lm.y * img_h)

                            x, y = int(lm.x * img_w), int(lm.y * img_h)

                            # Get 2D coordinates
                            face_2d.append([x, y])

                            # Get 3D coordinates
                            face_3d.append([x, y, lm.z])
                    #cv2.circle(original, (int(middle_forehead[0]), int(middle_forehead[1])), 2, (255, 0, 0), 1)

                    if point1[0] < img_w and point2[0] < img_w and point3[0] < img_w and point4[0] < img_w:
                        if point1[1] < img_h and point2[1] < img_h and point3[1] < img_h and point4[1] < img_h:
                            tmp, frame = create_cut_out(original, (point1, point2, point3, point4), cutout_size)
                    # Convert it to NumPy array
                    face_2d = np.array(face_2d, dtype=np.float64)
                    face_3d = np.array(face_3d, dtype=np.float64)

                    # Camera matrix, maybe delete
                    focal_length = 1 * img_w

                    cam_matrix = np.array([ [focal_length, 0, img_h / 2],
                                            [0, focal_length, img_w / 2],
                                            [0, 0, 1]])

                    # Distortion parameters
                    dist_matrix = np.zeros((4, 1), dtype=np.float64)

                    # Solve PnP
                    success, rot_vec, trans_vec = cv2.solvePnP(face_3d, face_2d, cam_matrix, dist_matrix)

                    # Get rotational matrix
                    rmat, jac = cv2.Rodrigues(rot_vec)

                    # Get angles
                    angles, mtxR, mtxQ, Qx, Qy, Qz = cv2.RQDecomp3x3(rmat)

                    # Get the y rotation degree
                    x = angles[0] * 360
                    y = angles[1] * 360
                    z = angles[2] * 360

                    # Display the nose direction
                    nose_3d_projection, jacobian = cv2.projectPoints(nose_3d, rot_vec, trans_vec, cam_matrix, dist_matrix)

                    p1 = (int(nose_2d[0]), int(nose_2d[1]))
                    p2 = (int(nose_2d[0] + y * 10), int(nose_2d[1] - y * 10))

                    p1 = (int(point1[0]), int(point1[1]))
                    p2 = (int(point2[0]), int(point2[1]))
                    p3 = (int(point3[0]), int(point3[1]))
                    p4 = (int(point4[0]), int(point4[1]))

                    #cv2.line(image, p1, p2, (255, 0, 0), 3)
                    cv2.circle(image, p1, 5, (255, 0, 0), 3)
                    cv2.circle(image, p2, 5, (255, 0, 0), 3)
                    cv2.circle(image, p3, 5, (255, 0, 0), 3)
                    cv2.circle(image, p4, 5, (255, 0, 0), 3)
                
                    # Add the text on the image
                    cv2.putText(image, "x: " + str(np.round(x,2)), (500, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                    cv2.putText(image, "y: " + str(np.round(y,2)), (500, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                    cv2.putText(image, "z: " + str(np.round(z,2)), (500, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

                mp_drawing.draw_landmarks(
                    image = image,
                    landmark_list = face_landmarks,
                    connections = mp_face_mesh.FACEMESH_CONTOURS,
                    landmark_drawing_spec = drawing_spec,
                    connection_drawing_spec = drawing_spec
                )
        else:
            break

        cv2.imshow('Head Pose Estimation', image)
        cv2.imshow('Cutout', tmp)
        cv2.imshow('Cutout', frame)
        writer.write(tmp)
        
        if cv2.waitKey(5) & 0xFF == 27:
            break
        if cv2.waitKey(5) & 0xFF == 8:
            chosen_landmark += 1



    video.release()
    writer.release()
    cv2.destroyAllWindows()
