#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
离线测试脚本：直接复用 detect_circle 逻辑进行测试
运行环境：MaixPy (Sipeed MAIX)
"""
import sys
import os
# import time
from maix import image, time
import cv2
import numpy as np
import math

# ==================== 全局配置 (与 test_main.py 保持一致) ====================
REAL_RADIUS_CM = 20  # 靶心实际半径（厘米）

def detect_circle_v3(frame, laser_point=None, img_cv=None):
    """检测图像中的靶心（优先清晰轮廓，其次黄色区域）- 返回椭圆参数版本
    增加红色圆圈检测，验证黄色圆圈是否为真正的靶心
    如果提供 laser_point，会选择最接近激光点的目标
    优化：
      1. 缩图到 MAX_DET_DIM 后再做 HSV/形态学，最长边 640->320 可获得 ~4x 加速
      2. 红色掩码在黄色轮廓循环外只计算一次，避免 N 次重复计算
      3. img_cv 可由外部传入（与其他线程共享转换结果），为 None 时自动转换
    Args:
        frame: 图像帧（img_cv 为 None 时使用）
        laser_point: 激光点坐标 (x, y)，用于多目标场景下的目标选择
        img_cv: 已转换的 numpy BGR/RGB 图像；不为 None 时跳过 image2cv 转换
    Returns:
        (result_img, best_center, best_radius, method, best_radius1, ellipse_params)
    """
    if img_cv is None:
        img_cv = image.image2cv(frame, False, False)
    from datetime import datetime
    print(f"[detect_circle_v3] begin {datetime.now()}")
    # -- 1. 缩图加速（与三角形路径保持一致）
    h_orig, w_orig = img_cv.shape[:2]
    MAX_DET_DIM = 480
    long_side = max(h_orig, w_orig)
    if long_side > MAX_DET_DIM:
        det_scale = MAX_DET_DIM / long_side
        img_det = cv2.resize(img_cv, (int(w_orig * det_scale), int(h_orig * det_scale)),
                             interpolation=cv2.INTER_LINEAR)
        inv_scale = 1.0 / det_scale  # 检测坐标 -> 原始坐标的倍率
    else:
        img_det = img_cv
        inv_scale = 1.0

    # 激光点映射到检测分辨率
    lp_det = None
    if laser_point is not None:
        lp_det = (laser_point[0] / inv_scale, laser_point[1] / inv_scale)
    best_center = best_radius = best_radius1 = method = None
    ellipse_params = None

    print(f"[detect_circle_v3] step 1 fin {datetime.now()}")

    # -- 2. HSV + 黄色掩码
    hsv = cv2.cvtColor(img_det, cv2.COLOR_RGB2HSV)
    h, s, v = cv2.split(hsv)
    s = np.clip(s * 1.1, 0, 255).astype(np.uint8)
    hsv = cv2.merge((h, s, v))
    lower_yellow = np.array([7, 80, 0])
    upper_yellow = np.array([32, 255, 255])
    mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask_yellow = cv2.morphologyEx(mask_yellow, cv2.MORPH_CLOSE, kernel)

    print(f"[detect_circle_v3] step 2 fin {datetime.now()}")

    # -- 3. 红色掩码：在循环外只算一次
    mask_red = cv2.bitwise_or(
        cv2.inRange(hsv, np.array([0, 50, 40]), np.array([10, 255, 255])),
        cv2.inRange(hsv, np.array([170, 50, 40]), np.array([180, 255, 255])),
    )
    kernel_red = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask_red = cv2.morphologyEx(mask_red, cv2.MORPH_CLOSE, kernel_red)
    contours_red, _ = cv2.findContours(mask_red, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 预先把红色轮廓筛选成 (center, radius) 列表，后续直接查表
    red_candidates = []
    for cnt_r in contours_red:
        ar = cv2.contourArea(cnt_r)
        if ar <= 10:
            continue
        pr = cv2.arcLength(cnt_r, True)
        if pr <= 0 or (4 * np.pi * ar) / (pr * pr) <= 0.3:
            continue
        if len(cnt_r) >= 5:
            (xr, yr), (wr, hr), _ = cv2.fitEllipse(cnt_r)
            red_candidates.append({"center": (int(xr), int(yr)), "radius": int(min(wr, hr) / 2)})
        else:
            (xr, yr), rr = cv2.minEnclosingCircle(cnt_r)
            red_candidates.append({"center": (int(xr), int(yr)), "radius": int(rr)})

    print(f"[detect_circle_v3] step 3 fin {datetime.now()}")

    # -- 4. 黄色轮廓循环（复用上面的红色候选列表）
    contours_yellow, _ = cv2.findContours(mask_yellow, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    valid_targets = []
    for cnt_yellow in contours_yellow:
        area = cv2.contourArea(cnt_yellow)
        if area <= 15:
            continue
        perimeter = cv2.arcLength(cnt_yellow, True)
        if perimeter <= 0:
            continue
        circularity = (4 * np.pi * area) / (perimeter * perimeter)
        if circularity <= 0.5:
            continue
        print(f"[target] -> 面积:{area:.1f}, 圆度:{circularity:.2f}")
        if len(cnt_yellow) >= 5:
            (x, y), (width, height), angle = cv2.fitEllipse(cnt_yellow)
            yellow_ellipse = ((x, y), (width, height), angle)
            yellow_center = (int(x), int(y))
            yellow_radius = int(min(width, height) / 2)
        else:
            (x, y), radius = cv2.minEnclosingCircle(cnt_yellow)
            yellow_center = (int(x), int(y))
            yellow_radius = int(radius)
            yellow_ellipse = None
        # 在预筛好的红色候选中匹配
        matched = False
        for rc in red_candidates:
            ddx = yellow_center[0] - rc["center"][0]
            ddy = yellow_center[1] - rc["center"][1]
            dist_centers = math.hypot(ddx, ddy)
            if dist_centers < yellow_radius * 1.5 and rc["radius"] > yellow_radius * 0.7:
                print(f"[target] -> 找到匹配的红圈: 黄心({yellow_center}), "
                                f"红心({rc['center']}), 距离:{dist_centers:.1f}, "
                                f"黄半径:{yellow_radius}, 红半径:{rc['radius']}")
                valid_targets.append({
                    "center": yellow_center,
                    "radius": yellow_radius,
                    "ellipse": yellow_ellipse,
                    "area": area,
                })
                matched = True
                break
        if not matched :
            print("Debug -> 未找到匹配的红色圆圈，可能是误识别")

    print(f"[detect_circle_v3] step 4 fin {datetime.now()}")

    # -- 5. 选最佳目标，坐标还原到原始分辨率
    if valid_targets:
        if lp_det:
            best_target = min(valid_targets,
                              key=lambda t: (t["center"][0] - lp_det[0]) ** 2
                                            + (t["center"][1] - lp_det[1]) ** 2)
            method = "v3_ellipse_red_validated_laser_selected"
        else:
            best_target = max(valid_targets, key=lambda t: t["area"])
            method = "v3_ellipse_red_validated"
        bc = best_target["center"]
        br = best_target["radius"]
        be = best_target["ellipse"]
        if inv_scale != 1.0:
            best_center = (int(bc[0] * inv_scale), int(bc[1] * inv_scale))
            best_radius = int(br * inv_scale)
            if be is not None:
                (ex, ey), (ew, eh), ea = be
                be = ((ex * inv_scale, ey * inv_scale),
                      (ew * inv_scale, eh * inv_scale), ea)
        else:
            best_center = bc
            best_radius = br
        ellipse_params = be
        best_radius1 = best_radius * 5
    result_img = image.cv2image(img_cv, False, False)
    print(f"[detect_circle_v3] step 5 fin {datetime.now()}")
    return result_img, best_center, best_radius, method, best_radius1, ellipse_params


def run_offline_test(image_path):
    """读取图片，检测圆，绘制结果，保存图片"""

    # 1. 检查文件是否存在
    if not os.path.exists(image_path):
        print(f"[ERROR] 找不到图片文件: {image_path}")
        return

    # 2. 使用 maix.image 读取图片 (适配 MaixPy v4)
    try:
        # 使用 image.load 读取文件，返回 Image 对象
        img = image.load(image_path)
        print(f"[INFO] 成功读取图片: {image_path} (尺寸: {img.width()}x{img.height()})")
    except Exception as e:
        print(f"[ERROR] 读取图片失败: {e}")
        print("提示：请确认 MaixPy 版本是否为 v4，且图片路径正确。")
        return

    # 3. 调用 detect_circle_v3 函数
    print("[INFO] 正在调用 detect_circle_v3 进行检测...")
    start_time = time.ticks_ms()

    result_img, center, radius, method, radius1, ellipse_params = detect_circle_v3(img)

    cost_time = time.ticks_ms() - start_time
    print(f"[INFO] 检测完成，耗时: {cost_time}ms")
    print(f"       结果 -> 圆心: {center}, 半径: {radius}, 方法: {method}")
    if ellipse_params:
        (ell_center, (width, height), angle) = ellipse_params
        print(
            f"       椭圆 -> 中心: ({ell_center[0]:.1f}, {ell_center[1]:.1f}), 长轴: {max(width, height):.1f}, 短轴: {min(width, height):.1f}, 角度: {angle:.1f}°")

    # 4. 绘制辅助线（可选，用于调试）
    if center and radius:
        # 为了绘制椭圆，需要转换回 cv2 图像
        img_cv = image.image2cv(result_img, False, False)

        cx, cy = center

        # 如果有椭圆参数，绘制椭圆
        if ellipse_params:
            (ell_center, (width, height), angle) = ellipse_params
            cx_ell, cy_ell = int(ell_center[0]), int(ell_center[1])

            # 确定长轴和短轴
            if width >= height:
                # width 是长轴，height 是短轴
                axes_major = width
                axes_minor = height
                major_angle = angle  # 长轴角度就是 angle
                minor_angle = angle + 90  # 短轴角度 = 长轴角度 + 90度
            else:
                # height 是长轴，width 是短轴
                axes_major = height
                axes_minor = width
                major_angle = angle + 90  # 长轴角度 = width角度 + 90度
                minor_angle = angle  # 短轴角度就是 angle

            # 使用 OpenCV 绘制椭圆（绿色，线宽2）
            cv2.ellipse(img_cv,
                        (cx_ell, cy_ell),  # 中心点
                        (int(width / 2), int(height / 2)),  # 半宽、半高
                        angle,  # 旋转角度（OpenCV需要原始angle）
                        0, 360,  # 起始和结束角度
                        (0, 255, 0),  # 绿色 (RGB格式)
                        2)  # 线宽

            # 绘制椭圆中心点（红色）
            cv2.circle(img_cv, (cx_ell, cy_ell), 3, (255, 0, 0), -1)

            import math
            # 绘制短轴（蓝色线条）
            minor_length = axes_minor / 2
            minor_angle_rad = math.radians(minor_angle)
            dx_minor = minor_length * math.cos(minor_angle_rad)
            dy_minor = minor_length * math.sin(minor_angle_rad)
            pt1_minor = (int(cx_ell - dx_minor), int(cy_ell - dy_minor))
            pt2_minor = (int(cx_ell + dx_minor), int(cy_ell + dy_minor))
            cv2.line(img_cv, pt1_minor, pt2_minor, (0, 0, 255), 2)  # 蓝色 (RGB格式)
        else:
            # 如果没有椭圆参数，绘制圆形（红色）
            cv2.circle(img_cv, (cx, cy), radius, (0, 0, 255), 2)
            cv2.circle(img_cv, (cx, cy), 2, (0, 0, 255), -1)

        # 转换回 maix image
        result_img = image.cv2image(img_cv, False, False)

        # 定义颜色对象用于文字
        try:
            color_black = image.Color.from_rgb(0, 0, 0)
        except AttributeError:
            color_black = image.Color(0, 0, 0)

        # D. 添加文字信息
        FOCAL_LENGTH_PIX = 1900
        d = (REAL_RADIUS_CM * FOCAL_LENGTH_PIX) / radius1 / 100.0
        info_str = f"R:{radius} M:{method} D:{d:.2f}"
        print(info_str)

        # 计算文字位置，防止超出图片边界
        r_outer = int(radius * 11.0) if radius else 100
        text_y = cy - r_outer - 20 if cy > r_outer + 20 else cy + r_outer + 20

        # 调用 draw_string
        result_img.draw_string(0, 0, info_str, color=color_black, scale=1.0)

    # 5. 保存结果图片
    base, ext = os.path.splitext(image_path)
    output_path = f"{base}_result{ext}"
    try:
        result_img.save(output_path, quality=100)
        print(f"[SUCCESS] 结果已保存至: {output_path}")
    except Exception as e:
        print(f"[ERROR] 保存图片失败: {e}")


if __name__ == "__main__":
    # ================= 配置区域 =================

    # 1. 设置要测试的图片路径
    # 建议将图片放在与脚本同级目录，或者使用绝对路径
    TARGET_IMAGE = "/root/phot/None_314_258_0_0041.bmp"

    TARGET_DIR = "/root/phot"  # 修改为你想要读取的目录路径

    # 支持的图片格式
    IMAGE_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.bmp']

    # ================= 执行区域 =================
    if 'TARGET_DIR' in locals():
        # 读取目录下所有图片文件，过滤掉 _result.jpg 后缀的文件
        image_files = []
        if os.path.exists(TARGET_DIR) and os.path.isdir(TARGET_DIR):
            for filename in os.listdir(TARGET_DIR):
                # 检查文件扩展名
                if any(filename.lower().endswith(ext) for ext in IMAGE_EXTENSIONS):
                    # 过滤掉 _result.jpg 后缀的文件
                    if not filename.endswith('_result.jpg'):
                        filepath = os.path.join(TARGET_DIR, filename)
                        if os.path.isfile(filepath):
                            image_files.append(filepath)

            # 按文件名排序（可选）
            image_files.sort()

            print(f"[INFO] 在目录 {TARGET_DIR} 中找到 {len(image_files)} 张图片")

            # 处理每张图片
            for img_path in image_files:
                print(f"\n{'=' * 10} 开始处理: {img_path} {'=' * 10}")
                run_offline_test(img_path)
        else:
            print(f"[ERROR] 目录不存在或不是有效目录: {TARGET_DIR}")

    else:
        run_offline_test(TARGET_IMAGE)