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="arc"):
    points_np = np.array(points)
    if points_np.shape[0] < 3:
        return []

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

    shape = {
        "label": label,
        "points": points,
        "shape_type": "polygon",
        "flags": {},
        "xmin": int(xs.min()),
        "ymin": int(ys.min()),
        "xmax": int(xs.max()),
        "ymax": int(ys.max())
    }
    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') == 'arc' and len(item.get('points', [])) == 1]

        if len(arc_points_raw) < 3:
            print(f"{file_name} 中 arc 点数不足 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="arc")
                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} 没有成功拟合的 arc 区域，跳过")
            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✅ 所有 arc 区域已保存为 labelme 格式（图像 + json）到目录：{output_dir}")


# ===== 修改为你自己的输入输出路径 =====
input_dir = r"G:\python_ws_g\data\test\jsonjpg_weizhuang12"
output_dir = r"G:\python_ws_g\data\test\circle"
point_step_deg = 0.5
draw_radius = 2

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