#!/usr/bin/env python3
# filepath: decimate_splat.py
"""
Decimate a Gaussian Splat PLY file and output as PLY or SPLAT format.
"""

import argparse
import numpy as np
from io import BytesIO
from pathlib import Path
from plyfile import PlyData, PlyElement


def load_gaussian_ply(ply_file_path: str) -> PlyData:
    """Load a Gaussian splat PLY file."""
    return PlyData.read(ply_file_path)


def compute_importance_scores(vert) -> np.ndarray:
    """
    Compute importance scores for each Gaussian.
    Higher scores = more important (larger and more opaque).
    """
    scales = np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
    opacities = 1 / (1 + np.exp(-vert["opacity"]))
    return scales * opacities


def decimate_ply(plydata: PlyData, keep_ratio: float) -> PlyData:
    """
    Decimate the PLY data by keeping only a fraction of the Gaussians.
    Keeps the most important Gaussians based on scale and opacity.
    """
    vert = plydata["vertex"]
    total_points = len(vert.data)
    keep_count = max(1, int(total_points * keep_ratio))
    
    # Compute importance and get indices of top Gaussians
    importance = compute_importance_scores(vert)
    sorted_indices = np.argsort(-importance)[:keep_count]
    
    # Sort indices to maintain some spatial coherence
    sorted_indices = np.sort(sorted_indices)
    
    # Create new vertex data with only kept points
    new_vertex_data = vert.data[sorted_indices]
    
    # Create new PlyElement and PlyData
    new_vertex_element = PlyElement.describe(new_vertex_data, "vertex")
    new_plydata = PlyData([new_vertex_element])
    
    return new_plydata


def convert_ply_to_splat(plydata: PlyData) -> bytes:
    """
    Convert PLY data to SPLAT format for the antimatter15 viewer.
    Returns the splat data as bytes.
    """
    vert = plydata["vertex"]
    
    sorted_indices = np.argsort(
        -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
        / (1 + np.exp(-vert["opacity"]))
    )
    
    buffer = BytesIO()
    for idx in sorted_indices:
        v = plydata["vertex"][idx]
        position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32)
        scales = np.exp(
            np.array([v["scale_0"], v["scale_1"], v["scale_2"]], dtype=np.float32)
        )
        color = np.array([
            0.5 + 0.28209479177387814 * v["f_dc_0"],
            0.5 + 0.28209479177387814 * v["f_dc_1"],
            0.5 + 0.28209479177387814 * v["f_dc_2"],
            1 / (1 + np.exp(-v["opacity"])),
        ])
        rot = np.array([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], dtype=np.float32)
        buffer.write(position.tobytes())
        buffer.write(scales.tobytes())
        buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
        buffer.write(
            ((rot / np.linalg.norm(rot)) * 128 + 128).clip(0, 255).astype(np.uint8).tobytes()
        )
    
    return buffer.getvalue()


def main():
    parser = argparse.ArgumentParser(
        description="Decimate a Gaussian Splat PLY file and output as PLY or SPLAT format."
    )
    parser.add_argument(
        "input",
        type=str,
        help="Input PLY file path"
    )
    parser.add_argument(
        "-o", "--output",
        type=str,
        help="Output file path (default: input_decimated.ply or .splat)"
    )
    parser.add_argument(
        "-r", "--ratio",
        type=float,
        default=0.5,
        help="Ratio of points to keep (0.0-1.0, default: 0.5)"
    )
    parser.add_argument(
        "-f", "--format",
        type=str,
        choices=["ply", "splat"],
        default="ply",
        help="Output format: 'ply' or 'splat' (default: ply)"
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Print verbose output"
    )
    
    args = parser.parse_args()
    
    # Validate ratio
    if not 0.0 < args.ratio <= 1.0:
        parser.error("Ratio must be between 0.0 (exclusive) and 1.0 (inclusive)")
    
    # Determine output path
    input_path = Path(args.input)
    if args.output:
        output_path = Path(args.output)
    else:
        suffix = ".splat" if args.format == "splat" else ".ply"
        output_path = input_path.with_stem(f"{input_path.stem}_decimated").with_suffix(suffix)
    
    if args.verbose:
        print(f"Loading: {args.input}")
    
    # Load PLY file
    plydata = load_gaussian_ply(args.input)
    original_count = len(plydata["vertex"].data)
    
    if args.verbose:
        print(f"Original point count: {original_count:,}")
    
    # Decimate
    decimated_plydata = decimate_ply(plydata, args.ratio)
    new_count = len(decimated_plydata["vertex"].data)
    
    if args.verbose:
        print(f"Decimated point count: {new_count:,} ({args.ratio * 100:.1f}%)")
    
    # Write output
    if args.format == "splat":
        splat_data = convert_ply_to_splat(decimated_plydata)
        with open(output_path, "wb") as f:
            f.write(splat_data)
    else:
        decimated_plydata.write(str(output_path))
    
    if args.verbose:
        print(f"Saved to: {output_path}")
    else:
        print(f"Decimated {original_count:,} → {new_count:,} points, saved to {output_path}")


if __name__ == "__main__":
    main()