"""Encode-decode round trip on a few objects, CUDA. Writes decoded meshes next to this script."""
import sys, time
from pathlib import Path

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

sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp

D = Path("/home/joao/code/data")
M = D / "meshes/data"
OBJ = {
    "dewalt_battery": D / "vision/Dewalt_battery.ply",
    "dress": M / "MANNEQUIN_SUMMER_DRESS_-_ORANGE.ply",
    "bunny": M / "bun_zipper.ply",
    "bowl_of": M / "objectfolder/16k/2/model.obj",
    "bowl_gso": M / "meshes_clean/gso__Threshold_Bead_Cereal_Bowl_White.ply",
    "mug": M / "meshes_clean/gso__Room_Essentials_Mug_White_Yellow.ply",
    "dino": M / "meshes_clean/gso__Schleich_Therizinosaurus_ln9cruulPqc.ply",
    "shoe": M / "meshes_clean/gso__Timberland_Mens_Classic_2Eye_Boat_Shoe.ply",
    "loaf_pan": M / "objectfolder/16k/38/model.obj",
    "horse": M / "meshes_clean/gso__Breyer_Horse_Of_The_Year_2015.ply",
}
OUT = Path(__file__).parent / "decoded"
OUT.mkdir(exist_ok=True)


def load(p):
    g = trimesh.load(p, force="mesh", process=True)
    if isinstance(g, trimesh.PointCloud) or len(getattr(g, "faces", [])) == 0:
        return "points", np.asarray(g.vertices), None
    if g.is_watertight:
        return "mesh", g, g
    return "points", np.asarray(g.sample(200_000, seed=0) if "seed" in g.sample.__code__.co_varnames else g.sample(200_000)), g


def surf(x, n=50_000):
    if isinstance(x, trimesh.Trimesh):
        return np.asarray(x.sample(n))
    x = np.asarray(x)
    return x[np.random.default_rng(0).choice(len(x), min(n, len(x)), replace=False)]


print(fcp.__file__, "device:", fcp.get_device())
print(f"{'object':16s} {'input':6s} {'atoms':>5s} {'bytes':>6s} {'enc s':>6s} {'dec s':>6s} {'mean/diag':>9s} {'p95/diag':>9s}")
for name, p in OBJ.items():
    kind, geom, ref = load(p)
    ref_pts = surf(ref if ref is not None else geom)
    diag = np.linalg.norm(np.ptp(ref_pts, axis=0))
    fcp.encode(geom, device="cuda")  # warm-up (kernel compilation)
    t = time.perf_counter(); d = fcp.encode(geom, device="cuda"); te = time.perf_counter() - t
    fcp.decode(d, device="cuda")
    t = time.perf_counter(); m = fcp.decode(d, device="cuda"); td = time.perf_counter() - t
    if m is None:
        print(f"{name:16s} {kind:6s} decode empty"); continue
    m.export(OUT / f"{name}.ply")
    q = surf(m)
    a = cKDTree(q).query(ref_pts)[0]; b = cKDTree(ref_pts).query(q)[0]
    e = np.concatenate([a, b]) / diag
    print(f"{name:16s} {kind:6s} {d.n_atoms if hasattr(d, 'n_atoms') else len(d):5d} {d.nbytes:6d} {te:6.3f} {td:6.3f} "
          f"{e.mean():9.5f} {np.percentile(e, 95):9.5f}")