import os
import json
import shutil
import numpy as np
import cv2
from scipy.optimize import leastsq
from tqdm import tqdm


def calc_R(x, y, xc, yc):
    return np.sqrt((x - xc) ** 2 + (y - yc) ** 2)


def f(c, x, y):
    Ri = calc_R(x, y, *c)
    return Ri - Ri.mean()

def fit_arc_and_create_mask(points, shape, point_step_deg=0.5, radius=2):
    x_data = points[:, 0]
    y_data = points[:, 1]
    center_estimate = np.mean(x_data), np.mean(y_data)
    try:
        center, _ = leastsq(f, center_estimate, args=(x_data, y_data))
        xc, yc = center
        Ri = calc_R(x_data, y_data, xc, yc)
        R = Ri.mean()

        h, w = shape
        if R <= 0 or not np.isfinite(R) or not np.isfinite(xc) or not np.isfinite(yc):
            raise ValueError("拟合圆非法")

        # 生成完整圆的点
        circle_angles = np.linspace(0, 2 * np.pi, int(360 / point_step_deg), endpoint=False)
        full_circle_points = np.stack([
            xc + R * np.cos(circle_angles),
            yc + R * np.sin(circle_angles)
        ], axis=-1).astype(np.float32)

        # 只保留图像内部的点
        in_bounds_mask = (
            (full_circle_points[:, 0] >= 0) & (full_circle_points[:, 0] < w) &
            (full_circle_points[:, 1] >= 0) & (full_circle_points[:, 1] < h)
        )
        clipped_points = full_circle_points[in_bounds_mask]

        if len(clipped_points) < 3:
            print(f"拟合失败使用的三个点为：\n{points}")
            raise ValueError("拟合圆点全部在图像外")

        return None, clipped_points.tolist()

    except Exception as e:
        print(f"⚠️ 圆拟合失败：{e}")
        return None, []


def build_shapes_from_points(points, label="circle"):
    points_np = np.array(points)
    if points_np.shape[0] < 3:
        return []

    xs = points_np[:, 0]
    ys = points_np[:, 1]

    xmin, xmax = xs.min(), xs.max()
    ymin, ymax = ys.min(), ys.max()

    eps = 1e-3

    # 边界点：每类只能找一个
    def find_one(points, cond):
        for pt in points:
            if cond(pt):
                return pt
        return None

    pt_xmin = find_one(points, lambda pt: abs(pt[0] - xmin) < eps)
    pt_xmax = find_one(points, lambda pt: abs(pt[0] - xmax) < eps)
    pt_ymin = find_one(points, lambda pt: abs(pt[1] - ymin) < eps)
    pt_ymax = find_one(points, lambda pt: abs(pt[1] - ymax) < eps)

    boundary_points = []
    for pt in [pt_xmin, pt_xmax, pt_ymin, pt_ymax]:
        if pt is not None and pt not in boundary_points:
            boundary_points.append(pt)

    if len(boundary_points) != 4:
        print("⚠️ 边界点不足 4 个，跳过该 shape")
        return []

    shape = {
        "label": label,
        "points": boundary_points,
        "shape_type": "polygon",
        "flags": {},
        "xmin": int(xmin),
        "ymin": int(ymin),
        "xmax": int(xmax),
        "ymax": int(ymax)
    }
    return [shape]




def process_folder_labelme(input_dir, output_dir, point_step_deg=0.5, radius=2):
    os.makedirs(output_dir, exist_ok=True)

    for file_name in tqdm(os.listdir(input_dir)):
        if not file_name.endswith(".json"):
            continue

        json_path = os.path.join(input_dir, file_name)
        image_path = json_path.replace(".json", ".jpg")
        if not os.path.exists(image_path):
            print(f"图像不存在，跳过：{image_path}")
            continue

        with open(json_path, "r") as f:
            label_data = json.load(f)

        arc_points_raw = [item['points'][0] for item in label_data['shapes']
                          if item.get('label') == 'circle' and len(item.get('points', [])) == 1]

        if len(arc_points_raw) < 3:
            print(f"{file_name} 中 circle 点数不足 3，跳过")
            continue

        image = cv2.imread(image_path)
        h, w = image.shape[:2]

        arc_shapes = []
        num_groups = len(arc_points_raw) // 3

        for i in range(num_groups):
            group_points = arc_points_raw[i*3:(i+1)*3]
            if len(group_points) < 3:
                print(f"{file_name} 第 {i+1} 组点数不足 3，跳过")
                continue
            points = np.array(group_points)
            try:
                _, arc_pts = fit_arc_and_create_mask(points, (h, w), point_step_deg, radius)
                shapes = build_shapes_from_points(arc_pts, label="circle")
                arc_shapes.extend(shapes)
            except Exception as e:
                print(f"{file_name} 第 {i+1} 组拟合失败，跳过。错误：{e}")
                continue

        if len(arc_shapes) == 0:
            print(f"{file_name} 没有成功拟合的 circle 区域，跳过")
            continue

        output_json = {
            "version": "5.0.1",
            "flags": {},
            "shapes": arc_shapes,
            "imagePath": os.path.basename(image_path),
            "imageHeight": h,
            "imageWidth": w
        }

        base_name = os.path.splitext(os.path.basename(image_path))[0]
        out_json_path = os.path.join(output_dir, base_name + ".json")
        out_img_path = os.path.join(output_dir, base_name + ".jpg")

        with open(out_json_path, "w") as f:
            json.dump(output_json, f, indent=4)

        shutil.copy(image_path, out_img_path)

    print(f"\n✅ 所有 circle 区域已保存为 labelme 格式（图像 + json）到目录：{output_dir}")


# ===== 修改为你自己的输入输出路径 =====
input_dir = r"\\192.168.50.222\share\zyh\data\rgb_453_source_json\dataset_for_guanban"
output_dir = r"\\192.168.50.222\share\zyh\data\rgb_453_source_json\4point"
point_step_deg = 0.5
draw_radius = 2

if __name__ == "__main__":
    process_folder_labelme(input_dir, output_dir, point_step_deg, draw_radius)