#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os, sys, json, csv, math
from pathlib import Path

# 让 matplotlib 不弹窗
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# 引入你的 geometry.py
sys.path.append("/Users/zongyikun/Desktop/qwen_vl_demo")
import geometry  # 你刚才贴的那个 geometry.py

def mean_mse(params):
    vals = [float(p.get("mse_fit", 0.0)) for p in params if isinstance(p, dict)]
    return sum(vals) / len(vals) if vals else None

def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--params_dir", required=True, help="params_from_svg 目录")
    ap.add_argument("--outlines_dir", required=False, help="outline PNG 目录（可选）")
    ap.add_argument("--out_csv", required=True, help="结果 CSV 路径")
    ap.add_argument("--vis_dir", required=False, help="可视化输出目录（可选）")
    ap.add_argument("--no_iou", action="store_true", help="不算 IoU（无 outlines 时可用）")
    ap.add_argument("--raster", type=int, default=160, help="IoU 光栅分辨率")
    args = ap.parse_args()

    params_dir = Path(args.params_dir)
    outlines_dir = Path(args.outlines_dir) if args.outlines_dir else None
    vis_dir = Path(args.vis_dir) if args.vis_dir else None
    if vis_dir:
        vis_dir.mkdir(parents=True, exist_ok=True)

    rows = []
    files = sorted(params_dir.glob("*.json"))
    print(f"[INFO] found {len(files)} param file(s) under {params_dir}")

    for jp in files:
        id_ = jp.stem
        try:
            params = json.loads(jp.read_text(encoding="utf-8"))
        except Exception as e:
            print(f"[SKIP read error] {id_}: {e}")
            continue

        # 渲染并（可选）叠加 outline
        outline_path = None
        if outlines_dir:
            cand = outlines_dir / f"{id_}.png"
            if cand.exists():
                outline_path = str(cand)

        # 用 geometry 生成可视化（如需）
        if vis_dir:
            # 复用 geometry 的绘图但保存图片
            fig = plt.figure()
            try:
                geometry.draw_from_annotation_json(str(jp), outline_png_path=outline_path, title=id_)
                out_png = vis_dir / f"{id_}.png"
                plt.savefig(out_png, dpi=150, bbox_inches="tight")
            finally:
                plt.close(fig)

        # 计算 IoU（如需）
        if not args.no_iou and outline_path:
            # geometry.draw_from_annotation_json 已经会 apply 变换并返回 response
            resp = geometry.draw_from_annotation_json(str(jp), outline_png_path=None, title=None)
            plt.close('all')  # 不要窗口
            iou = geometry.approx_iou_with_outline(resp, outline_path, raster=args.raster)
        else:
            iou = None

        mse_mean = mean_mse(params)
        rows.append({"id": id_, "iou": iou, "mse_mean": mse_mean})

    # 写 CSV
    out_csv = Path(args.out_csv)
    with out_csv.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["id", "iou", "mse_mean"])
        w.writeheader()
        for r in rows:
            w.writerow(r)

    print(f"[DONE] wrote {len(rows)} rows -> {out_csv}")
    if vis_dir:
        print(f"[VIS] saved previews -> {vis_dir}")

if __name__ == "__main__":
    main()