#!/usr/bin/env python
"""Align post-hoc (non-geometry-grounded) baseline meshes (Pixal3D / RVG) to the
GT canonical [-0.5,0.5]^3 frame, PRESERVING texture/vertex colour, so the NO-ICP
appeval harness can score them fairly.

For each pred glb: compute cube-rotation + ICP alignment to the GT canon mesh
(align.align_mesh on a geometry copy), then apply that same transform to the
FULL textured mesh (apply_alignment only touches .vertices) and export.

Usage:
  python align_baselines.py --pred-dir DIR --gt-dir DIR --out DIR \
      --selection SEL.json --limit 30
"""
from __future__ import annotations
import argparse, json, sys
from pathlib import Path

import numpy as np
import trimesh

sys.path.insert(0, str(Path(__file__).resolve().parent))
from faithfulness import voxelize_points
from align import align_mesh, apply_alignment


def gt_voxels(gt_mesh: trimesh.Trimesh, n: int = 64, samples: int = 200_000):
    """Full occupancy of the GT canon mesh (already in [-0.5,0.5]^3)."""
    pts = gt_mesh.sample(samples)
    return voxelize_points(np.asarray(pts), n)


def load_textured(path: Path) -> trimesh.Trimesh:
    m = trimesh.load(str(path), force="mesh", process=False)
    return m


def geom_copy(m: trimesh.Trimesh) -> trimesh.Trimesh:
    return trimesh.Trimesh(np.asarray(m.vertices), np.asarray(m.faces),
                           process=False)


def resolve_gt(gt_dir: Path, obj: str):
    for c in (gt_dir / f"{obj}_canon.glb", gt_dir / obj / "mesh.glb",
              gt_dir / f"{obj}.glb"):
        if c.exists():
            return c
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pred-dir", type=Path, required=True)
    ap.add_argument("--gt-dir", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--selection", type=Path, required=True)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--n", type=int, default=64)
    ap.add_argument("--no-icp", dest="icp", action="store_false", default=True,
                    help="AXIS-ONLY alignment: unit-cube normalization + best-of-24 "
                         "discrete octahedral orientation (voxel-F1 scored), and SKIP "
                         "the trimesh ICP R/t/scale refinement. Default keeps ICP on "
                         "(existing behaviour unchanged).")
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    args = ap.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    sel = json.loads(args.selection.read_text())["selections"]
    objs = [s["object"] for s in sel]
    if args.limit:
        objs = objs[:args.limit]
    if args.nshards > 1:
        objs = objs[args.shard::args.nshards]

    report = []
    for i, obj in enumerate(objs, 1):
        pred_p = args.pred_dir / f"{obj}.glb"
        outp = args.out / f"{obj}.glb"
        if outp.exists():
            print(f"[{i}/{len(objs)}] {obj} SKIP (aligned exists)", flush=True)
            continue
        if not pred_p.exists():
            print(f"[{i}/{len(objs)}] {obj} MISSING pred", flush=True)
            report.append({"object": obj, "error": "missing pred"})
            continue
        gt_p = resolve_gt(args.gt_dir, obj)
        if gt_p is None:
            print(f"[{i}/{len(objs)}] {obj} MISSING gt", flush=True)
            report.append({"object": obj, "error": "missing gt"})
            continue
        try:
            tex = load_textured(pred_p)
            gt = trimesh.load(str(gt_p), force="mesh", process=False)
            gt_ref = gt_voxels(gt, n=args.n)
            geo = geom_copy(tex)
            ar = align_mesh(geo, gt_ref, gt, n=args.n, icp=args.icp)
            aligned = apply_alignment(tex, ar)   # touches only .vertices
            aligned.export(str(outp))
            report.append({"object": obj, "f1_raw": round(ar.f1_raw, 4),
                           "f1_icp": round(ar.f1_icp, 4),
                           "best_f1": round(ar.best_f1(), 4)})
            print(f"[{i}/{len(objs)}] {obj} OK best_f1={ar.best_f1():.3f} "
                  f"(raw={ar.f1_raw:.3f} icp={ar.f1_icp:.3f})", flush=True)
        except Exception as e:
            import traceback; traceback.print_exc()
            report.append({"object": obj, "error": f"{type(e).__name__}: {e}"})
            print(f"[{i}/{len(objs)}] {obj} FAIL {e}", flush=True)

    (args.out / "align_report.json").write_text(json.dumps(report, indent=2))
    ok = [r for r in report if "best_f1" in r]
    if ok:
        f1s = [r["best_f1"] for r in ok]
        low = [r["object"] for r in ok if r["best_f1"] < 0.4]
        print(f"\nALIGN DONE n={len(ok)} mean_best_f1={np.mean(f1s):.3f} "
              f"min={np.min(f1s):.3f} | low(<0.4): {low}")


if __name__ == "__main__":
    main()