import cv2
import numpy as np
from ultralytics import YOLO
import time
import os

# ====================== 用户配置 ======================
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/point/exp_pose2/weights/best.pt"
IMAGE_PATH = "/home/hx/yolo/output_masks/3.png"
SAVE_RESULT = True
SAVE_PATH = "/home/hx/yolo/test_image/result_with_angle2.jpg"
SHOW_IMAGE = True

# 关键点颜色：红、绿、蓝（对应 kpt 1, 2, 3）
KPT_COLORS = [
    (0, 0, 255),    # 红色 - kpt 1
    (0, 255, 0),    # 绿色 - kpt 2
    (255, 0, 0),    # 蓝色 - kpt 3
]

# ====================== 计算两向量夹角（0~180°）=======================
def calculate_angle_between_vectors(v1, v2):
    """
    计算两个向量之间的夹角（单位：度），范围 0 ~ 180°
    """
    v1 = np.array(v1)
    v2 = np.array(v2)

    if np.linalg.norm(v1) == 0 or np.linalg.norm(v2) == 0:
        return 0.0

    v1_u = v1 / np.linalg.norm(v1)
    v2_u = v2 / np.linalg.norm(v2)

    dot = np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)
    angle_rad = np.arccos(dot)
    angle_deg = np.degrees(angle_rad)
    return angle_deg

# ====================== 主程序 =======================
if __name__ == "__main__":
    print("正在加载模型...")
    try:
        model = YOLO(MODEL_PATH)
        print(f"✅ 模型加载完成: {MODEL_PATH}")
    except Exception as e:
        print(f"❌ 模型加载失败: {e}")
        exit(1)

    # 读取图像
    img = cv2.imread(IMAGE_PATH)
    if img is None:
        raise FileNotFoundError(f"❌ 无法加载图像: {IMAGE_PATH}")

    h, w = img.shape[:2]
    print(f"✅ 图像加载成功: {IMAGE_PATH} ({w}x{h})")

    # --- 推理 ---
    results = model(img, conf=0.25, imgsz=1280)
    result = results[0]

    # 存储每个实例的关键点和置信度
    instances = []

    if result.keypoints is not None and result.boxes is not None:
        kpts_list = result.keypoints.xy.cpu().numpy()  # (N, K, 2)
        confs = result.boxes.conf.cpu().numpy()        # (N,)
        bboxes = result.boxes.xyxy.cpu().numpy()       # (N, 4)

        # 组合：置信度 + 关键点 + 边框
        for i in range(len(kpts_list)):
            if len(kpts_list[i]) >= 2:  # 至少有两个关键点
                kpt1 = (int(kpts_list[i][0][0]), int(kpts_list[i][0][1]))
                kpt2 = (int(kpts_list[i][1][0]), int(kpts_list[i][1][1]))
                instances.append({
                    'conf': confs[i],
                    'kpt1': kpt1,
                    'kpt2': kpt2,
                    'bbox': bboxes[i]
                })

        # 按置信度降序排序
        instances.sort(key=lambda x: x['conf'], reverse=True)

        # 只保留置信度最高的两个实例
        top_two = instances[:2]

        if len(top_two) == 2:
            inst1, inst2 = top_two

            # 绘制第一个实例的关键点和连线
            cv2.circle(img, inst1['kpt1'], 10, KPT_COLORS[0], -1)
            cv2.circle(img, inst1['kpt2'], 10, KPT_COLORS[1], -1)
            cv2.line(img, inst1['kpt1'], inst1['kpt2'], (255, 255, 0), 3)  # 黄色线
            cv2.putText(img, "1", (inst1['kpt1'][0]+10, inst1['kpt1'][1]-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[0], 2)
            cv2.putText(img, "2", (inst1['kpt2'][0]+10, inst1['kpt2'][1]-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[1], 2)

            # 绘制第二个实例的关键点和连线
            cv2.circle(img, inst2['kpt1'], 10, KPT_COLORS[0], -1)
            cv2.circle(img, inst2['kpt2'], 10, KPT_COLORS[1], -1)
            cv2.line(img, inst2['kpt1'], inst2['kpt2'], (0, 255, 255), 3)  # 青色线
            cv2.putText(img, "1", (inst2['kpt1'][0]+10, inst2['kpt1'][1]-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[0], 2)
            cv2.putText(img, "2", (inst2['kpt2'][0]+10, inst2['kpt2'][1]-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[1], 2)

            # 构造两个向量：kpt2 - kpt1
            vec1 = np.array([inst1['kpt2'][0] - inst1['kpt1'][0],
                             inst1['kpt2'][1] - inst1['kpt1'][1]])

            vec2 = np.array([inst2['kpt2'][0] - inst2['kpt1'][0],
                             inst2['kpt2'][1] - inst2['kpt1'][1]])

            # 计算夹角
            angle = calculate_angle_between_vectors(vec1, vec2)
            print(f"✅ 置信度最高的两个实例，其关键点1-2连线夹角: {angle:.1f}°")

            # 在图像上显示角度
            cv2.putText(img, f"Angle: {angle:.1f}°", (50, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, lineType=cv2.LINE_AA)

        else:
            print(f"⚠️ 检测到 {len(top_two)} 个有效实例（需要至少2个）")
    else:
        print("⚠️ 未检测到关键点或边界框")

    # --- 保存结果 ---
    os.makedirs(os.path.dirname(SAVE_PATH), exist_ok=True)
    if SAVE_RESULT:
        cv2.imwrite(SAVE_PATH, img)
        print(f"✅ 结果图像已保存至: {SAVE_PATH}")

    # --- 显示图像 ---
    if SHOW_IMAGE:
        cv2.imshow("Keypoint Angle Between Two Best Instances", img)
        print("👉 按任意键关闭窗口...")
        cv2.waitKey(0)
        cv2.destroyAllWindows()