from dataclasses import dataclass
from typing import Literal

import torch
from jaxtyping import Float
from torch import Tensor

from ...dataset import DatasetCfg
from ..types import Gaussians
from .decoder import Decoder, DecoderOutput
from .gsplat_splatting import (
    GsplatRasterizeMode,
    GsplatRenderMode,
    render_depth_gsplat,
    render_gsplat,
)


@dataclass
class DecoderSplattingGsplatCfg:
    name: Literal["splatting_gsplat", "gsplat"]
    render_mode: GsplatRenderMode = "RGB"
    rasterize_mode: GsplatRasterizeMode = "antialiased"
    use_covariances: bool = True
    packed: bool = True
    eps2d: float = 0.1
    sparse_grad: bool = False
    absgrad: bool = False
    channel_chunk: int = 32
    scale_invariant: bool = False


class DecoderSplattingGsplat(Decoder[DecoderSplattingGsplatCfg]):
    background_color: Float[Tensor, "3"]

    def __init__(
        self,
        cfg: DecoderSplattingGsplatCfg,
        dataset_cfg: DatasetCfg,
    ) -> None:
        super().__init__(cfg, dataset_cfg)
        self.register_buffer(
            "background_color",
            torch.tensor(dataset_cfg.background_color, dtype=torch.float32),
            persistent=False,
        )

    def forward(
        self,
        gaussians: Gaussians,
        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],
        depth_mode=None,
    ) -> DecoderOutput:
        render_mode = self.cfg.render_mode
        if depth_mode is not None:
            render_mode = "RGB+ED" if depth_mode == "relative_disparity" else "RGB+D"

        color, depth, accumulated_alpha = render_gsplat(
            extrinsics,
            intrinsics,
            near,
            far,
            image_shape,
            self.background_color[None].expand(extrinsics.shape[0], -1),
            gaussians.means,
            gaussians.covariances,
            gaussians.harmonics,
            gaussians.opacities,
            gaussian_scales=getattr(gaussians, "scales", None),
            gaussian_rotations=getattr(
                gaussians,
                "rotations_unnorm",
                getattr(gaussians, "rotations", None),
            ),
            render_mode=render_mode,
            use_covariances=self.cfg.use_covariances,
            eps2d=self.cfg.eps2d,
            packed=self.cfg.packed,
            sparse_grad=self.cfg.sparse_grad,
            absgrad=self.cfg.absgrad,
            rasterize_mode=self.cfg.rasterize_mode,
            channel_chunk=self.cfg.channel_chunk,
            scale_invariant=self.cfg.scale_invariant,
        )
        assert color is not None

        if depth_mode in ("disparity", "relative_disparity"):
            assert depth is not None
            depth = 1.0 / depth.clamp_min(1e-8)
        elif depth_mode == "log":
            assert depth is not None
            depth = depth.clamp_min(near[..., None, None]).clamp_max(
                far[..., None, None]
            ).log()
        elif depth_mode is None and "D" not in render_mode:
            depth = None

        return DecoderOutput(
            color,
            depth,
            accumulated_alpha=accumulated_alpha,
        )

    def render_depth(
        self,
        gaussians: Gaussians,
        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],
        mode="depth",
    ) -> Float[Tensor, "batch view height width"]:
        return render_depth_gsplat(
            extrinsics,
            intrinsics,
            near,
            far,
            image_shape,
            self.background_color[None].expand(extrinsics.shape[0], -1),
            gaussians.means,
            gaussians.covariances,
            gaussians.harmonics,
            gaussians.opacities,
            gaussian_scales=getattr(gaussians, "scales", None),
            gaussian_rotations=getattr(
                gaussians,
                "rotations_unnorm",
                getattr(gaussians, "rotations", None),
            ),
            mode=mode,
            use_covariances=self.cfg.use_covariances,
            eps2d=self.cfg.eps2d,
            packed=self.cfg.packed,
            sparse_grad=self.cfg.sparse_grad,
            absgrad=self.cfg.absgrad,
            rasterize_mode=self.cfg.rasterize_mode,
            channel_chunk=self.cfg.channel_chunk,
            scale_invariant=self.cfg.scale_invariant,
        )