#!/usr/bin/env python3
"""
Compare geometric scale of two GLB exports (pred vs GT) before running full metrics.

Uses trimesh (same stack as calculate_metric_3d.py). For each mesh we report
vertex count, axis-aligned bbox min/max, extents, diagonal, center; then
pred/GT ratios for diagonal and each axis extent, and center delta.

Batch mode mirrors calculate_metric_3d.py path layout.
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import trimesh


@dataclass
class BBoxReport:
    n_vertices: int
    n_faces: int
    min_xyz: np.ndarray
    max_xyz: np.ndarray
    extents: np.ndarray
    diagonal: float
    center: np.ndarray


def load_glb_as_mesh(path: Path) -> trimesh.Trimesh:
    if not path.is_file():
        raise FileNotFoundError(str(path))
    loaded = trimesh.load(str(path), force=None, ignore_broken=True)
    if isinstance(loaded, trimesh.Scene):
        meshes: List[trimesh.Trimesh] = []
        for geom in loaded.geometry.values():
            if isinstance(geom, trimesh.Trimesh):
                meshes.append(geom)
        if not meshes:
            raise ValueError(f"No Trimesh geometry in scene: {path}")
        return trimesh.util.concatenate(meshes)
    if isinstance(loaded, trimesh.Trimesh):
        return loaded
    raise TypeError(f"Unsupported trimesh type {type(loaded)} for {path}")


def bbox_report(mesh: trimesh.Trimesh) -> BBoxReport:
    v = np.asarray(mesh.vertices, dtype=np.float64)
    if v.size == 0:
        raise ValueError("empty mesh vertices")
    vmin = v.min(axis=0)
    vmax = v.max(axis=0)
    extents = vmax - vmin
    diagonal = float(np.linalg.norm(extents))
    center = 0.5 * (vmin + vmax)
    faces = int(len(mesh.faces)) if hasattr(mesh, "faces") else 0
    return BBoxReport(
        n_vertices=int(v.shape[0]),
        n_faces=faces,
        min_xyz=vmin,
        max_xyz=vmax,
        extents=extents,
        diagonal=diagonal,
        center=center,
    )


def safe_ratio(a: float, b: float) -> float:
    if not math.isfinite(a) or not math.isfinite(b) or abs(b) < 1e-12:
        return float("nan")
    return float(a / b)


def compare_reports(pred: BBoxReport, gt: BBoxReport) -> Dict[str, Any]:
    ratios_xyz = [safe_ratio(float(pe), float(ge)) for pe, ge in zip(pred.extents, gt.extents)]
    return {
        "diagonal_ratio_pred_over_gt": safe_ratio(pred.diagonal, gt.diagonal),
        "extent_ratio_xyz_pred_over_gt": ratios_xyz,
        "center_delta_pred_minus_gt": (pred.center - gt.center).tolist(),
    }


def print_report(tag: str, rep: BBoxReport) -> None:
    print(f"\n[{tag}] {rep.n_vertices} verts, {rep.n_faces} faces")
    print(f"  min xyz: {rep.min_xyz}")
    print(f"  max xyz: {rep.max_xyz}")
    print(f"  extents (Lx,Ly,Lz): {rep.extents}")
    print(f"  diagonal: {rep.diagonal:.6f}")
    print(f"  center:   {rep.center}")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Compare pred vs GT GLB bounding-box scale.")
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--pred_glb", type=Path, help="Single predicted GLB path.")
    g.add_argument(
        "--eval_results_dir",
        type=Path,
        help="Eval output root (batch); pairs built like calculate_metric_3d.py.",
    )

    p.add_argument("--gt_glb", type=Path, help="Single GT GLB path (use with --pred_glb).")
    p.add_argument("--gt_dir", type=Path, help="GT root for batch mode (use with --eval_results_dir).")
    p.add_argument(
        "--dataset_root_dir",
        type=Path,
        default=Path("/mnt/zsn/data/3DEditVerse"),
        help="Used to locate test_data_info.json in batch mode.",
    )
    p.add_argument(
        "--test_data_info_path",
        type=Path,
        default=None,
        help="Override test_data_info.json path (default: <dataset_root_dir>/test_data_info.json).",
    )
    p.add_argument("--wo_mixamo", action="store_true", help="Skip mixamo keys in batch mode.")
    p.add_argument(
        "--warn_ratio",
        type=float,
        default=1.05,
        help="Warn if diagonal or any extent ratio falls outside [1/r, r] (default 1.05).",
    )
    p.add_argument("--max_samples", type=int, default=None, help="Optional cap for batch mode.")
    p.add_argument("--save_json", type=Path, default=None, help="Optional path to write batch summary JSON.")
    return p.parse_args()


def _batch_pairs(args: argparse.Namespace) -> List[Tuple[str, Path, Path]]:
    if args.gt_dir is None:
        raise SystemExit("batch mode requires --gt_dir")
    info_path = args.test_data_info_path or (args.dataset_root_dir / "test_data_info.json")
    with open(info_path, "r", encoding="utf-8") as f:
        test_data_info = json.load(f)
    pairs: List[Tuple[str, Path, Path]] = []
    for key in test_data_info.get("flux_edit", []):
        pred = args.eval_results_dir / "flux_edit" / key / "edit.glb"
        gt = args.gt_dir / "flux_edit" / key / "edit.glb"
        pairs.append((f"flux_edit/{key}", pred, gt))
    for key in test_data_info.get("alpaca", []):
        pred = args.eval_results_dir / "alpaca" / key / "edit.glb"
        gt = args.gt_dir / "alpaca" / key / "edit.glb"
        pairs.append((f"alpaca/{key}", pred, gt))
    if not args.wo_mixamo:
        with open(args.dataset_root_dir / "dataset_info.json", "r", encoding="utf-8") as f:
            all_data_info = json.load(f)
        for mixamo_key in test_data_info.get("mixamo", []):
            object_name, ori_idx, edit_idx = mixamo_key
            pred_key = f"{object_name}_{ori_idx}_{edit_idx}"
            gt_key = all_data_info["mixamo"][object_name][int(edit_idx)]["ss_latents_path"].split("/")[-1][:-4]
            pred = args.eval_results_dir / "mixamo" / pred_key / "edit.glb"
            gt = args.gt_dir / "mixamo" / gt_key / "ori.glb"
            pairs.append((f"mixamo/{pred_key}", pred, gt))
    if args.max_samples is not None:
        pairs = pairs[: args.max_samples]
    return pairs


def main() -> None:
    args = parse_args()
    warn_r = float(args.warn_ratio)
    if warn_r <= 1.0:
        raise SystemExit("--warn_ratio must be > 1")

    if args.pred_glb is not None:
        if args.gt_glb is None:
            raise SystemExit("--pred_glb requires --gt_glb")
        pred_mesh = load_glb_as_mesh(args.pred_glb)
        gt_mesh = load_glb_as_mesh(args.gt_glb)
        pr = bbox_report(pred_mesh)
        gr = bbox_report(gt_mesh)
        print_report("PRED", pr)
        print_report("GT", gr)
        cmp = compare_reports(pr, gr)
        print("\n[COMPARE]")
        print(json.dumps(cmp, indent=2))
        diag_r = cmp["diagonal_ratio_pred_over_gt"]
        bad = False
        if math.isfinite(diag_r) and (diag_r > warn_r or diag_r < 1.0 / warn_r):
            print(f"WARNING: diagonal ratio {diag_r:.6f} outside [{1/warn_r:.4f}, {warn_r:.4f}]")
            bad = True
        for i, r in enumerate(cmp["extent_ratio_xyz_pred_over_gt"]):
            if math.isfinite(r) and (r > warn_r or r < 1.0 / warn_r):
                print(f"WARNING: extent ratio axis {i} pred/gt = {r:.6f} outside [{1/warn_r:.4f}, {warn_r:.4f}]")
                bad = True
        if bad:
            raise SystemExit(2)
        print("\nOK: pred and GT appear to share a similar axis-aligned scale (within warn_ratio).")
        return

    pairs = _batch_pairs(args)
    rows: List[Dict[str, Any]] = []
    worst_dev = 1.0
    worst_key = ""
    missing = 0
    failed = 0
    for key, pred_path, gt_path in pairs:
        if not pred_path.is_file() or not gt_path.is_file():
            missing += 1
            rows.append({"key": key, "status": "missing_file", "pred": str(pred_path), "gt": str(gt_path)})
            continue
        try:
            pr = bbox_report(load_glb_as_mesh(pred_path))
            gr = bbox_report(load_glb_as_mesh(gt_path))
            cmp = compare_reports(pr, gr)
            diag_r = cmp["diagonal_ratio_pred_over_gt"]
            row: Dict[str, Any] = {
                "key": key,
                "status": "ok",
                "pred": str(pred_path),
                "gt": str(gt_path),
                "pred_diag": pr.diagonal,
                "gt_diag": gr.diagonal,
                **cmp,
            }
            extent_rs = cmp["extent_ratio_xyz_pred_over_gt"]
            bad = False
            if math.isfinite(diag_r) and (diag_r > warn_r or diag_r < 1.0 / warn_r):
                bad = True
            for r in extent_rs:
                if math.isfinite(r) and (r > warn_r or r < 1.0 / warn_r):
                    bad = True
            if bad:
                row["status"] = "scale_mismatch"
                failed += 1
            rows.append(row)
            if math.isfinite(diag_r) and diag_r > 0:
                dev = max(diag_r, 1.0 / diag_r)
                if dev > worst_dev:
                    worst_dev = dev
                    worst_key = key
        except Exception as exc:
            failed += 1
            rows.append({"key": key, "status": "error", "error": repr(exc), "pred": str(pred_path), "gt": str(gt_path)})

    finite_diags = [
        r["diagonal_ratio_pred_over_gt"]
        for r in rows
        if r.get("status") == "ok" and math.isfinite(float(r.get("diagonal_ratio_pred_over_gt", float("nan"))))
    ]
    summary = {
        "n_pairs": len(pairs),
        "missing_files": missing,
        "scale_mismatch_or_error": failed,
        "warn_ratio": warn_r,
        "diagonal_ratio_median": float(np.median(finite_diags)) if finite_diags else None,
        "worst_symmetric_diag_deviation": worst_dev,
        "worst_key": worst_key,
    }
    print(json.dumps(summary, indent=2))
    if args.save_json is not None:
        args.save_json.parent.mkdir(parents=True, exist_ok=True)
        with open(args.save_json, "w", encoding="utf-8") as f:
            json.dump({"summary": summary, "rows": rows}, f, indent=2)
        print(f"wrote {args.save_json}")

    if missing:
        print(f"NOTE: {missing} pairs missing pred or gt GLB on disk.", file=sys.stderr)
    if failed:
        raise SystemExit(2)


if __name__ == "__main__":
    main()