# -*- coding: utf-8 -*-
"""
批量把 dataset/onepiece_from_svg 下的单块 JSON 渲染为 PNG，
输出到 dataset/onepiece_images。

兼容两种输入结构：
- 顶层是单个 piece 的 dict（旧式）
- 顶层是 list[dict]（新式，仅取第一个/或逐个渲染）
"""
import os, json, sys
import matplotlib
matplotlib.use("Agg")  # 后端渲染，不弹窗
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

# --- 路径 ---
ROOT = "/Users/zongyikun/Desktop/qwen_vl_demo/kilogram-main"
IN_DIR  = os.path.join(ROOT, "dataset/onepiece_from_svg")
OUT_DIR = os.path.join(ROOT, "dataset/onepiece_images")
os.makedirs(OUT_DIR, exist_ok=True)

# 引入 geometry.py 的七巧板定义
sys.path.append(os.path.join(ROOT))
from geometry import tangram_pieces, _NAME_TO_PIECE, _coerce_to_response

def _normalize_ann(obj):
    """将文件内容转成 _coerce_to_response 能够识别的结构。
    - 若是 dict（单块），转成 list[dict]
    - 若是 list，直接返回 list
    """
    if isinstance(obj, dict):
        return [obj]
    if isinstance(obj, list):
        return obj
    return []

def draw_and_save(response, out_png, title=None):
    # 画布
    fig, ax = plt.subplots(figsize=(6, 6), dpi=120)
    ax.set_aspect('equal')
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.set_xticks(range(0, 11))
    ax.set_yticks(range(0, 11))
    ax.grid(True, which='both', linestyle='--', linewidth=0.5, color='gray', alpha=0.5)
    if title:
        ax.set_title(title)

    # 重置模板（确保每次从原始顶点开始）
    for p in tangram_pieces:
        p.transformed_vertices = p.original_vertices.copy()

    # 逐块绘制（one-piece 文件通常只有一个）
    for inst in response:
        name = inst.get("name")
        piece = _NAME_TO_PIECE.get(name)
        if not piece:
            continue
        piece.set_by_instruction(inst)
        poly = Polygon(piece.transformed_vertices, closed=True,
                       facecolor=piece.color, alpha=0.6, edgecolor='black')
        ax.add_patch(poly)
        cx, cy = piece.transformed_vertices.mean(axis=0)
        ax.text(cx, cy, name, ha='center', va='center', fontsize=8, color='black')

    fig.tight_layout()
    fig.savefig(out_png, bbox_inches="tight")
    plt.close(fig)

def main():
    cnt_in, cnt_out = 0, 0
    for fname in sorted(os.listdir(IN_DIR)):
        if not fname.endswith(".json"):
            continue
        in_path = os.path.join(IN_DIR, fname)
        with open(in_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        cnt_in += 1

        # 统一结构 -> 再用 geometry._coerce_to_response 规范化为内部 response
        ann_list = _normalize_ann(data)
        response = _coerce_to_response(ann_list)
        if not response:
            print(f"[SKIP] {fname} 无可用 piece（可能字段不兼容）")
            continue

        # 仅渲染第一块（如果你希望每个 list 元素都渲染一个 PNG，请在这里循环）
        # 但在我们的 onepiece_from_svg 中，文件本身就表示单块
        out_png = os.path.join(OUT_DIR, os.path.splitext(fname)[0] + ".png")
        draw_and_save(response, out_png, title=os.path.splitext(fname)[0])
        cnt_out += 1
        if cnt_out % 50 == 0:
            print(f"[INFO] 已输出 {cnt_out} 张")

    print(f"\n✅ Done! 输入 JSON: {cnt_in}，输出 PNG: {cnt_out}")
    print(f"输出目录: {OUT_DIR}")

if __name__ == "__main__":
    main()