"""Point clouds of the report's objects, 10k and 2048 points, seeded. Writes ../pointclouds/<name>_<tag>_<n>.ply (points
only). Tags:
  clean, noise0p2, noise0p6, noise1   area-uniform samples of the clean mesh and of each noisy mesh (`make_noisy.py`:
                                      smooth bumps on the surface, so these are clean samples of a bumpy surface);
  pn0p12, pn0p6, pn1p2                PCPNet's protocol: area-uniform samples of the clean mesh, each point moved by
                                      isotropic Gaussian noise of sigma = 0.12 / 0.6 / 1.2 % of the bounding-box diagonal;
  dgrad, dstripe                      variable density (PCPNet-like), clean mesh: dgrad keeps a point with probability
                                      falling linearly from 1 to 0.1 along x; dstripe alternates stripes along y, a tenth
                                      of the extent wide, kept with probability 1 and 0.2.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/make_clouds.py [tags...]
"""
import sys
from pathlib import Path

import numpy as np
import trimesh

sys.path.insert(0, str(Path(__file__).parent))
from make_noisy import SIGMAS, label
from parts_report import OBJ, ROOT

SIZES = (10_000, 2048)
OUT = ROOT / "pointclouds"
PN = (0.0012, 0.006, 0.012)  # per-point noise, of the diagonal (PCPNet's levels)
TAGS = ["clean"] + [f"noise{label(s)}" for s in SIGMAS] + [f"pn{label(s)}" for s in PN] + ["dgrad", "dstripe"]


def keep_prob(P, tag):
    lo, hi = P.min(0), P.max(0)
    if tag == "dgrad":
        return 1.0 - 0.9 * (P[:, 0] - lo[0]) / max(hi[0] - lo[0], 1e-12)
    u = (P[:, 1] - lo[1]) / max(hi[1] - lo[1], 1e-12)
    return np.where(np.floor(u * 10) % 2 == 0, 1.0, 0.2)


def cloud(name, tag, n, seed):
    """The tag's cloud of n points."""
    if tag.startswith("noise"):
        return np.asarray(trimesh.load(ROOT / "meshes" / f"{name}_{tag}.ply", force="mesh").sample(n, seed=seed))
    g = trimesh.load(OBJ[name], force="mesh")
    if tag == "clean":
        return np.asarray(g.sample(n, seed=seed))
    rng = np.random.default_rng(seed + 1000)
    if tag.startswith("pn"):
        sigma = {f"pn{label(s)}": s for s in PN}[tag] * float(np.linalg.norm(g.extents))
        return np.asarray(g.sample(n, seed=seed)) + rng.normal(0.0, sigma, (n, 3))
    P = np.asarray(g.sample(20 * n, seed=seed))  # thinned to n by the density
    p = keep_prob(P, tag)
    idx = rng.choice(len(P), n, replace=False, p=p / p.sum())
    return P[idx]


def main(tags):
    OUT.mkdir(exist_ok=True)
    for name in OBJ:
        for tag in tags:
            for k, n in enumerate(SIZES):
                trimesh.PointCloud(cloud(name, tag, n, 10 * k + 1)).export(OUT / f"{name}_{tag}_{n}.ply")
        print(name, "done", flush=True)


if __name__ == "__main__":
    main(sys.argv[1:] or TAGS)