import cv2
import math
import random
import numpy as np
import mediapipe as mp

from scipy.spatial.distance import euclidean as dist
from sklearn.preprocessing import LabelEncoder, MinMaxScaler

# feature definitions


DIMS = (224,224,3) # dimensions of the image
RIGHT = [[33, 133], [160, 144], [159, 145], [158, 153]] # right eye landmark positions
LEFT = [[263, 362], [387, 373], [386, 374], [385, 380]] # left eye landmark positions
MOUTH = [[61, 291], [39, 181], [0, 17], [269, 405]] # mouth landmark coordinates

EYE_AR_THRESH = 0.45
PROB_THRESH = 0.3
EYE_AR_CONSEC_FRAMES = 15

MOUTH_AR_THRESH = 0.33
MOUTH_AR_CONSEC_FRAMES = 20

MP_FACE_DETECTION = mp.solutions.face_detection
MP_DRAWING = mp.solutions.drawing_utils
MP_DRAWING_STYLES = mp.solutions.drawing_styles
MP_FACE_MESH = mp.solutions.face_mesh
DRAWING_SPEC = MP_DRAWING.DrawingSpec(thickness=1, circle_radius=1)

def get_ear(landmarks,eye):
    ''' Calculate the ratio of the eye length to eye width. 
    :param landmarks: Face Landmarks returned from FaceMesh MediaPipe model
    :param eye: List containing positions which correspond to the eye
    :return: Eye aspect ratio value
    '''
    N1 = dist(landmarks[eye[1][0]], landmarks[eye[1][1]])
    N2 = dist(landmarks[eye[2][0]], landmarks[eye[2][1]])
    N3 = dist(landmarks[eye[3][0]], landmarks[eye[3][1]])
    D = dist(landmarks[eye[0][0]], landmarks[eye[0][1]])
    return (N1 + N2 + N3) / (3 * D)

def get_eye_feature(landmarks):
    ''' Calculate the eye feature as the average of the eye aspect ratio for the two eyes
    :param landmarks: Face Landmarks returned from FaceMesh MediaPipe model
    :return: Eye feature value
    '''
    return (get_ear(landmarks,LEFT) + get_ear(landmarks,RIGHT))

def get_mouth_feature(landmarks):
    ''' Calculate mouth feature as the ratio of the mouth length to mouth width
    :param landmarks: Face Landmarks returned from FaceMesh MediaPipe model
    :return: Mouth feature value
    '''
    n_1 = dist(landmarks[MOUTH[1][0]], landmarks[MOUTH[1][1]])
    n_2 = dist(landmarks[MOUTH[2][0]], landmarks[MOUTH[2][1]])
    n_3 = dist(landmarks[MOUTH[3][0]], landmarks[MOUTH[3][1]])
    dst = dist(landmarks[MOUTH[0][0]], landmarks[MOUTH[0][1]])
    return (n_1 + n_2 + n_3)/(3*dst)

# image processing


def process_mp_img(frame):
    """
    returns features and/or processed image
    """
    with MP_FACE_MESH.FaceMesh(
        min_detection_confidence=0.3,
        min_tracking_confidence=0.8) as face_mesh:
        # convert the img to RGB and process it with MediaPipe Face Detection
        results = face_mesh.process(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB))

        if results.multi_face_landmarks is not None:
            landmark_pos = []
            for i, data in enumerate(results.multi_face_landmarks[0].landmark):
                landmark_pos.append([data.x, data.y, data.z])
            landmark_pos = np.array(landmark_pos)

            # draw face detections of each face
            annotated_img = frame.copy()
            for face_landmarks in results.multi_face_landmarks:
                # Calculate eye and mouth features
                eye_feature = get_eye_feature(landmark_pos)
                mouth_feature = get_mouth_feature(landmark_pos)

                # Binary classification: drowsy (1) or non-drowsy (0)
                drowsy = (eye_feature <= EYE_AR_THRESH) or (mouth_feature > MOUTH_AR_THRESH)       
                # face mesh
                MP_DRAWING.draw_landmarks(
                    image=annotated_img,
                    landmark_list=face_landmarks,
                    connections=MP_FACE_MESH.FACEMESH_TESSELATION,
                    landmark_drawing_spec=None,
                    connection_drawing_spec=MP_DRAWING_STYLES
                    .get_default_face_mesh_tesselation_style()
                )
                # eyes and mouth regions
                MP_DRAWING.draw_landmarks(
                    image=annotated_img,
                    landmark_list=face_landmarks,
                    connections=MP_FACE_MESH.FACEMESH_CONTOURS,
                    landmark_drawing_spec=None,
                    connection_drawing_spec=MP_DRAWING_STYLES
                    .get_default_face_mesh_contours_style()
                )
    return annotated_img, eye_feature, mouth_feature, drowsy



def mediapipe_process(frames):
    """
    Process all videos using MediaPipe and returns a 
    dictionary with the eye and mouth features in 
    the format {frame_number: {"eye_feature":0, "mouth_feature":0, "drowsy":0}}
    """
    mp_features = {}
    eye_features_all = []
    mouth_features_all = []
    # Extract eye and mouth features for all videos
    for frame in frames:
        mp_features[frame] = {"eye_feature": 0, "mouth_feature": 0, "drowsy": 0}
        _,eye_feature,mouth_feature,drowsy = process_mp_img(frame)
        mp_features[frame]["eye_feature"] = eye_feature
        mp_features[frame]["mouth_feature"] = mouth_feature
        mp_features[frame]["drowsy"] = drowsy
        eye_features_all.append(eye_feature)
        mouth_features_all.append(mouth_feature)

    # Calculate mean and standard deviation for normalization
    eye_mean, eye_std = np.mean(eye_features_all), np.std(eye_features_all)
    mouth_mean, mouth_std = np.mean(mouth_features_all), np.std(mouth_features_all)

    # Normalize eye and mouth features for all videos
    for frame,features in mp_features.items():
        features["eye_feature"] = (features["eye_feature"] - eye_mean) / eye_std
        features[frame]["mouth_feature"] = (features["mouth_feature"] - mouth_mean) / mouth_std

    return mp_features