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)
    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()

    # 极角
    angles = np.arctan2(points[:, 1] - yc, points[:, 0] - xc)
    start_angle = np.min(angles)
    end_angle = np.max(angles)

    # 确保角度顺序正确
    if end_angle - start_angle > np.pi:
        start_angle, end_angle = end_angle, start_angle + 2 * np.pi

    arc_angles = np.arange(start_angle, end_angle, np.deg2rad(point_step_deg))
    arc_points = np.stack([
        xc + R * np.cos(arc_angles),
        yc + R * np.sin(arc_angles)
    ], axis=-1).astype(np.int32)

    # 创建 mask 并画出拟合的弧线（扩展为像素块）
    mask = np.zeros(shape, dtype=np.uint8)
    for pt in arc_points:
        cv2.circle(mask, tuple(pt), radius=radius, color=255, thickness=-1)

    return mask, arc_points.tolist()


def build_shapes_from_mask(mask, label="arc"):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    shapes = []

    for contour in contours:
        contour = contour.squeeze()
        if len(contour.shape) != 2 or contour.shape[0] < 3:
            continue  # 不是合法 polygon

        points = contour.tolist()
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]

        shape = {
            "label": label,
            "points": points,
            "shape_type": "polygon",
            "flags": {},
            "xmin": int(min(xs)),
            "ymin": int(min(ys)),
            "xmax": int(max(xs)),
            "ymax": int(max(ys))
        }
        shapes.append(shape)

    return shapes


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)

        points_list = [item['points'][0] for item in label_data['shapes'] if item['label'] == 'arc']
        if len(points_list) != 3:
            print(f"{file_name} 中 arc 点数不足 3，跳过")
            continue

        image = cv2.imread(image_path)
        h, w = image.shape[:2]
        points = np.array(points_list)
        mask, arc_points = fit_arc_and_create_mask(points, (h, w), point_step_deg, radius)

        shapes = build_shapes_from_mask(mask, label="arc")
        if len(shapes) == 0:
            print(f"{file_name} 拟合失败，跳过")
            continue

        output_json = {
            "version": "5.0.1",
            "flags": {},
            "shapes": 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✅ 所有 arc 区域已保存为 labelme 格式（图像 + json）到目录：{output_dir}")


# ===== 修改为你自己的输入输出路径 =====
input_dir = r"G:\python_ws_g\data\arc"
output_dir = r"G:\python_ws_g\data\arcout"
point_step_deg = 0.2
draw_radius = 3


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