from math import isqrt
from typing import Literal

import torch
from einops import rearrange, repeat
from jaxtyping import Float
from torch import Tensor

from gsplat.rendering import rasterization


GsplatRenderMode = Literal["RGB", "RGB+D", "RGB+ED"]
GsplatRasterizeMode = Literal["classic", "antialiased"]


def normalized_intrinsics_to_pixel_K(
    intrinsics: Float[Tensor, "batch view 3 3"],
    image_shape: tuple[int, int],
) -> Float[Tensor, "batch view 3 3"]:
    h, w = image_shape
    K = intrinsics.clone().float()
    K[..., 0, 0] *= w
    K[..., 1, 1] *= h
    K[..., 0, 2] *= w
    K[..., 1, 2] *= h
    return K


def _identity_quats_like(
    means: Float[Tensor, "batch gaussian 3"],
) -> Float[Tensor, "batch gaussian 4"]:
    quats = torch.zeros((*means.shape[:-1], 4), device=means.device, dtype=means.dtype)
    # gsplat uses wxyz quaternions.
    quats[..., 0] = 1.0
    return quats


def _xyzw_to_wxyz(
    quats: Float[Tensor, "batch gaussian 4"],
) -> Float[Tensor, "batch gaussian 4"]:
    return torch.cat((quats[..., 3:4], quats[..., :3]), dim=-1)


@torch.autocast("cuda", enabled=False)
def render_gsplat(
    extrinsics: Float[Tensor, "batch view 4 4"],
    intrinsics: Float[Tensor, "batch view 3 3"],
    near: Float[Tensor, "batch view"],
    far: Float[Tensor, "batch view"],
    image_shape: tuple[int, int],
    background_color: Float[Tensor, "batch 3"],
    gaussian_means: Float[Tensor, "batch gaussian 3"],
    gaussian_covariances: Float[Tensor, "batch gaussian 3 3"] | None,
    gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"],
    gaussian_opacities: Float[Tensor, "batch gaussian"],
    gaussian_scales: Float[Tensor, "batch gaussian 3"] | None = None,
    gaussian_rotations: Float[Tensor, "batch gaussian 4"] | None = None,
    render_mode: GsplatRenderMode = "RGB",
    use_covariances: bool = True,
    eps2d: float = 0.1,
    packed: bool = True,
    sparse_grad: bool = False,
    absgrad: bool = False,
    rasterize_mode: GsplatRasterizeMode = "antialiased",
    channel_chunk: int = 32,
    scale_invariant: bool = False,
) -> tuple[
    Float[Tensor, "batch view 3 height width"] | None,
    Float[Tensor, "batch view height width"] | None,
    Float[Tensor, "batch view height width"],
]:
    b, v, _, _ = extrinsics.shape
    h, w = image_shape

    extrinsics = extrinsics.float()
    intrinsics = intrinsics.float()
    near = near.float()
    far = far.float()
    background_color = background_color.float()
    gaussian_means = gaussian_means.float()
    gaussian_sh_coefficients = gaussian_sh_coefficients.float()
    gaussian_opacities = gaussian_opacities.float()

    covars = gaussian_covariances.float() if gaussian_covariances is not None else None
    scales = gaussian_scales.float() if gaussian_scales is not None else None
    # The adapter stores quaternions as xyzw, while gsplat expects wxyz.
    quats = (
        _xyzw_to_wxyz(gaussian_rotations.float())
        if gaussian_rotations is not None
        else None
    )

    if use_covariances:
        if covars is None:
            raise ValueError("gsplat covariance mode requires gaussians.covariances.")
        if scales is None:
            scales = torch.ones_like(gaussian_means)
        if quats is None:
            quats = _identity_quats_like(gaussian_means)
    else:
        covars = None
        if scales is None or quats is None:
            raise ValueError(
                "gsplat scale/quaternion mode requires gaussians.scales and "
                "gaussians.rotations. Prefer use_covariances=true."
            )

    if scale_invariant:
        scale = 1.0 / near[:, :1].clamp_min(1e-8)
        extrinsics = extrinsics.clone()
        extrinsics[..., :3, 3] = extrinsics[..., :3, 3] * scale[..., None]
        gaussian_means = gaussian_means * scale[:, :, None]
        near = near * scale
        far = far * scale
        if covars is not None:
            covars = covars * (scale[:, :, None, None] ** 2)
        elif scales is not None:
            scales = scales * scale[:, :, None]

    Ks = normalized_intrinsics_to_pixel_K(intrinsics, image_shape)
    viewmats = torch.linalg.inv(extrinsics)
    sh_degree = isqrt(gaussian_sh_coefficients.shape[-1]) - 1
    shs = rearrange(
        gaussian_sh_coefficients, "b g c n -> b g n c"
    ).contiguous()

    render_rgb = "RGB" in render_mode
    render_depth = "D" in render_mode

    # gsplat's packed path has limited background broadcasting across render
    # modes. Our datasets use black backgrounds, so omitting the background keeps
    # the memory-efficient packed path available without changing semantics.
    backgrounds = (
        None
        if packed
        else repeat(background_color, "b c -> b v c", v=v).to(Ks)
    )

    render_out, alpha, _ = rasterization(
        means=gaussian_means,
        quats=quats,
        scales=scales,
        opacities=gaussian_opacities,
        colors=shs,
        viewmats=viewmats,
        Ks=Ks,
        width=w,
        height=h,
        near_plane=float(near.detach().min().item()),
        far_plane=float(far.detach().max().item()),
        eps2d=eps2d,
        sh_degree=sh_degree,
        packed=packed,
        backgrounds=backgrounds,
        render_mode=render_mode,
        sparse_grad=sparse_grad,
        absgrad=absgrad,
        rasterize_mode=rasterize_mode,
        channel_chunk=channel_chunk,
        covars=covars,
    )

    color = (
        rearrange(render_out[..., :3], "b v h w c -> b v c h w")
        if render_rgb
        else None
    )
    depth = render_out[..., -1] if render_depth else None
    return color, depth, alpha.squeeze(-1)


