# Modified for SplatAnything: inference-only release.
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from plyfile import PlyData, PlyElement

def save_image_png(path: Path, image_tensor: torch.Tensor) -> None:
    # image_tensor: [H, W, 3]
    img = (image_tensor.detach().cpu() * 255.0).to(torch.uint8).numpy()
    Image.fromarray(img).save(str(path))

def save_depth_png(path: Path, depth_tensor: torch.Tensor) -> None:
    # depth_tensor: [H, W]
    d = depth_tensor.detach()
    d = d - d.min()
    d = d / (d.max() + 1e-9)
    img = (d.clamp(0, 1) * 255.0).to(torch.uint8).cpu().numpy()
    Image.fromarray(img, mode="L").save(str(path))

def save_depth_npy(path: Path, depth_tensor: torch.Tensor) -> None:
    # depth_tensor: [H, W]
    # Save actual depth values in numpy format
    d = depth_tensor.detach().cpu().numpy()
    np.save(str(path), d)

def _build_vertex_ply_element(pts: np.ndarray, colors: np.ndarray) -> PlyElement:
    """Build a PLY vertex element from points and colors arrays.

    Args:
        pts: Point coordinates, shape [N, 3], dtype float32
        colors: RGB colors, shape [N, 3], dtype uint8

    Returns:
        PlyElement describing the vertices
    """
    vertex_dtype = [("x", "f4"), ("y", "f4"), ("z", "f4"),
                    ("red", "u1"), ("green", "u1"), ("blue", "u1")]
    vertex_elements = np.empty(len(pts), dtype=vertex_dtype)
    vertex_elements["x"] = pts[:, 0]
    vertex_elements["y"] = pts[:, 1]
    vertex_elements["z"] = pts[:, 2]
    vertex_elements["red"] = colors[:, 0]
    vertex_elements["green"] = colors[:, 1]
    vertex_elements["blue"] = colors[:, 2]
    return PlyElement.describe(vertex_elements, "vertex")

def save_scene_ply(path: Path,
                   points_xyz: torch.Tensor,
                   point_colors: torch.Tensor,
                   valid_mask: torch.Tensor = None) -> None:
    """Save point cloud to PLY format"""
    pts = points_xyz.detach().cpu().to(torch.float32).numpy().reshape(-1, 3)
    colors = point_colors.detach().cpu().to(torch.uint8).numpy().reshape(-1, 3)

    # Filter out invalid points (NaN, Inf)
    if valid_mask is None:
        valid_mask = np.isfinite(pts).all(axis=1)
    else:
        valid_mask = valid_mask.detach().cpu().numpy().reshape(-1)
    pts = pts[valid_mask]
    colors = colors[valid_mask]

    # Handle empty point cloud
    if len(pts) == 0:
        pts = np.array([[0, 0, 0]], dtype=np.float32)
        colors = np.array([[255, 255, 255]], dtype=np.uint8)

    PlyData([_build_vertex_ply_element(pts, colors)]).write(str(path))

def _build_gs_ply_data(means, scales, rotations, rgbs, opacities, quantile_threshold):
    """Build Gaussian splat PLY data with scale-based filtering.

    Args:
        means: Gaussian centers [N, 3]
        scales: Gaussian scales [N, 3]
        rotations: Gaussian rotations as quaternions [N, 4]
        rgbs: RGB colors [N, 3]
        opacities: Opacity values [N]
        quantile_threshold: Percentile threshold for scale filtering (e.g. 0.98 or 0.90).
            None -> no filtering, export ALL Gaussians (inspect the full cloud).

    Returns:
        PlyData object ready to be written or returned
    """
    # None -> keep every Gaussian (no scale-based pruning); else drop the largest-scale tail.
    if quantile_threshold is None:
        filter_mask = torch.ones(scales.shape[0], dtype=torch.bool, device=scales.device)
    else:
        scale_threshold = torch.quantile(scales.max(dim=-1)[0], quantile_threshold, dim=0)
        filter_mask = scales.max(dim=-1)[0] <= scale_threshold

    means = means[filter_mask].reshape(-1, 3)
    scales = scales[filter_mask].reshape(-1, 3)
    rotations = rotations[filter_mask].reshape(-1, 4)
    rgbs = rgbs[filter_mask].reshape(-1, 3)
    opacities = opacities[filter_mask].reshape(-1)

    attributes = ["x", "y", "z", "nx", "ny", "nz"]
    for i in range(3):
        attributes.append(f"f_dc_{i}")
    attributes.append("opacity")
    for i in range(3):
        attributes.append(f"scale_{i}")
    for i in range(4):
        attributes.append(f"rot_{i}")

    dtype_full = [(attribute, "f4") for attribute in attributes]
    elements = np.empty(means.shape[0], dtype=dtype_full)

    # 3DGS / SuperSplat PLY storage convention (see process_ply_to_splat):
    #   color   = 0.5 + C0 * f_dc      -> store the SH DC coefficient (rgbs arg)
    #   opacity = sigmoid(opacity)     -> store the logit (inverse_sigmoid)
    #   scale   = exp(scale)           -> store log(scale)
    # rgbs is the SH DC coefficient (NOT activated RGB); callers pass sh[..., 0, :].
    opacity_logit = torch.logit(opacities.clamp(1e-6, 1 - 1e-6))
    attributes_data = (
        means.float().detach().cpu().numpy(),
        torch.zeros_like(means).float().detach().cpu().numpy(),
        rgbs.detach().cpu().contiguous().numpy(),
        opacity_logit[..., None].detach().cpu().numpy(),
        scales.log().detach().cpu().numpy(),
        rotations.detach().cpu().numpy(),
    )
    attributes_data = np.concatenate(attributes_data, axis=1)
    elements[:] = list(map(tuple, attributes_data))

    return PlyData([PlyElement.describe(elements, "vertex")])

def save_gs_ply(path: Path,
                means: torch.Tensor,
                scales: torch.Tensor,
                rotations: torch.Tensor,
                rgbs: torch.Tensor,
                opacities: torch.Tensor,
                quantile_threshold: float = 0.98) -> None:
    """
    Export Gaussian splat data to PLY format.

    Args:
        path: Output PLY file path
        means: Gaussian centers [N, 3]
        scales: Gaussian scales [N, 3]
        rotations: Gaussian rotations as quaternions [N, 4]
        rgbs: RGB colors [N, 3]
        opacities: Opacity values [N]
        quantile_threshold: scale-tail pruning quantile (default 0.98 drops the largest 2%
            of Gaussians by scale). Pass None to export the FULL cloud with no filtering.
    """
    # Ensure float32 for quantile and numpy conversion (bf16 not supported)
    means, scales, rotations, rgbs, opacities = (
        t.float() for t in (means, scales, rotations, rgbs, opacities)
    )
    plydata = _build_gs_ply_data(means, scales, rotations, rgbs, opacities,
                                 quantile_threshold=quantile_threshold)
    plydata.write(str(path))