import json
import os
import glob

def labelme_to_yolo_keypoints_batch(json_dir, output_dir, target_box_label="J1", class_id=0, img_shape=None, keypoints_per_instance=4):
    """
    批量转换 LabelMe JSON → YOLO Pose 格式 (.txt)
    - 每 keypoints_per_instance 个关键点对应一个 target_box_label 实例
    - 关键点必须与框一一对应
    - 转换失败或数据不匹配时，删除 JSON 和对应图片
    """
    if img_shape is None:
        raise ValueError("必须提供 img_shape 参数，例如 (1440, 2506)")

    os.makedirs(output_dir, exist_ok=True)
    json_files = glob.glob(os.path.join(json_dir, "*.json"))
    json_files = [f for f in json_files if os.path.isfile(f) and not f.endswith("_mask.json")]

    if not json_files:
        print(f"❌ 在 {json_dir} 中未找到任何 JSON 文件")
        return

    img_h, img_w = img_shape
    converted_count = 0
    deleted_count = 0

    print(f"🔍 开始转换：目标框='{target_box_label}', 每实例 {keypoints_per_instance} 个关键点")

    for json_file in json_files:
        success = False
        base_name = os.path.splitext(os.path.basename(json_file))[0]
        output_path = os.path.join(output_dir, f"{base_name}.txt")
        image_file_to_delete = None

        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)

            # 获取图片路径
            image_path = data.get("imagePath")
            if image_path:
                image_file_to_delete = os.path.join(json_dir, os.path.basename(image_path))
            else:
                print(f"⚠️  {base_name}: JSON 中无 imagePath，无法定位图片")
                image_file_to_delete = None

            # 提取关键点
            keypoint_labels = {str(i) for i in range(1, keypoints_per_instance + 1)}  # 支持 1,2,3,4...
            keypoints = []
            for shape in data.get("shapes", []):
                label = shape["label"]
                if shape["shape_type"] == "point" and label in keypoint_labels:
                    x, y = shape["points"][0]
                    nx = x / img_w
                    ny = y / img_h
                    # 归一化并裁剪到 [0,1]
                    nx = max(0.0, min(1.0, nx))
                    ny = max(0.0, min(1.0, ny))
                    keypoints.append((label, nx, ny))

            # 提取 J1 矩形框（每个框是一个实例）
            j1_boxes = [
                s for s in data.get("shapes", [])
                if s["label"] == target_box_label and s["shape_type"] == "rectangle"
            ]
            num_instances = len(j1_boxes)
            total_keypoints = len(keypoints)
            expected_total = num_instances * keypoints_per_instance

            # 检查数量匹配
            if total_keypoints != expected_total:
                print(f"❌ {base_name}: 关键点数量不匹配！期望 {expected_total}，实际 {total_keypoints}")
                raise ValueError("关键点数量不匹配")
            if num_instances == 0:
                print(f"❌ {base_name}: 未找到任何 '{target_box_label}' 框")
                raise ValueError("无目标框")

            # 提取并归一化每个框的坐标
            bboxes = []
            for box in j1_boxes:
                x1, y1 = box["points"][0]
                x2, y2 = box["points"][1]
                # 计算中心点和宽高
                x_center = (x1 + x2) / 2 / img_w
                y_center = (y1 + y2) / 2 / img_h
                w = abs(x2 - x1) / img_w
                h = abs(y2 - y1) / img_h
                # 裁剪到 [0,1]
                x_center = max(0.0, min(1.0, x_center))
                y_center = max(0.0, min(1.0, y_center))
                w = max(0.0, min(1.0, w))
                h = max(0.0, min(1.0, h))
                bboxes.append((x_center, y_center, w, h))

            # 开始写入 YOLO 文件
            with open(output_path, 'w', encoding='utf-8') as f_out:
                for i in range(num_instances):
                    # 获取该实例对应的 4 个关键点
                    start_idx = i * keypoints_per_instance
                    end_idx = start_idx + keypoints_per_instance
                    group = keypoints[start_idx:end_idx]

                    if len(group) != keypoints_per_instance:
                        print(f"❌ {base_name}: 实例 {i+1} 缺少关键点")
                        raise ValueError("实例关键点不足")

                    # 按标签排序关键点 (1,2,3,4)
                    sorted_group = sorted(group, key=lambda x: x[0])

                    # 构造 YOLO 行：class + bbox + keypoints
                    yolo_line = [
                        str(class_id),
                        f"{bboxes[i][0]:.6f}",  # x_center
                        f"{bboxes[i][1]:.6f}",  # y_center
                        f"{bboxes[i][2]:.6f}",  # w
                        f"{bboxes[i][3]:.6f}"   # h
                    ]
                    for _, kx, ky in sorted_group:
                        yolo_line.extend([f"{kx:.6f}", f"{ky:.6f}", "2"])  # v=2 表示可见

                    f_out.write(" ".join(yolo_line) + "\n")

            print(f"✅ 已转换: {os.path.basename(json_file)} -> {num_instances} 个实例")
            success = True
            converted_count += 1

        except Exception as e:
            print(f"❌ 转换失败 {base_name}: {e}")

            # 删除无效文件
            try:
                if os.path.exists(output_path):
                    os.remove(output_path)
                if image_file_to_delete and os.path.exists(image_file_to_delete):
                    os.remove(image_file_to_delete)
                    print(f"🗑️  已删除图片: {os.path.basename(image_file_to_delete)}")
                if os.path.exists(json_file):
                    os.remove(json_file)
                    print(f"🗑️  已删除 JSON: {os.path.basename(json_file)}")
                deleted_count += 1
            except Exception as del_e:
                print(f"💥 删除文件时出错: {del_e}")

    print("\n" + "="*60)
    print(f"🎉 批量转换完成！")
    print(f"✅ 成功保留: {converted_count} 个文件")
    print(f"❌ 异常删除: {deleted_count} 个文件（JSON + 图片）")
    print(f"📁 输出目录: {output_dir}")
    print(f"📦 每实例关键点数: {keypoints_per_instance}")
    print(f"🏷️  目标框标签: {target_box_label}")
    print("="*60)


# ================== 用户配置区 ==================
JSON_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/yolodataset/point1/f11"
OUTPUT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/yolodataset/point1/labels_keypoints"
TARGET_BOX_LABEL = "J1"
CLASS_ID = 1
IMG_SHAPE = (1440, 2506)  # (height, width)
KEYPOINTS_PER_INSTANCE = 3

# ================== 执行转换 ==================
if __name__ == "__main__":
    labelme_to_yolo_keypoints_batch(
        json_dir=JSON_DIR,
        output_dir=OUTPUT_DIR,
        target_box_label=TARGET_BOX_LABEL,
        class_id=CLASS_ID,
        img_shape=IMG_SHAPE,
        keypoints_per_instance=KEYPOINTS_PER_INSTANCE
    )