"""Shape metrics against ShapeNetCore.v2.PC15k GT, following MinD-3D's tools/get_cd.py protocol
(2048 points, pc_norm), plus a shared orientation search so methods with different canonical frames
are compared fairly."""
import os
import glob
import json
import argparse
import collections

import numpy as np
import trimesh
from scipy.spatial import cKDTree
from scipy.optimize import linear_sum_assignment

N_POINTS = 2048
FSCORE_TAU = 0.02


def pc_norm(pc):
    pc = pc - pc.mean(0)
    return pc / (2 * np.max(np.linalg.norm(pc, axis=1)))


def rot_y(deg):
    t = np.deg2rad(deg)
    c, s = np.cos(t), np.sin(t)
    return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])


UP_MAPS = {
    "y_up": np.eye(3),
    "z_up": np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]], dtype=np.float64),
}
CANDIDATES = [(up, az, rot_y(az) @ m) for up, m in UP_MAPS.items() for az in range(0, 360, 45)]


def chamfer(a, b):
    d_ab = cKDTree(b).query(a)[0]
    d_ba = cKDTree(a).query(b)[0]
    return np.mean(d_ab ** 2) + np.mean(d_ba ** 2), d_ab, d_ba


def fscore(d_ab, d_ba, tau):
    p = np.mean(d_ab < tau)
    r = np.mean(d_ba < tau)
    return 0.0 if p + r == 0 else 2 * p * r / (p + r)


def emd(a, b):
    cost = np.linalg.norm(a[:, None] - b[None], axis=-1)
    r, c = linear_sum_assignment(cost)
    return cost[r, c].mean()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gen_dir", required=True)
    parser.add_argument("--gt_root", default="/home/hubin/data/ShapeNetCore.v2.PC15k/ShapeNetCore.v2.PC15k")
    parser.add_argument("--test_list", default="/home/hubin/data/fMRI-Shape/annotations/core_test_list.txt")
    parser.add_argument("--out", required=True)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    gt_paths = {os.path.basename(p)[:-4]: p for p in glob.glob(f"{args.gt_root}/*/_/*/*.npy")}
    ids = [l.strip() for l in open(args.test_list) if l.strip()]
    rows = []
    for obj in ids:
        cat, uid = obj.split("/")
        gen = [p for p in glob.glob(os.path.join(args.gen_dir, f"{cat}_{uid}.*")) if p.endswith((".ply", ".obj"))]
        if uid not in gt_paths or not gen:
            continue
        rng = np.random.default_rng(args.seed)
        gt = np.load(gt_paths[uid])
        gt = pc_norm(gt[rng.choice(len(gt), N_POINTS, replace=False)].astype(np.float64))
        mesh = trimesh.load(gen[0], force="mesh")
        pts = trimesh.sample.sample_surface(mesh, N_POINTS, seed=args.seed)[0].astype(np.float64)
        pts = pc_norm(pts)

        best = None
        for up, az, R in CANDIDATES:
            cd, d_ab, d_ba = chamfer(pts @ R.T, gt)
            if best is None or cd < best[0]:
                best = (cd, up, az, R, d_ab, d_ba)
        cd, up, az, R, d_ab, d_ba = best
        rows.append({
            "id": obj, "cat": cat, "cd": cd, "emd": emd(pts @ R.T, gt),
            "fscore": fscore(d_ab, d_ba, FSCORE_TAU), "up": up, "azimuth": az,
        })
        print(f"{obj} CD={cd * 1e3:.3f}e-3 EMD={rows[-1]['emd']:.4f} F={rows[-1]['fscore']:.3f} ({up},{az})", flush=True)

    by_cat = collections.defaultdict(list)
    for r in rows:
        by_cat[r["cat"]].append(r)
    summary = {
        "n": len(rows),
        "cd_x1e3": float(np.mean([r["cd"] for r in rows]) * 1e3),
        "emd": float(np.mean([r["emd"] for r in rows])),
        "fscore@0.02": float(np.mean([r["fscore"] for r in rows])),
        "per_category": {c: {"n": len(v), "cd_x1e3": float(np.mean([r["cd"] for r in v]) * 1e3),
                             "emd": float(np.mean([r["emd"] for r in v]))} for c, v in sorted(by_cat.items())},
        "chosen_orientation": collections.Counter(f"{r['up']}/{r['azimuth']}" for r in rows).most_common(),
    }
    with open(args.out, "w") as f:
        json.dump({"summary": summary, "rows": rows}, f, indent=1)
    print(json.dumps({k: v for k, v in summary.items() if k != "per_category"}, indent=1))


if __name__ == "__main__":
    main()