import cv2
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

YAW_THRESHOLD = 13
PITCH_THRESHOLD = 5

def get_aoi(yaw, pitch):
    ret = ''
    if yaw > YAW_THRESHOLD:
        ret += 'Left '
    elif yaw > -YAW_THRESHOLD:
        ret += 'Center '
    else:
        ret += 'Right '

    if pitch > PITCH_THRESHOLD:
        ret += 'Lower'
    elif pitch > -PITCH_THRESHOLD:
        ret += 'Center'
    else:
        ret += 'Upper'

    return ret

areas = ['left upper', 'center', 'right upper', 'left lower', 'right lower']

# 1. クラスのセットアップ (FaceLandmarkerに変更)
BaseOptions = python.BaseOptions
FaceLandmarker = vision.FaceLandmarker
FaceLandmarkerOptions = vision.FaceLandmarkerOptions
VisionRunningMode = vision.RunningMode

# 2. 顔ランドマーカーの設定
# 変換行列（顔の向きデータ）を出力するようにフラグを立てます
options = FaceLandmarkerOptions(
    base_options=BaseOptions(model_asset_path='face_landmarker.task'),
    running_mode=VisionRunningMode.VIDEO,
    output_facial_transformation_matrixes=True,  # 顔の向き取得に必須
    num_faces=1  # 検出する顔の最大数
)

cap = cv2.VideoCapture(0)

# FaceLandmarkerのインスタンスを生成
with FaceLandmarker.create_from_options(options) as landmarker:
    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame_timestamp_ms = int(cap.get(cv2.CAP_PROP_POS_MSEC))

        # 稀にタイムスタンプが0や重複になるのを防ぐ安全策
        if frame_timestamp_ms == 0:
            frame_timestamp_ms = 1

        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)

        # 3. 検出の実行
        result = landmarker.detect_for_video(mp_image, frame_timestamp_ms)

        # 4. 変換行列から顔の向き（オイラー角）を計算
        if result.facial_transformation_matrixes:
            for matrix in result.facial_transformation_matrixes:
                # 4x4の変換行列から、左上の3x3の回転行列のみを抽出
                rot_mat = matrix[:3, :3]

                # OpenCVを使って回転行列をオイラー角（度数法）に分解
                angles, _, _, _, _, _ = cv2.RQDecomp3x3(rot_mat)

                # ピッチ(上下), ヨー(左右), ロール(傾き)の取得
                # ※カメラや座標系によって軸の順番は変わる場合があります
                pitch = angles[0]
                yaw = angles[1]
                roll = angles[2]

                # オイラー角からAOIを推定
                estimate_aoi = get_aoi(yaw, pitch)
                text_aoi = f"Estimated AOI: {estimate_aoi}"

                # 画面上に角度を描画
                text_pitch = f"Pitch (Up/Down): {pitch:.1f}"
                text_yaw = f"Yaw (Left/Right): {yaw:.1f}"
                text_roll = f"Roll (Tilt): {roll:.1f}"

                cv2.putText(frame, text_aoi, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

                cv2.putText(frame, text_pitch, (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
                cv2.putText(frame, text_yaw, (20, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
                cv2.putText(frame, text_roll, (20, 130), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)

        cv2.imshow('Area of Interest Estimation', frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

cap.release()
cv2.destroyAllWindows()