def render_depth_gsplat(
    extrinsics: Float[Tensor, "batch view 4 4"],
    intrinsics: Float[Tensor, "batch view 3 3"],
    near: Float[Tensor, "batch view"],
    far: Float[Tensor, "batch view"],
    image_shape: tuple[int, int],
    background_color: Float[Tensor, "batch 3"],
    gaussian_means: Float[Tensor, "batch gaussian 3"],
    gaussian_covariances: Float[Tensor, "batch gaussian 3 3"] | None,
    gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"],
    gaussian_opacities: Float[Tensor, "batch gaussian"],
    gaussian_scales: Float[Tensor, "batch gaussian 3"] | None = None,
    gaussian_rotations: Float[Tensor, "batch gaussian 4"] | None = None,
    mode: Literal["depth", "disparity", "relative_disparity", "log"] = "depth",
    use_covariances: bool = True,
    eps2d: float = 0.1,
    packed: bool = True,
    sparse_grad: bool = False,
    absgrad: bool = False,
    rasterize_mode: GsplatRasterizeMode = "antialiased",
    channel_chunk: int = 32,
    scale_invariant: bool = False,
) -> Float[Tensor, "batch view height width"]:
    render_mode = "RGB+ED" if mode == "relative_disparity" else "RGB+D"
    _, depth, _ = render_gsplat(
        extrinsics,
        intrinsics,
        near,
        far,
        image_shape,
        background_color,
        gaussian_means,
        gaussian_covariances,
        gaussian_sh_coefficients,
        gaussian_opacities,
        gaussian_scales=gaussian_scales,
        gaussian_rotations=gaussian_rotations,
        render_mode=render_mode,
        use_covariances=use_covariances,
        eps2d=eps2d,
        packed=packed,
        sparse_grad=sparse_grad,
        absgrad=absgrad,
        rasterize_mode=rasterize_mode,
        channel_chunk=channel_chunk,
        scale_invariant=scale_invariant,
    )
    assert depth is not None

    if mode in ("disparity", "relative_disparity"):
        depth = 1.0 / depth.clamp_min(1e-8)
    elif mode == "log":
        depth = depth.clamp_min(near[..., None, None]).clamp_max(
            far[..., None, None]
        ).log()

    return depth