import os
import json
import cv2
import numpy as np
import math

# === 配置 ===
input_folder = r"/data/share/zyh/master_dataset/dataset_net/pokou_251115_251121/resize"  # 输入 JSON 和图片
output_folder = os.path.join(os.path.dirname(input_folder), "show")
os.makedirs(output_folder, exist_ok=True)


# === 工具函数：生成椭圆轨迹点 ===
def ellipse_points(cx, cy, rx, ry, theta, num=100):
    """
    返回椭圆轨迹上的 num 个点
    theta 单位为弧度
    """
    t = np.linspace(0, 2 * np.pi, num)
    x = rx * np.cos(t)
    y = ry * np.sin(t)
    # 旋转
    xr = x * np.cos(theta) - y * np.sin(theta)
    yr = x * np.sin(theta) + y * np.cos(theta)
    # 平移
    xr += cx
    yr += cy
    return np.stack([xr, yr], axis=1).astype(np.int32)


# === 遍历输出文件夹 ===
for file in os.listdir(input_folder):
    if not file.endswith(".json"):
        continue
    json_path = os.path.join(input_folder, file)
    img_name = file.replace(".json", ".jpg")
    img_path = os.path.join(input_folder, img_name)

    if not os.path.exists(img_path):
        print(f"?? Image not found: {img_name}")
        continue

    img = cv2.imread(img_path)
    if img is None:
        continue

    with open(json_path, "r", encoding="utf-8") as jf:
        data = json.load(jf)

    # === 绘制 shapes ===
    for shape in data.get("shapes", []):
        if shape.get("label") == "arc":
            pts = np.array(shape.get("points", []), dtype=np.int32)
            ends = np.array(shape.get("ends", []), dtype=np.int32)
            params = shape.get("params", [0, 0, 0, 0, 0])
            cx, cy, rx, ry, theta = params

            # 绘制三点
            for p in pts:
                cv2.circle(img, (int(p[0]), int(p[1])), 5, (0, 0, 255), -1)  # 红色

            # 绘制端点
            for p in ends:
                cv2.circle(img, (int(p[0]), int(p[1])), 7, (0, 255, 0), -1)  # 绿色

            # 绘制椭圆轨迹
            ep = ellipse_points(cx, cy, rx, ry, theta)
            for i in range(len(ep) - 1):
                cv2.line(img, tuple(ep[i]), tuple(ep[i + 1]), (255, 0, 0), 2)  # 蓝色轨迹

        elif shape.get("shape_type") == "line":
            pts = shape.get("points", [])
            if len(pts) >= 2:
                cv2.line(img, tuple(map(int, pts[0])), tuple(map(int, pts[1])), (0, 255, 255), 2)  # 黄色线

    # === 保存结果 ===
    save_path = os.path.join(output_folder, img_name)
    cv2.imwrite(save_path, img)
    print(f"Saved visualization for {img_name} -> {save_path}")

print("\nAll done! Visualizations saved in:", output_folder)