import argparse
import os
import re

import cv2
import numpy as np

# mmdet / mmpose
from mmdet.apis import inference_detector, init_detector
from mmpose.apis import inference_topdown
from mmpose.apis import init_model as init_pose_estimator
from mmpose.evaluation.functional import nms
from mmpose.structures import merge_data_samples
from mmpose.utils import adapt_mmdet_pipeline
from PIL import Image, ImageDraw


def extract_keypoints_rtmpose(pose_results):
    """
    RTMposeの推論結果 (pose_results) から、もっとも平均可視スコアの高い
    インスタンスを選んで keypoints([17,2]) を返す。検出できなければ None。
    """
    if not pose_results:
        return None

    max_avg_visible = 0
    best_instance = None
    for result in pose_results:
        pred_instances = result.pred_instances
        for instance in pred_instances:
            avg_visible = np.mean(instance.keypoints_visible)
            if avg_visible > max_avg_visible:
                max_avg_visible = avg_visible
                best_instance = instance

    if best_instance is None:
        return None

    # best_instance.keypoints の shape が (1, num_kpts, 2) の場合がある
    return best_instance.keypoints[0]


def pillow_draw_circle(draw, center, radius, fill=None, outline=None, width=1):
    """(center=(x,y)) を中心とする円を描画。"""
    x, y = int(center[0]), int(center[1])
    left_up = (x - radius, y - radius)
    right_down = (x + radius, y + radius)
    draw.ellipse([left_up, right_down], fill=fill, outline=outline, width=width)


def draw_glow_marker(draw, center, main_color, radius=15):
    """
    main_color で塗りつぶし、白枠を付けて光っているように描画
    """
    x, y = int(center[0]), int(center[1])
    # 白枠 (外周3ピクセル分)
    outer_radius = radius + 3
    pillow_draw_circle(
        draw, (x, y), outer_radius, fill=None, outline=(255, 255, 255), width=2
    )
    # 内部を塗りつぶす
    pillow_draw_circle(draw, (x, y), radius, fill=main_color, outline=None, width=0)


def main():
    parser = argparse.ArgumentParser(description="RTMpose: detect shoulders & hips.")
    parser.add_argument(
        "--img_dir",
        type=str,
        required=True,
        help="入力画像フォルダのパス（PNG/JPG 画像が格納されている）",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="pose_out",
        help="肩・腰描画後の画像を保存するフォルダ",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="推論に使用するデバイス (cuda / cpu など)",
    )

    # ここは「以下のソースコードの設定のままでいいです」とあった部分をそのまま利用
    det_config = "modules/rtmpose/mmdetection_cfg/rtmdet_m_640-8xb32_coco-person.py"
    det_checkpoint = (
        "models/rtmpose/rtmdet_m_8xb32-100e_coco-obj365-person-235e8209.pth"
    )
    pose_config = (
        "modules/rtmpose/configs/body_2d_keypoint/rtmpose/body8/"
        "rtmpose-l_8xb256-420e_body8-256x192.py"
    )
    pose_checkpoint = "models/rtmpose/rtmpose-l_simcc-aic-coco_pt-aic-coco_420e-256x192-f016ffe0_20230126.pth"

    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # (A) 人体検出器 (RTDet) の初期化
    print("[INFO] Initializing RTDet with config:")
    print("  det_config:", det_config)
    print("  det_checkpoint:", det_checkpoint)
    detector = init_detector(det_config, det_checkpoint, device=args.device)
    detector.cfg = adapt_mmdet_pipeline(detector.cfg)

    # (B) RTMpose (姿勢推定) の初期化
    print("[INFO] Initializing RTMpose with config:")
    print("  pose_config:", pose_config)
    print("  pose_checkpoint:", pose_checkpoint)
    pose_estimator = init_pose_estimator(
        pose_config, pose_checkpoint, device=args.device
    )

    # 画像ファイル一覧
    images = sorted(
        [
            f
            for f in os.listdir(args.img_dir)
            if f.lower().endswith((".png", ".jpg", ".jpeg"))
        ],
        key=lambda x: int(re.search(r"(\d+)", x).group())
        if re.search(r"(\d+)", x)
        else x,
    )
    if not images:
        print(f"[WARNING] 指定フォルダに画像が見つかりません: {args.img_dir}")
        return

    # 元コードで「肩・腰」に使われていた色 (pose_color_rgb = (33, 95, 154))
    pose_color_rgb = (33, 95, 154)

    for img_name in images:
        img_path = os.path.join(args.img_dir, img_name)
        frame_bgr = cv2.imread(img_path)
        if frame_bgr is None:
            print(f"[WARNING] 画像の読み込みに失敗: {img_path}")
            continue

        # (1) 物体検出
        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        det_result = inference_detector(detector, frame_rgb)
        pred_instance = det_result.pred_instances.cpu().numpy()

        # label=0 (person), スコア>0.3 の bboxをフィルタ
        bboxes = np.concatenate(
            (pred_instance.bboxes, pred_instance.scores[:, None]), axis=1
        )
        person_bboxes = bboxes[
            np.logical_and(pred_instance.labels == 0, pred_instance.scores > 0.3)
        ]

        # NMS (バージョンによって引数形式が違う場合あり)
        keep_idx = nms(person_bboxes, 0.3)
        person_bboxes = person_bboxes[keep_idx, :4]

        if len(person_bboxes) == 0:
            # 人物が検出されなければ何も描画せず保存
            out_path = os.path.join(args.output_dir, img_name)
            cv2.imwrite(out_path, frame_bgr)
            continue

        # (2) RTMpose で姿勢推定
        pose_results = inference_topdown(pose_estimator, frame_rgb, person_bboxes)
        data_samples = merge_data_samples(pose_results)

        # (3) 肩・腰 キーポイントだけを取り出す
        keypoints = extract_keypoints_rtmpose(pose_results)
        if keypoints is None:
            out_path = os.path.join(args.output_dir, img_name)
            cv2.imwrite(out_path, frame_bgr)
            continue

        # index: left_shoulder=5, right_shoulder=6, left_hip=11, right_hip=12 (COCO)
        left_shoulder = keypoints[5]
        right_shoulder = keypoints[6]
        left_hip = keypoints[11]
        right_hip = keypoints[12]

        # (4) 肩と腰を元画像に描画 (Pillow)
        pil_img = Image.fromarray(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_img)

        # 任意の可視性チェックなどは省略。ここでは座標がある4点を描画。
        for pt in [left_shoulder, right_shoulder, left_hip, right_hip]:
            draw_glow_marker(draw, (pt[0], pt[1]), pose_color_rgb, radius=15)

        # Pillow → BGR
        out_img_rgb = np.array(pil_img)
        out_img_bgr = cv2.cvtColor(out_img_rgb, cv2.COLOR_RGB2BGR)

        # (5) 保存
        out_path = os.path.join(args.output_dir, img_name)
        cv2.imwrite(out_path, out_img_bgr)

    print("All done. Results saved to:", args.output_dir)


if __name__ == "__main__":
    main()