#!/usr/bin/env python3
"""Convert a Gaussian-splat PLY into a lightweight colored point-cloud GLB."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import trimesh


SH_C0 = 0.28209479177387814
ALIGN_XY_SCALE = 3.672624762730334
ALIGN_YAW_RADIANS = 2.8292228430345236
ALIGN_XY_TRANSLATION = np.array([-1.4822845855746674, 0.8021416967574936], dtype=np.float32)
ALIGN_Z_SCALE = 3.5309628207501587
ALIGN_Z_OFFSET = 0.9007954323633903
VIEWER_FLOOR_OFFSET_Y = 0.41


def read_gaussian_ply(path: Path) -> tuple[np.ndarray, np.ndarray]:
    payload = path.read_bytes()
    header_end = payload.index(b"end_header\n") + len(b"end_header\n")
    header = payload[:header_end].decode("utf-8")
    vertex_count = int(next(line.split()[2] for line in header.splitlines() if line.startswith("element vertex ")))
    properties = [line.split()[2] for line in header.splitlines() if line.startswith("property ")]
    stride = len(properties) * 4
    expected_size = header_end + vertex_count * stride
    if len(payload) < expected_size:
        raise ValueError("PLY vertex payload is truncated")

    values = np.frombuffer(payload, dtype="<f4", count=vertex_count * len(properties), offset=header_end)
    values = values.reshape(vertex_count, len(properties))
    property_index = {name: index for index, name in enumerate(properties)}

    positions = values[:, [property_index["x"], property_index["y"], property_index["z"]]].copy()
    dc = values[:, [property_index["f_dc_0"], property_index["f_dc_1"], property_index["f_dc_2"]]]
    colors = np.clip((0.5 + SH_C0 * dc) * 255.0, 0, 255).astype(np.uint8)
    return positions, colors


def align_to_viewer(positions: np.ndarray) -> np.ndarray:
    # The CloudCompare registration included roughly 2.5 degrees of roll and
    # pitch. A 2D occupancy map is planar, so retain only its fitted XY yaw.
    cosine = np.cos(ALIGN_YAW_RADIANS)
    sine = np.sin(ALIGN_YAW_RADIANS)
    rotation = np.array([[cosine, -sine], [sine, cosine]], dtype=np.float32)
    aligned_xy = (positions[:, :2] @ rotation.T) * ALIGN_XY_SCALE + ALIGN_XY_TRANSLATION
    aligned_z = positions[:, 2] * ALIGN_Z_SCALE + ALIGN_Z_OFFSET
    # ROS map is Z-up; the viewer uses Three.js Y-up.
    viewer_positions = np.column_stack((aligned_xy[:, 0], aligned_z, -aligned_xy[:, 1])).astype(np.float32)
    viewer_positions[:, 1] += VIEWER_FLOOR_OFFSET_Y
    return viewer_positions


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--max-points", type=int, default=398_137)
    parser.add_argument("--seed", type=int, default=731)
    args = parser.parse_args()

    positions, colors = read_gaussian_ply(args.input)
    positions = align_to_viewer(positions)
    if args.max_points < len(positions):
        generator = np.random.default_rng(args.seed)
        selected = generator.choice(len(positions), size=args.max_points, replace=False)
        positions = positions[selected]
        colors = colors[selected]

    rgba = np.column_stack((colors, np.full(len(colors), 255, dtype=np.uint8)))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    trimesh.points.PointCloud(positions, colors=rgba).export(args.output, file_type="glb")
    print(f"Wrote {len(positions):,} points to {args.output}")


if __name__ == "__main__":
    main()