# pose_grid_overlay.py

import cv2
import mediapipe as mp
import numpy as np
import os

def draw_pose_on_grid(image_path: str, grid_size=(3, 3), save_path="outputs/pose_grid.jpg"):
    # 1. 이미지 불러오기
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"이미지를 불러올 수 없습니다: {image_path}")
    
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, _ = image.shape

    # 2. 저장 경로 디렉토리 보장
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    # 3. 포즈 추론
    with mp.solutions.pose.Pose(static_image_mode=True, model_complexity=2) as pose_model:
        results = pose_model.process(image_rgb)

        # 4. 격자 그리기
        for i in range(1, grid_size[0]):
            x = int(w * i / grid_size[0])
            cv2.line(image, (x, 0), (x, h), (200, 200, 200), 1)
        for i in range(1, grid_size[1]):
            y = int(h * i / grid_size[1])
            cv2.line(image, (0, y), (w, y), (200, 200, 200), 1)

        # 5. 포즈 keypoints 그리기
        if results.pose_landmarks:
            for lm in results.pose_landmarks.landmark:
                cx, cy = int(lm.x * w), int(lm.y * h)
                if 0 <= cx < w and 0 <= cy < h:  # 이미지 영역 내부만 그리기
                    cv2.circle(image, (cx, cy), 5, (255, 0, 0), -1)
        else:
            raise RuntimeError("❌ 포즈를 감지할 수 없습니다. 인물이 뚜렷하게 나온 이미지를 사용해보세요.")

    # 6. 저장
    cv2.imwrite(save_path, image)