import csv
from datetime import datetime

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

# データ保存用バッファ [(対象領域, Yaw値, Pitch値, Roll値)]
data_log = []
check_areas = ['left upper', 'center', 'right upper', 'left lower', 'right lower']
yaw, pitch, roll = 0, 0, 0
current_area_index = 0

# 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)

        now_area = check_areas[current_area_index]

        text_now_area = f"Checking Area: {now_area}"
        cv2.putText(frame, text_now_area, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 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]

                # 画面上に角度を描画
                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_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 Record', frame)

        current_yaw = yaw if yaw is not None else 0.0
        current_pitch = pitch if pitch is not None else 0.0
        current_roll = roll if roll is not None else 0.0

        # 'r'が押されたら記録し，次の領域へ
        if cv2.waitKey(1) & 0xFF == ord('r'):
            print(f"Checked Area: {now_area}")
            data_log.append([
                now_area,
                round(current_yaw, 2),
                round(current_pitch, 2),
                round(current_roll, 2)
            ])
            current_area_index += 1

        if current_area_index >= len(check_areas):
            print('All Area Checked.')
            break

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

cap.release()
cv2.destroyAllWindows()

filename = f"AOI_record_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
with open(filename, 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(["Area", "Yaw", "Pitch", "Roll"])
    writer.writerows(data_log)