from pathlib import Path

import numpy as np
import torch
from jaxtyping import Float
from torch import Tensor


SH_C0 = 0.28209479177387814


def export_splat(
    means: Float[Tensor, "gaussian 3"],
    scales: Float[Tensor, "gaussian 3"],
    rotations: Float[Tensor, "gaussian 4"],
    harmonics: Float[Tensor, "gaussian 3 d_sh"],
    opacities: Float[Tensor, " gaussian"],
    path: Path,
    scale_threshold: float | None = None,
) -> None:
    """Write the compact 32-byte-per-Gaussian format used by gsplat.js."""
    if harmonics.ndim == 3 and harmonics.shape[-1] == 3 and harmonics.shape[-2] != 3:
        harmonics = harmonics.transpose(-1, -2)

    if scale_threshold is not None:
        if scale_threshold <= 0:
            raise ValueError("scale_threshold must be positive")
        keep = scales.max(dim=-1).values <= scale_threshold
        if not keep.any():
            raise ValueError(
                f"No Gaussians remain below scale threshold {scale_threshold:.4f}"
            )
        means = means[keep]
        scales = scales[keep]
        rotations = rotations[keep]
        harmonics = harmonics[keep]
        opacities = opacities[keep]

    positions_np = means.detach().float().cpu().contiguous().numpy().astype("<f4")
    scales_np = scales.detach().float().cpu().contiguous().numpy().astype("<f4")

    colors = (0.5 + SH_C0 * harmonics[..., 0]).clamp(0.0, 1.0)
    rgb_np = (
        colors.detach().float().cpu().mul(255).round().to(torch.uint8).numpy()
    )
    alpha_np = (
        opacities.detach()
        .float()
        .cpu()
        .clamp(0.0, 1.0)
        .mul(255)
        .round()
        .to(torch.uint8)
        .numpy()
    )

    # OF3GS trains directly with gsplat, whose quaternion convention is wxyz.
    rotations = torch.nn.functional.normalize(rotations.float(), dim=-1)
    rotations_np = (
        rotations.detach()
        .cpu()
        .mul(128)
        .add(128)
        .clamp(0, 255)
        .round()
        .to(torch.uint8)
        .numpy()
    )

    count = positions_np.shape[0]
    rows = np.empty((count, 32), dtype=np.uint8)
    rows[:, 0:12] = positions_np.view(np.uint8).reshape(count, 12)
    rows[:, 12:24] = scales_np.view(np.uint8).reshape(count, 12)
    rows[:, 24:27] = rgb_np
    rows[:, 27] = alpha_np
    rows[:, 28:32] = rotations_np

    path.parent.mkdir(exist_ok=True, parents=True)
    path.write_bytes(rows.tobytes())