#!/usr/bin/env python3
"""3DGS-PLY → .splat (antimatter15-Format, 32 Bytes/Splat) für krpano.

Usage: ply2splat.py input.ply [output.splat]
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, "/home/panomity.de/sharp")
from sharp_batch import read_gs_ply  # noqa: E402

SH_C0 = 0.28209479177387814


def main() -> None:
    src = Path(sys.argv[1])
    dst = Path(sys.argv[2]) if len(sys.argv) > 2 else src.with_suffix(".splat")

    names, data = read_gs_ply(src)
    idx = {n: i for i, n in enumerate(names)}

    pos = data[:, [idx["x"], idx["y"], idx["z"]]].astype(np.float32)
    scales = np.exp(data[:, [idx["scale_0"], idx["scale_1"], idx["scale_2"]]]).astype(np.float32)

    rgb = 0.5 + SH_C0 * data[:, [idx["f_dc_0"], idx["f_dc_1"], idx["f_dc_2"]]]
    alpha = 1.0 / (1.0 + np.exp(-data[:, idx["opacity"]]))
    rgba = np.clip(np.concatenate([rgb, alpha[:, None]], axis=1) * 255.0, 0, 255).astype(np.uint8)

    quat = data[:, [idx["rot_0"], idx["rot_1"], idx["rot_2"], idx["rot_3"]]]
    quat = quat / np.linalg.norm(quat, axis=1, keepdims=True)
    quat_u8 = np.clip(quat * 128.0 + 128.0, 0, 255).astype(np.uint8)

    # Nach Wichtigkeit sortieren (Volumen × Deckkraft), wie üblich beim .splat-Format.
    importance = scales.prod(axis=1) * alpha
    order = np.argsort(-importance)

    n = data.shape[0]
    out = np.empty(n, dtype=[("pos", np.float32, 3), ("scale", np.float32, 3),
                             ("rgba", np.uint8, 4), ("quat", np.uint8, 4)])
    out["pos"] = pos[order]
    out["scale"] = scales[order]
    out["rgba"] = rgba[order]
    out["quat"] = quat_u8[order]
    out.tofile(dst)
    print(f"{dst}: {n} Splats, {dst.stat().st_size/1e6:.1f} MB")


if __name__ == "__main__":
    main()