"""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("/home/joao/code/data/meshes/fcp_2710_tests")
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)]
for name, p in OBJ.items():
    kind, geom, ref = load(p)
    d = fcp.encode(geom, device="cuda")
    fcp.save(d, OUT / f"{name}.fcp")
    d2 = fcp.load(OUT / f"{name}.fcp")
    assert d2.to_bytes() == d.to_bytes()
    m = fcp.decode(d2, device="cuda")
    m.export(OUT / f"{name}.ply")
    old = trimesh.load(Path(__file__).parent / "decoded" / f"{name}.ply", process=False)
    same = old.vertices.shape == m.vertices.shape and np.array_equal(np.float32(old.vertices), np.float32(m.vertices))
    print(f"{name:16s} {kind:6s} {d.nbytes:6d} B  decoded mesh identical to the earlier run: {same}")