"""Differentiable Gaussian rasterizer wrapper around gsplat.

We render, in a *single* pass (design §3.3), per view:
  * ``rgb`` or ``feature`` channels (alpha-composited colors),
  * ``depth`` = alpha-composited expected depth  D_hat = sum_i T_i a_i z_i / sum_i T_i a_i
    (gsplat ``RGB+ED`` mode), which is the GS-native depth used by L_mapdepth,
  * ``alpha`` (accumulated opacity), and
  * optional ``aux`` channels (e.g. a soft lane-membership map for L_lane).

gsplat consumes ``world2cam`` view matrices and scalar-first ``wxyz`` quats, so
the conventions in :mod:`mapgs.geometry.transforms` carry through unchanged.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import torch

from mapgs.geometry.transforms import se3_inverse
from mapgs.render.gaussians import Gaussians

try:
    from gsplat import rasterization as _gsplat_rasterization
    _HAS_GSPLAT = True
except Exception:  # pragma: no cover
    _HAS_GSPLAT = False


@dataclass
class RenderOutput:
    color: torch.Tensor              # [V, C, H, W]  (C=3 rgb or feature_dim)
    depth: torch.Tensor              # [V, H, W]     expected (alpha-normalized) depth
    alpha: torch.Tensor             # [V, H, W]     accumulated opacity
    aux: Optional[torch.Tensor] = None  # [V, A, H, W] auxiliary composited channels (e.g. lane)

    def rgb(self) -> torch.Tensor:
        """First 3 channels of ``color`` (valid when not in feature mode)."""
        return self.color[:, :3]


class GaussianRasterizer:
    """Thin, differentiable wrapper. One scene (one Gaussian set) -> many views."""

    def __init__(self, near: float = 0.01, far: float = 500.0, background: float = 0.0,
                 eps2d: float = 0.3):
        if not _HAS_GSPLAT:
            raise ImportError(
                "gsplat is required for rendering. Install with `pip install gsplat`."
            )
        self.near = near
        self.far = far
        self.background = background
        self.eps2d = eps2d

    def render(
        self,
        gaussians: Gaussians,
        K: torch.Tensor,            # [V, 3, 3]
        cam2world: torch.Tensor,    # [V, 4, 4]
        height: int,
        width: int,
        aux_colors: Optional[torch.Tensor] = None,  # [N, A] per-gaussian aux channels
    ) -> RenderOutput:
        V = K.shape[0]
        device = gaussians.means.device
        K = K.to(device)
        cam2world = cam2world.to(device)
        viewmats = se3_inverse(cam2world)  # world2cam

        if aux_colors is None:
            aux_colors = gaussians.aux  # ride-along aux channels (e.g. lane indicator)

        C = gaussians.colors.shape[-1]
        colors = gaussians.colors
        A = 0
        if aux_colors is not None:
            A = aux_colors.shape[-1]
            colors = torch.cat([colors, aux_colors], dim=-1)  # [N, C+A]

        bg = None
        if self.background != 0.0:
            bg = torch.full((V, colors.shape[-1]), float(self.background), device=device)

        # gsplat: render all composited channels + 1 expected-depth channel.
        out, alpha, _info = _gsplat_rasterization(
            means=gaussians.means,
            quats=gaussians.quats,
            scales=gaussians.scales,
            opacities=gaussians.opacities,
            colors=colors,
            viewmats=viewmats,
            Ks=K,
            width=width,
            height=height,
            near_plane=self.near,
            far_plane=self.far,
            eps2d=self.eps2d,
            render_mode="RGB+ED",
            backgrounds=bg,
            packed=False,
        )
        # out: [V, H, W, C+A+1]; alpha: [V, H, W, 1]
        depth = out[..., -1]                       # [V, H, W]
        composited = out[..., : C + A]             # [V, H, W, C+A]
        color = composited[..., :C].permute(0, 3, 1, 2).contiguous()
        aux = None
        if A > 0:
            aux = composited[..., C:].permute(0, 3, 1, 2).contiguous()
        return RenderOutput(
            color=color,
            depth=depth,
            alpha=alpha[..., 0],
            aux=aux,
        )

    def render_depth(self, gaussians: Gaussians, K, cam2world, height, width) -> torch.Tensor:
        """Depth-only convenience (still does a full pass; used by L_extrap z-buffer)."""
        return self.render(gaussians, K, cam2world, height, width).depth