#!/usr/bin/env python3
"""Rewrite a SUM-schema PLY into the spelling desktop viewers expect.

The training files carry colour as r/g/b float32 in [0,1] — that is what the
SUM Parts loader reads and what the network was trained on. Most mesh/point
viewers (CloudCompare, MeshLab, Mapple) look for red/green/blue uint8 instead,
and when they don't find it they either render the cloud flat grey or list the
floats as scalar fields.

This converts colour spelling only. Coordinates and labels pass through
untouched, so the geometry you inspect is exactly the geometry the model saw.

Usage:
    python ply_for_viewer.py in.ply out.ply
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
from plyfile import PlyData, PlyElement


def main() -> None:
    if len(sys.argv) != 3:
        print(__doc__)
        raise SystemExit(2)

    src, dst = Path(sys.argv[1]), Path(sys.argv[2])
    if not src.exists():
        raise SystemExit(f"{src}: not found")

    v = PlyData.read(str(src))["vertex"]
    props = [p.name for p in v.properties]
    print(f"in  : {src.name}  {len(v):,} points  props={props}")

    rgb_set = next((s for s in (("r", "g", "b"), ("red", "green", "blue"))
                    if all(c in props for c in s)), None)
    if rgb_set is None:
        raise SystemExit("no colour channels found")

    rgb = np.stack([np.asarray(v[c], dtype=np.float64) for c in rgb_set], axis=1)
    if rgb.max() <= 1.0:
        rgb = rgb * 255.0
        print("      colour was float 0-1 -> scaling to 0-255")
    rgb = np.clip(rgb, 0, 255).astype(np.uint8)

    dtype = [("x", "f4"), ("y", "f4"), ("z", "f4"),
             ("red", "u1"), ("green", "u1"), ("blue", "u1")]
    if "label" in props:
        dtype.append(("label", "i4"))

    arr = np.empty(len(v), dtype=dtype)
    arr["x"], arr["y"], arr["z"] = v["x"], v["y"], v["z"]
    arr["red"], arr["green"], arr["blue"] = rgb[:, 0], rgb[:, 1], rgb[:, 2]
    if "label" in props:
        arr["label"] = np.asarray(v["label"]).astype(np.int32)

    dst.parent.mkdir(parents=True, exist_ok=True)
    PlyData([PlyElement.describe(arr, "vertex")], text=False).write(str(dst))

    print(f"out : {dst}  ({dst.stat().st_size / 1e6:.1f} MB)")
    print(f"      mean RGB {rgb.mean(axis=0).round(1).tolist()}")


if __name__ == "__main__":
    main()