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

TARGET_SIZE = 640  # 模型输入尺寸

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

# --------------------
# 根据角点分布，选取左右边缘角点
# --------------------
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)
    x_min, x_max = 0, w

    left_thresh = x_min + int(w * left_ratio)
    right_thresh = x_max - int(w * right_ratio)

    # 左右候选角点
    left_candidates = corners[corners[:, 0] <= left_thresh]
    right_candidates = corners[corners[:, 0] >= right_thresh]

    # --------------------
    # 进一步按 y 变化筛选
    # --------------------
    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):
    """
    pts: (N,2) array
    dist_thresh: 点到拟合直线的最大允许距离
    返回 (m, b) 直线参数, 以及拟合用到的点
    """
    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
    x2, y2 = x[mask], y[mask]

    if len(x2) < 2:
        return (m, b), pts  # 保底返回

    # 二次拟合
    m, b = np.polyfit(y2, x2, 1)
    return (m, b), np.stack([x2, y2], axis=1)


# --------------------
# 推理 ROI 并可视化 mask + 边缘角点 + 拟合直线
# --------------------
def infer_mask_with_selected_corners(image_path, model_path, output_dir="./output"):
    model = YOLO(model_path)

    image_path = Path(image_path)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    orig_img = cv2.imread(str(image_path))
    overlay_img = orig_img.copy()

    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]

        # 可视化 mask
        if result.masks is not None and len(result.masks.data) > 0:
            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)

            # 绿色 mask 覆盖
            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, left_inliers = fit_line_with_outlier_removal(left_pts)
            right_line, right_inliers = fit_line_with_outlier_removal(right_pts)

            # 可视化角点
            for cx, cy in left_inliers:
                cv2.circle(overlay_img[y:y+h, x:x+w], (int(cx), int(cy)), 5, (0, 0, 255), -1)
            for cx, cy in right_inliers:
                cv2.circle(overlay_img[y:y+h, x:x+w], (int(cx), int(cy)), 5, (255, 0, 0), -1)

            # 可视化拟合直线
            if left_line is not None:
                m, b = left_line
                y1, y2 = 0, h
                x1, x2 = int(m * y1 + b), int(m * y2 + b)
                cv2.line(overlay_img[y:y+h, x:x+w], (x1, y1), (x2, y2), (0, 0, 255), 3)

            if right_line is not None:
                m, b = right_line
                y1, y2 = 0, h
                x1, x2 = int(m * y1 + b), int(m * y2 + b)
                cv2.line(overlay_img[y:y+h, x:x+w], (x1, y1), (x2, y2), (255, 0, 0), 3)

    # 保存结果
    save_path = output_dir / f"mask_edge_corners_{image_path.name}"
    cv2.imwrite(str(save_path), overlay_img)
    print(f"✅ 保存结果: {save_path}")

    return overlay_img


# ===================== 使用示例 =====================
if __name__ == "__main__":
    IMAGE_PATH = "../test_image/1.png"
    MODEL_PATH = "best.pt"
    infer_mask_with_selected_corners(IMAGE_PATH, MODEL_PATH)