import cv2
import numpy as np
from pathlib import Path
from ultralytics import YOLO

# --------------------
# 配置参数
# --------------------
IMAGE_PATH = "/home/hx/yolo/yemian/test_image/1.png"
MODEL_PATH = "best.pt"
OUTPUT_PATH = "./output/single_result.jpg"

TARGET_SIZE = 640

# 新增：用于计算像素到实际尺寸换算比例的函数
def calculate_pixel_to_real_ratio(real_length_mm, pixel_length):
    """
    计算像素到实际尺寸的换算比例。
    参数:
    - real_length_mm: 实际物理长度（例如：毫米）
    - pixel_length: 对应的实际物理长度在图像中的像素数
    返回:
    - PIXEL_TO_REAL_RATIO: 每个像素代表的实际物理长度（单位：mm/像素）
    """
    if pixel_length == 0:
        raise ValueError("像素长度不能为0")
    return real_length_mm / pixel_length


# 在主函数infer_single_image之前设置一个默认值
PIXEL_TO_REAL_RATIO = 1.0  # 默认值，之后会被真实计算的比例替换

# 假设我们知道某物体的真实长度是real_length_mm毫米，在图像中占pixel_length像素
real_length_mm = 100  # 物体的实际长度（单位：毫米）
pixel_length = 200  # 物体在图像中的像素长度

try:
    PIXEL_TO_REAL_RATIO = calculate_pixel_to_real_ratio(real_length_mm, pixel_length)
    print(f"换算比例已设定: {PIXEL_TO_REAL_RATIO:.4f} mm/像素")
except ValueError as e:
    print(e)

# 全局 ROI 定义：(x, y, w, h)
ROIS = [
    (859, 810, 696, 328),
]

# --------------------
# 辅助函数（保持不变）
# --------------------
def select_edge_corners(corners, w, left_ratio=0.2, right_ratio=0.2, y_var_thresh=5):
    if corners is None:
        return [], []
    corners = np.int32(corners).reshape(-1, 2)
    left_thresh = int(w * left_ratio)
    right_thresh = w - int(w * right_ratio)

    left_candidates = corners[corners[:, 0] <= left_thresh]
    right_candidates = corners[corners[:, 0] >= right_thresh]

    def filter_by_y_variation(pts):
        if len(pts) < 2:
            return pts
        pts_sorted = pts[np.argsort(pts[:, 1])]
        diffs = np.abs(np.diff(pts_sorted[:, 1]))
        keep_idx = np.where(diffs > y_var_thresh)[0]
        selected = [pts_sorted[i] for i in keep_idx] + [pts_sorted[i + 1] for i in keep_idx]
        return np.array(selected) if len(selected) > 0 else pts_sorted

    left_final = filter_by_y_variation(left_candidates)
    right_final = filter_by_y_variation(right_candidates)
    return left_final, right_final


def fit_line_with_outlier_removal(pts, dist_thresh=10):
    if pts is None or len(pts) < 2:
        return None, pts
    pts = np.array(pts)
    x, y = pts[:, 0], pts[:, 1]
    m, b = np.polyfit(y, x, 1)  # x = m*y + b
    x_fit = m * y + b
    dists = np.abs(x - x_fit)
    mask = dists < dist_thresh
    if mask.sum() < 2:
        return (m, b), pts
    m, b = np.polyfit(y[mask], x[mask], 1)
    inliers = np.stack([x[mask], y[mask]], axis=1)
    return (m, b), inliers


# --------------------
# 单图推理主函数
# --------------------
def infer_single_image(image_path, model_path, output_path):
    orig_img = cv2.imread(str(image_path))
    if orig_img is None:
        print(f"❌ 无法读取图像: {image_path}")
        return None

    overlay_img = orig_img.copy()
    x_diff_pixel = None  # 像素单位的差值

    model = YOLO(model_path)
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)

    for idx, (x, y, w, h) in enumerate(ROIS):
        roi_img = orig_img[y:y+h, x:x+w]
        resized_img = cv2.resize(roi_img, (TARGET_SIZE, TARGET_SIZE))

        results = model(source=resized_img, imgsz=TARGET_SIZE, verbose=False)
        result = results[0]

        if result.masks is None or len(result.masks.data) == 0:
            print("❌ 未检测到 mask")
            continue

        mask = result.masks.data[0].cpu().numpy()
        mask_bin = (mask > 0.5).astype(np.uint8)
        mask_bin = cv2.resize(mask_bin, (w, h), interpolation=cv2.INTER_NEAREST)

        color_mask = np.zeros_like(roi_img, dtype=np.uint8)
        color_mask[mask_bin == 1] = (0, 255, 0)
        overlay_img[y:y+h, x:x+w] = cv2.addWeighted(roi_img, 0.7, color_mask, 0.3, 0)

        mask_gray = (mask_bin * 255).astype(np.uint8)
        corners = cv2.goodFeaturesToTrack(mask_gray, maxCorners=200, qualityLevel=0.01, minDistance=5)

        left_pts, right_pts = select_edge_corners(corners, w)
        left_line, _ = fit_line_with_outlier_removal(left_pts)
        right_line, _ = fit_line_with_outlier_removal(right_pts)

        if left_line and right_line:
            y_ref = h * 0.6
            m1, b1 = left_line
            m2, b2 = right_line
            x1 = m1 * y_ref + b1
            x2 = m2 * y_ref + b2
            x_diff_pixel = abs(x2 - x1)

            # 绘制参考线和文字（仍用像素值显示）
            cv2.line(overlay_img[y:y+h, x:x+w], (0, int(y_ref)), (w, int(y_ref)), (0, 255, 255), 2)
            cv2.putText(overlay_img[y:y+h, x:x+w],
                        f"x_diff={x_diff_pixel:.1f}px",
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                        1, (0, 255, 255), 2)

            # 绘制左右拟合线
            for (m, b), color in [(left_line, (0, 0, 255)), (right_line, (255, 0, 0))]:
                y1, y2 = 0, h
                x1_line, x2_line = int(m * y1 + b), int(m * y2 + b)
                cv2.line(overlay_img[y:y+h, x:x+w], (x1_line, y1), (x2_line, y2), color, 3)

    cv2.imwrite(output_path, overlay_img)
    print(f"✅ 结果已保存至: {output_path}")

    if x_diff_pixel is not None:
        x_diff_real = x_diff_pixel * PIXEL_TO_REAL_RATIO
        print(f"📊 x差值（像素） = {x_diff_pixel:.2f} px")
        print(f"📏 x差值（实际） = {x_diff_real:.2f} mm")  # 可改为 cm 或其他单位
    else:
        print("⚠️ 未能计算 x 差值")

    return x_diff_pixel


# =====================
# 运行入口
# =====================
if __name__ == "__main__":
    infer_single_image(IMAGE_PATH, MODEL_PATH, OUTPUT_PATH)