"""Volume from a point cloud: FCP (encode from the points, default slicer, then the descriptor's volume) against
screened Poisson reconstruction (Open3D, true normals, depth 8 and 9, then the mesh's volume) and the convex hull (an
upper bound). The same 20k area-uniform points per object (seed 3, saved in out/pc/).

    env -u PYTHONPATH ~/code/fcp_devel/.venv/bin/python scripts/volume_pc.py points
    env -u PYTHONPATH ~/code/fcp_devel/.venv/bin/python scripts/volume_pc.py fcp
    env -u PYTHONPATH ~/code/data/meshes/fcp_2910_collision/.venv/bin/python scripts/volume_pc.py poisson
    env -u PYTHONPATH ~/code/fcp_devel/.venv/bin/python scripts/volume_pc.py table
"""
import json
import sys
import time
from pathlib import Path

import numpy as np

step = sys.argv[1]
sys.argv = [sys.argv[0], "x"]
sys.path.insert(0, str(Path(__file__).parent))
ROOT = Path(__file__).resolve().parent.parent
PC = ROOT / "out/pc"
N = 20000


def med(f, n=5):
    f()
    t = []
    for _ in range(n):
        t0 = time.perf_counter()
        f()
        t.append(time.perf_counter() - t0)
    return 1e3 * float(np.median(t))


def vol(V, F):
    a, b, c = V[F[:, 0]], V[F[:, 1]], V[F[:, 2]]
    return float(np.einsum("ij,ij->", a, np.cross(b, c)) / 6)


if step == "points":
    import trimesh
    from hull import OBJECTS

    PC.mkdir(parents=True, exist_ok=True)
    for name, path in OBJECTS.items():
        m = trimesh.load(path, force="mesh")
        P, fi = trimesh.sample.sample_surface(m, N, seed=3)
        np.savez(PC / f"{name}.npz", P=P, N=m.face_normals[fi], vol=m.volume)
elif step == "fcp":
    from volume import volume  # puts the package on the path
    import fcp

    res = {}
    for f in sorted(PC.glob("*.npz")):
        z = np.load(f)
        P = np.asarray(z["P"], np.float32)
        for atoms in (1024, 6144):
            d = fcp.encode(P, atoms=atoms, device="cpu")
            res[f"{f.stem}/fcp_{atoms}"] = dict(vol=volume(d) / float(z["vol"]),
                                                ms=med(lambda: volume(fcp.encode(P, atoms=atoms, device="cpu"))))
        print(f.stem, res[f"{f.stem}/fcp_6144"], flush=True)
    (ROOT / "out/volume_pc_fcp.json").write_text(json.dumps(res))
elif step == "poisson":
    import open3d as o3d
    from scipy.spatial import ConvexHull

    res = {}
    for f in sorted(PC.glob("*.npz")):
        z = np.load(f)
        pc = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(z["P"]))
        pc.normals = o3d.utility.Vector3dVector(z["N"])
        for depth in (8, 9):
            def run():
                m, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pc, depth=depth, n_threads=1)
                return vol(np.asarray(m.vertices), np.asarray(m.triangles))
            res[f"{f.stem}/poisson_d{depth}"] = dict(vol=run() / float(z["vol"]), ms=med(run, 3))
        res[f"{f.stem}/hull"] = dict(vol=ConvexHull(z["P"]).volume / float(z["vol"]), ms=med(lambda: ConvexHull(z["P"]).volume))
        print(f.stem, {k.split('/')[1]: v for k, v in res.items() if k.startswith(f.stem)}, flush=True)
    (ROOT / "out/volume_pc_poisson.json").write_text(json.dumps(res))
elif step == "table":
    import pandas as pd

    r = {**json.loads((ROOT / "out/volume_pc_fcp.json").read_text()), **json.loads((ROOT / "out/volume_pc_poisson.json").read_text())}
    d = pd.DataFrame([dict(obj=k.split("/")[0], method=k.split("/")[1], err=100 * (v["vol"] - 1), ms=v["ms"]) for k, v in r.items()])
    d.to_csv(ROOT / "out/volume_pc.csv", index=False)
    print(d.pivot(index="obj", columns="method", values="err").round(2).to_string())
    print(d.pivot(index="obj", columns="method", values="ms").round(1).to_string())
    print(d.groupby("method").agg(err_med=("err", lambda x: x.abs().median()), err_max=("err", lambda x: x.abs().max()), ms_med=("ms", "median")).round(2))