import os
import json
import shutil
import numpy as np
from tqdm import tqdm


def random_12_points(h, w, num_points=12, margin=20):
    xs = np.random.uniform(margin, w - margin, num_points)
    ys = np.random.uniform(margin, h - margin, num_points)
    points = np.stack([xs, ys], axis=-1).astype(np.float32)
    return points.tolist()


def build_point_shape(pt, label="arc"):
    return {
        "label": label,
        "points": [pt],
        "shape_type": "point",
        "flags": {},
        "attributes": {},
        "group_id": None,
        "description": "",
        "difficult": False,
        "kie_linking": []
    }


def process_json_folder(input_dir, output_dir, num_points=12, margin=20):
    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)
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        arc_points_raw = [item['points'][0] for item in data['shapes']
                          if item.get('label') == 'arc' and item.get('shape_type') == 'point']

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

        num_groups = len(arc_points_raw) // 3
        new_shapes = []
        h = data.get("imageHeight", 2000)
        w = data.get("imageWidth", 2000)

        for i in range(num_groups):
            # 这里不使用拟合，直接随机生成12个点
            try:
                circle_pts = random_12_points(h, w, num_points=num_points, margin=margin)
                for pt in circle_pts:
                    shape = build_point_shape(pt, label="arc")
                    new_shapes.append(shape)

            except Exception as e:
                print(f"{file_name} 第 {i+1} 组生成失败，跳过。错误：{e}")
                continue

        if not new_shapes:
            continue

        new_json = {
            "version": data.get("version", "5.0.1"),
            "flags": {},
            "shapes": new_shapes,
            "imagePath": data["imagePath"],
            "imageHeight": h,
            "imageWidth": w,
            "imageData": None
        }

        base_name = os.path.splitext(file_name)[0]
        out_json_path = os.path.join(output_dir, base_name + ".json")
        out_img_path = os.path.join(output_dir, data["imagePath"])

        with open(out_json_path, "w", encoding="utf-8") as f:
            json.dump(new_json, f, indent=4)

        old_img_path = os.path.join(input_dir, data["imagePath"])
        if os.path.exists(old_img_path):
            shutil.copy(old_img_path, out_img_path)

    print(f"\n✅ 已处理完成，输出保存至：{output_dir}")


# ========= 路径设置 =========
input_dir = r"G:\python_ws_g\data\test\jsonjpg"
output_dir = r"G:\python_ws_g\data\test\jsonjpg_weizhuang12"

if __name__ == "__main__":
    process_json_folder(input_dir, output_dir, num_points=12, margin=20)