# -*- coding: utf-8 -*-
"""
共享工具函数：
    - imread_unicode: 兼容 Windows 中文路径的图片读取
    - segment_object_grabcut: 用 GrabCut 从背景(转盘)中分割出物体，返回二值掩码
    - pick_two_points: 交互式点选2个点(用于标定时点选已知长度的两端)
    - select_roi_unicode: 封装 cv2.selectROI，用于框选物体大致范围供 GrabCut 使用
"""
import cv2
import numpy as np
import json
import os


def imread_unicode(path):
    """兼容 Windows 下含中文/特殊字符路径的图片读取（cv2.imread 对非ASCII路径常读取失败）。"""
    data = np.fromfile(path, dtype=np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_COLOR)
    return img


def imread_unicode_unchanged(path):
    """
    同上，但保留 alpha 通道(用于读取抠图后的透明背景PNG)。
    返回的图像若原图有透明通道，shape为 (H,W,4)，否则为 (H,W,3)。
    """
    data = np.fromfile(path, dtype=np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_UNCHANGED)
    return img


def mask_from_alpha(png_img, alpha_thresh=127):
    """
    从抠图PNG的alpha通道直接生成二值掩码，不需要GrabCut/手动框选。
    要求 png_img 是4通道(BGRA)图像。
    """
    if png_img is None:
        raise RuntimeError("图片读取失败(为None)。")
    if png_img.ndim != 3 or png_img.shape[2] != 4:
        raise RuntimeError(
            "该图片没有透明通道(不是4通道BGRA的PNG)，无法从alpha直接取掩码。"
            "请确认传入的是抠图后的透明背景PNG，而不是普通jpg/无透明通道的png。"
        )
    alpha = png_img[:, :, 3]
    mask = np.where(alpha > alpha_thresh, 255, 0).astype(np.uint8)
    # 轻微开闭运算去除抠图边缘的毛刺噪点，不改变整体轮廓
    kernel = np.ones((3, 3), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    return mask


def composite_on_white(bgra_img):
    """
    把带透明通道的抠图结果合成到白色背景上，得到一张普通3通道BGR图，
    用于在没有提供原图时也能生成清晰的可视化标注图。
    """
    b, g, r, a = cv2.split(bgra_img.astype(np.float32))
    alpha = a / 255.0
    white = np.full_like(b, 255.0)
    out_b = b * alpha + white * (1 - alpha)
    out_g = g * alpha + white * (1 - alpha)
    out_r = r * alpha + white * (1 - alpha)
    return cv2.merge([out_b, out_g, out_r]).astype(np.uint8)


def imwrite_unicode(path, img):
    """兼容 Windows 中文路径的图片写入。"""
    ext = os.path.splitext(path)[1]
    ok, buf = cv2.imencode(ext, img)
    if ok:
        buf.tofile(path)
    return ok


def save_json(path, data):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def load_json(path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def select_roi_unicode(window_name, img, max_display_width=1200):
    """
    框选物体大致范围(拖出一个矩形框住鞋子，尽量贴合但不用太精确)。
    大图会先等比缩小显示，避免窗口超出屏幕，选完自动换算回原图坐标。
    操作：鼠标拖拽画框 -> 按空格/回车确认 -> 按c取消重选。
    """
    h, w = img.shape[:2]
    scale = min(1.0, max_display_width / w)
    disp = cv2.resize(img, (int(w * scale), int(h * scale))) if scale < 1.0 else img.copy()

    print(f"[{window_name}] 请用鼠标拖拽框选出物体大致范围(不用太精确，比物体略大即可)，"
          f"框完按空格/回车确认。")
    r = cv2.selectROI(window_name, disp, showCrosshair=True, fromCenter=False)
    cv2.destroyWindow(window_name)
    x, y, rw, rh = r
    if rw == 0 or rh == 0:
        raise RuntimeError("未选择有效区域，请重新运行并框选物体。")
    # 换算回原图坐标
    return (int(x / scale), int(y / scale), int(rw / scale), int(rh / scale))


def segment_object_grabcut(img, rect, iter_count=5, margin=0.06):
    """
    用 GrabCut 从背景(转盘/桌面)中把物体分割出来。
    rect: (x, y, w, h) 物体大致所在的框选区域(来自 select_roi_unicode)
    margin: 在框选区域基础上再收缩一点作为"确定前景"的种子区域，提高分割稳定性
    返回: 二值掩码(255=物体, 0=背景)，与原图同尺寸
    """
    mask = np.zeros(img.shape[:2], np.uint8)
    bgd_model = np.zeros((1, 65), np.float64)
    fgd_model = np.zeros((1, 65), np.float64)

    cv2.grabCut(img, mask, rect, bgd_model, fgd_model, iter_count, cv2.GC_INIT_WITH_RECT)
    mask2 = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype(np.uint8)

    # 形态学开闭运算去除噪点、填补小空洞
    kernel = np.ones((5, 5), np.uint8)
    mask2 = cv2.morphologyEx(mask2, cv2.MORPH_OPEN, kernel)
    mask2 = cv2.morphologyEx(mask2, cv2.MORPH_CLOSE, kernel)
    return mask2


def largest_contour(mask):
    """从二值掩码中取面积最大的轮廓（即物体主体，排除噪点）。"""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        raise RuntimeError("未能从分割结果中找到有效轮廓，请重新框选或检查背景是否干净。")
    return max(contours, key=cv2.contourArea)


class TwoPointPicker:
    """
    交互式点选2个点，用于标定时点出已知长度参照物的两端
    (比如标定尺的0刻度和300mm刻度，或标定块的顶部和底部)。
    """
    def __init__(self, window_name, img, labels=("点1", "点2"), max_display_width=1400):
        self.window_name = window_name
        self.orig_img = img
        h, w = img.shape[:2]
        self.scale = min(1.0, max_display_width / w)
        self.img = cv2.resize(img, (int(w * self.scale), int(h * self.scale))) \
            if self.scale < 1.0 else img.copy()
        self.display = self.img.copy()
        self.points = []
        self.labels = labels

    def _on_mouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN and len(self.points) < 2:
            self.points.append((x, y))
            self._redraw()

    def _redraw(self):
        self.display = self.img.copy()
        for i, (x, y) in enumerate(self.points):
            cv2.circle(self.display, (x, y), 6, (0, 0, 255), -1)
            cv2.putText(self.display, self.labels[i], (x + 8, y - 8),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        if len(self.points) == 2:
            cv2.line(self.display, self.points[0], self.points[1], (0, 255, 0), 2)
        cv2.imshow(self.window_name, self.display)

    def run(self):
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.window_name, self._on_mouse)
        print(f"请依次点击: {self.labels[0]} -> {self.labels[1]}")
        print("点完2个点后按任意键确认，按 r 重来。")
        self._redraw()
        while True:
            key = cv2.waitKey(20) & 0xFF
            if key == ord('r'):
                self.points = []
                self._redraw()
            elif len(self.points) == 2 and key != 255:
                break
        cv2.destroyAllWindows()
        # 换算回原图坐标
        return [(x / self.scale, y / self.scale) for x, y in self.points]