# 文件名: test_with_geometry.py
import os
import json
import argparse
import matplotlib.pyplot as plt
from geometry import render_outline_from_json, load_gt_outline, compute_iou_mask

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--id", required=True, help="例如 page1-130")
    args = parser.parse_args()

    ann_file = f"annotations/{args.id}.json"
    gt_file = f"outlines/{args.id}.svg"  # 或者你的 GT 路径

    if not os.path.exists(ann_file):
        raise FileNotFoundError(f"找不到 {ann_file}")
    if not os.path.exists(gt_file):
        raise FileNotFoundError(f"找不到 {gt_file}")

    # 载入并渲染
    with open(ann_file) as f:
        ann_data = json.load(f)

    gt_mask = load_gt_outline(gt_file, raster=160)
    pred_mask = render_outline_from_json(ann_data, raster=160)

    iou = compute_iou_mask(gt_mask, pred_mask)

    # 可视化
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    axes[0].imshow(gt_mask, cmap="gray")
    axes[0].set_title("GT Outline")
    axes[1].imshow(pred_mask, cmap="gray")
    axes[1].set_title("Pred from JSON")
    overlay = (gt_mask.astype(int) - pred_mask.astype(int))
    axes[2].imshow(overlay, cmap="bwr")
    axes[2].set_title(f"Overlay IoU={iou:.3f}")

    for ax in axes:
        ax.axis("off")

    os.makedirs("debug", exist_ok=True)
    out_path = f"debug/{args.id}.png"
    plt.savefig(out_path)
    plt.close(fig)
    print(f"[DONE] Saved comparison to {out_path}")

if __name__ == "__main__":
    main() 
# -*- coding: utf-8 -*-
# 文件名: test_with_geometry.py
"""
把 annotations/<ID>.json 复原成几何拼图，与 outlines/<ID>.png 做可视化对比：
左=GT  中=Pred(由JSON复原)  右=误差着色(灰=TP, 红=FP, 蓝=FN)
同时自动在全局角度 {0,45,...,315} 中寻找最佳 IoU。
输出保存到 ./debug/<ID>.png
"""
import os
import json
import argparse
from typing import List, Dict, Any

# 数据路径（按你的工程结构）
ROOT = os.path.dirname(__file__)
ANN  = os.path.join(ROOT, "kilogram-main", "dataset", "annotations")
OUT  = os.path.join(ROOT, "kilogram-main", "dataset", "outlines")

# 映射到 geometry 引擎的名称
NAME_MAP = {
    ("big_triangle",   "L1"): "T1",
    ("big_triangle",   "L2"): "T2",
    ("medium_triangle", None):  "T3",
    ("small_triangle", "S1"): "T4",
    ("small_triangle", "S2"): "T5",
    ("square",          None):  "S",
    ("parallelogram",   None):  "P",
}

# ---- 将 annotation(类型制) 规范成 geometry 响应制 ----

def _canon_instances(pieces: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """为 big/small 两对赋 instance(L1/L2, S1/S2)，按质心 x→y 排序。"""
    def key_xy(p):
        x, y = p.get("pos", [0, 0])
        return (float(x), float(y))
    out = []
    bigs   = [p.copy() for p in pieces if p.get("type") == "big_triangle"]
    meds   = [p.copy() for p in pieces if p.get("type") == "medium_triangle"]
    smalls = [p.copy() for p in pieces if p.get("type") == "small_triangle"]
    squares= [p.copy() for p in pieces if p.get("type") == "square"]
    paras  = [p.copy() for p in pieces if p.get("type") == "parallelogram"]
    bigs.sort(key=key_xy)
    smalls.sort(key=key_xy)
    for i, p in enumerate(bigs):
        q = p.copy(); q["instance"] = "L1" if i == 0 else "L2"; out.append(q)
    for i, p in enumerate(smalls):
        q = p.copy(); q["instance"] = "S1" if i == 0 else "S2"; out.append(q)
    out += meds + squares + paras
    return out


def to_geometry_response(pieces: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    pcs = _canon_instances(pieces)
    out = []
    for p in pcs:
        t = p.get("type")
        inst = p.get("instance")
        name = NAME_MAP.get((t, inst)) or NAME_MAP.get((t, None))
        if not name:
            continue
        out.append({
            "name":  name,
            "pos":   p.get("pos", [0, 0]),
            "angle": float(p.get("angle", 0)),
            "flip":  bool(p.get("flip", False)),
            "scale": float(p.get("scale", 1.0)),
        })
    return out

# ---- 主逻辑：找最佳角度 + 生成可视化 ----

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--id", required=True, help="例如 page1-130")
    args = ap.parse_args()

    ann_path = os.path.join(ANN, f"{args.id}.json")
    out_path = os.path.join(OUT, f"{args.id}.png")  # 我们的评测与可视化都以 PNG 轮廓为准

    if not os.path.exists(ann_path):
        raise FileNotFoundError(f"找不到注释: {ann_path}")
    if not os.path.exists(out_path):
        raise FileNotFoundError(f"找不到轮廓图: {out_path}")

    ann = json.load(open(ann_path, "r", encoding="utf-8"))
    base = to_geometry_response(ann)  # 严格使用原始角度/缩放/位置

    # 动态导入我们已实现的评测+可视化函数
    import result_test as rt

    # 在 0..315(步长45°) 中搜索最佳全局角
    best_iou, best_deg, best_rsp = -1.0, 0, base
    for k in range(8):
        deg = 45 * k
        rsp = [{**p, "angle": (p.get("angle", 0) + deg) % 360} for p in base]
        iou_k = float(rt.iou_with_outline(out_path, rsp))
        if iou_k > best_iou:
            best_iou, best_deg, best_rsp = iou_k, deg, rsp

    # 生成三联图到 ./debug/<ID>.png
    dbg_dir = os.path.join(ROOT, "debug")
    os.makedirs(dbg_dir, exist_ok=True)
    dbg_png = os.path.join(dbg_dir, f"{args.id}.png")
    title = f"{args.id}  IoU={best_iou:.3f}  rot={best_deg}"
    rt.save_debug_overlay(out_path, best_rsp, dbg_png, title=title)

    print(f"[DONE] wrote {dbg_png}")


if __name__ == "__main__":
    main()