import cv2
import mediapipe as mp
import numpy as np
import sys

NOSE_TIP = 1
CHIN = 152
FOREHEAD = 10
LEFT_EYE_OUTER = 263
RIGHT_EYE_OUTER = 33
LEFT_MOUTH = 61
RIGHT_MOUTH = 291
TOP_LIP = 13
BOTTOM_LIP = 14
LEFT_EYE_TOP = 159
LEFT_EYE_BOTTOM = 145
RIGHT_EYE_TOP = 386
RIGHT_EYE_BOTTOM = 374

FACE_3D_MODEL = np.array([
    [0.0, 0.0, 0.0],
    [0.0, -330.0, -65.0],
    [-225.0, 170.0, -135.0],
    [225.0, 170.0, -135.0],
    [-150.0, -150.0, -125.0],
    [150.0, -150.0, -125.0],
], dtype=np.float64)


def dessiner_landmarks(frame, landmarks, frame_w, frame_h):
    for point in landmarks.landmark:
        x = int(point.x * frame_w)
        y = int(point.y * frame_h)
        cv2.circle(frame, (x, y), 1, (0, 200, 100), -1)


def dessiner_axe_tete(frame, landmarks, frame_w, frame_h):
    image_points = np.array([
        [landmarks[NOSE_TIP].x * frame_w, landmarks[NOSE_TIP].y * frame_h],
        [landmarks[CHIN].x * frame_w, landmarks[CHIN].y * frame_h],
        [landmarks[LEFT_EYE_OUTER].x * frame_w, landmarks[LEFT_EYE_OUTER].y * frame_h],
        [landmarks[RIGHT_EYE_OUTER].x * frame_w, landmarks[RIGHT_EYE_OUTER].y * frame_h],
        [landmarks[LEFT_MOUTH].x * frame_w, landmarks[LEFT_MOUTH].y * frame_h],
        [landmarks[RIGHT_MOUTH].x * frame_w, landmarks[RIGHT_MOUTH].y * frame_h],
    ], dtype=np.float64)

    focal_length = frame_w
    camera_matrix = np.array([
        [focal_length, 0, frame_w / 2],
        [0, focal_length, frame_h / 2],
        [0, 0, 1]
    ], dtype=np.float64)

    success, rotation_vec, translation_vec = cv2.solvePnP(
        FACE_3D_MODEL, image_points, camera_matrix, np.zeros((4, 1))
    )
    if not success:
        return None, None, None

    axis_length = 80
    axes_3d = np.float32([
        [axis_length, 0, 0],
        [0, axis_length, 0],
        [0, 0, axis_length]
    ])
    axes_2d, _ = cv2.projectPoints(axes_3d, rotation_vec, translation_vec, camera_matrix, np.zeros((4, 1)))
    nose = (int(image_points[0][0]), int(image_points[0][1]))

    cv2.line(frame, nose, tuple(axes_2d[0].ravel().astype(int)), (0, 0, 255), 2)
    cv2.line(frame, nose, tuple(axes_2d[1].ravel().astype(int)), (0, 255, 0), 2)
    cv2.line(frame, nose, tuple(axes_2d[2].ravel().astype(int)), (255, 0, 0), 2)

    rotation_mat, _ = cv2.Rodrigues(rotation_vec)
    angles, _, _, _, _, _ = cv2.RQDecomp3x3(rotation_mat)
    return angles[0], angles[1], angles[2]


def dessiner_signaux(frame, lm, pitch, yaw, roll):
    face_height = abs(lm[CHIN].y - lm[FOREHEAD].y)
    face_width = abs(lm[LEFT_EYE_OUTER].x - lm[RIGHT_EYE_OUTER].x)

    mouth_opening = abs(lm[TOP_LIP].y - lm[BOTTOM_LIP].y) / face_height if face_height > 0 else 0
    left_eye = abs(lm[LEFT_EYE_TOP].y - lm[LEFT_EYE_BOTTOM].y) / face_height if face_height > 0 else 0
    right_eye = abs(lm[RIGHT_EYE_TOP].y - lm[RIGHT_EYE_BOTTOM].y) / face_height if face_height > 0 else 0
    smile = abs(lm[LEFT_MOUTH].x - lm[RIGHT_MOUTH].x) / face_width if face_width > 0 else 0

    lignes = [
        f"Pitch : {pitch:.1f}",
        f"Yaw   : {yaw:.1f}",
        f"Roll  : {roll:.1f}",
        f"Bouche: {mouth_opening:.3f}",
        f"Oeil G: {left_eye:.3f}",
        f"Oeil D: {right_eye:.3f}",
        f"Sourire: {smile:.3f}",
    ]

    overlay = frame.copy()
    cv2.rectangle(overlay, (8, 8), (200, 18 + len(lignes) * 20), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.5, frame, 0.5, 0, frame)

    for i, texte in enumerate(lignes):
        cv2.putText(frame, texte, (12, 24 + i * 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 180), 1)


def visualiser_live(chemin_video):
    cap = cv2.VideoCapture(chemin_video)
    if not cap.isOpened():
        raise ValueError(f"Impossible d'ouvrir : {chemin_video}")

    frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    wait_ms = max(1, int(1000 / fps))

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

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

            if results.multi_face_landmarks:
                face_landmarks = results.multi_face_landmarks[0]
                lm = face_landmarks.landmark

                dessiner_landmarks(frame, face_landmarks, frame_w, frame_h)
                pitch, yaw, roll = dessiner_axe_tete(frame, lm, frame_w, frame_h)

                if pitch is not None:
                    dessiner_signaux(frame, lm, pitch, yaw, roll)

            cv2.imshow("Visualisation faciale - Q pour quitter", frame)
            if cv2.waitKey(wait_ms) & 0xFF == ord('q'):
                break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    chemin = sys.argv[1] if len(sys.argv) > 1 else "test.mp4"
    visualiser_live(chemin)