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

IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}


def draw_direction(img, cx, cy, angle_deg, length=80, color=(0, 255, 0)):
    """画主方向箭头"""
    theta = np.radians(angle_deg)
    x2 = int(cx + length * np.cos(theta))
    y2 = int(cy + length * np.sin(theta))
    cv2.arrowedLine(
        img,
        (int(cx), int(cy)),
        (x2, y2),
        color,
        2,
        tipLength=0.2
    )


def process_obb_images(
    model_path,
    image_dir,
    output_dir="./inference_results",
    conf_thresh=0.15,
    imgsz=640
):
    os.makedirs(output_dir, exist_ok=True)
    results_dict = {}

    print("加载 YOLO 模型...")
    model = YOLO(model_path)
    print("✅ 模型加载完成")

    image_files = [
        f for f in os.listdir(image_dir)
        if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS
    ]

    if not image_files:
        print(f"❌ 未找到图像文件：{image_dir}")
        return results_dict

    print(f"发现 {len(image_files)} 张图像待处理")

    for img_filename in image_files:
        img_path = os.path.join(image_dir, img_filename)
        print(f"\n正在处理：{img_filename}")

        img = cv2.imread(img_path)
        if img is None:
            print("❌ 读取失败，跳过")
            continue

        # ---------- OBB 推理 ----------
        results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode="obb")
        result = results[0]
        annotated_img = result.plot()

        boxes = result.obb
        angles_deg = []
        centers = []

        if boxes is None or len(boxes) == 0:
            print("❌ 未检测到目标")
        else:
            for i, box in enumerate(boxes):
                cx, cy, w, h, r_rad = box.xywhr.cpu().numpy()[0]

                direction = r_rad if w >= h else r_rad + np.pi / 2
                direction = direction % np.pi
                angle_deg = np.degrees(direction)

                angles_deg.append(angle_deg)
                centers.append((int(cx), int(cy)))

                print(f"  Box {i + 1} 主方向: {angle_deg:.2f} deg")

                # 主方向可视化
                draw_direction(annotated_img, cx, cy, angle_deg)
                cv2.circle(annotated_img, (int(cx), int(cy)), 4, (0, 0, 255), -1)

        # ---------- 两两夹角 ----------
        pairwise_angles_deg = []

        if len(angles_deg) >= 2:
            for i in range(len(angles_deg)):
                for j in range(i + 1, len(angles_deg)):
                    diff_rad = abs(
                        np.radians(angles_deg[i]) -
                        np.radians(angles_deg[j])
                    )
                    min_diff_rad = min(diff_rad, np.pi - diff_rad)
                    angle_ij = np.degrees(min_diff_rad)

                    pairwise_angles_deg.append(angle_ij)
                    print(
                        f"  Box {i + 1} 与 Box {j + 1} 夹角: {angle_ij:.2f} deg"
                    )

        # ---------- 右上角粗体 angle ----------
        if pairwise_angles_deg:
            max_angle = max(pairwise_angles_deg)
            h, w = annotated_img.shape[:2]

            text = f"angle: {max_angle:.1f} deg"

            # 粗体效果（多次叠加）
            for dx, dy in [(0, 0), (1, 0), (0, 1), (1, 1)]:
                cv2.putText(
                    annotated_img,
                    text,
                    (w - 300 + dx, 40 + dy),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1.2,
                    (0, 0, 255),
                    3
                )

        # ---------- 保存 ----------
        save_path = os.path.join(output_dir, "detected_" + img_filename)
        cv2.imwrite(save_path, annotated_img)
        print(f"✅ 保存完成: {save_path}")

        results_dict[img_filename] = {
            "angles_deg": angles_deg,
            "pairwise_angles_deg": pairwise_angles_deg
        }

    print("\n🎉 全部处理完成")
    return results_dict


# ------------------- 主入口 -------------------
if __name__ == "__main__":
    MODEL_PATH = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new3/weights/best.pt"
    IMAGE_SOURCE_DIR = r"/home/hx/yolo/angle_base_obb/test_image"
    OUTPUT_DIR = "./inference_results"

    results = process_obb_images(
        MODEL_PATH,
        IMAGE_SOURCE_DIR,
        OUTPUT_DIR
    )

    for img_name, info in results.items():
        print(f"\n{img_name}")
        print("主方向角:", info["angles_deg"])
        print("两两夹角:", info["pairwise_angles_deg"])