from dataclasses import dataclass

import torch
from einops import rearrange
from jaxtyping import Float
from torch import Tensor, nn
import torch.nn.functional as F

from ....geometry.projection import get_world_rays
from ....misc.sh_rotation import rotate_sh
from .gaussians import build_covariance


@dataclass
class Gaussians:
    means: Float[Tensor, "*batch 3"]
    covariances: Float[Tensor, "*batch 3 3"]
    scales: Float[Tensor, "*batch 3"]
    rotations: Float[Tensor, "*batch 4"]
    harmonics: Float[Tensor, "*batch 3 _"]
    opacities: Float[Tensor, " *batch"]
    rotations_unnorm: Float[Tensor, "*batch 4"]


@dataclass
class GaussianAdapterCfg:
    gaussian_scale_min: float
    gaussian_scale_max: float
    sh_degree: int
    exp_scale_bias: float
    no_rotate_sh: bool
    no_sh_mask: bool


class GaussianAdapter(nn.Module):
    cfg: GaussianAdapterCfg

    def __init__(self, cfg: GaussianAdapterCfg):
        super().__init__()
        self.cfg = cfg

        # Create a mask for the spherical harmonics coefficients. This ensures that at
        # initialization, the coefficients are biased towards having a large DC
        # component and small view-dependent components.
        self.register_buffer(
            "sh_mask",
            torch.ones((self.d_sh,), dtype=torch.float32),
            persistent=False,
        )
        for degree in range(1, self.cfg.sh_degree + 1):
            self.sh_mask[degree**2 : (degree + 1) ** 2] = 0.1 * 0.25**degree

    def forward(
        self,
        extrinsics: Float[Tensor, "*#batch 4 4"],
        intrinsics: Float[Tensor, "*#batch 3 3"] | None,
        coordinates: Float[Tensor, "*#batch 2"],
        depths: Float[Tensor, "*#batch"] | None,
        opacities: Float[Tensor, "*#batch"],
        raw_gaussians: Float[Tensor, "*#batch _"],
        image_shape: tuple[int, int],
        eps: float = 1e-8,
        point_cloud: Float[Tensor, "*#batch 3"] | None = None,
        input_images: Tensor | None = None,
    ) -> Gaussians:
        # Gaussian geometry is consumed by CUDA/gsplat rasterizers, so keep the
        # position, covariance, and SH rotation math in FP32 under mixed precision.
        with torch.autocast(device_type=raw_gaussians.device.type, enabled=False):
            extrinsics = extrinsics.float()
            intrinsics = intrinsics.float() if intrinsics is not None else None
            coordinates = coordinates.float()
            depths = depths.float() if depths is not None else None
            opacities = opacities.float()
            raw_gaussians = raw_gaussians.float()

            scales, rotations_unnorm, sh = raw_gaussians.split(
                (3, 4, 3 * self.d_sh), dim=-1
            )

            scales = torch.clamp(
                F.softplus(scales - self.cfg.exp_scale_bias),
                min=self.cfg.gaussian_scale_min,
                max=self.cfg.gaussian_scale_max,
            )

            # Normalize the quaternion features to yield a valid quaternion.
            rotations = rotations_unnorm / (
                rotations_unnorm.norm(dim=-1, keepdim=True) + eps
            )

            # Split RGB spherical-harmonic coefficients into the expected layout.
            sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3)
            sh = sh.broadcast_to((*opacities.shape, 3, self.d_sh))
            if self.cfg.no_sh_mask:
                sh = sh.clone()
            else:
                sh = sh * self.sh_mask

            if input_images is not None:
                # [B, V, H*W, 1, 1, 3]
                imgs = rearrange(
                    input_images.float(), "b v c h w -> b v (h w) () () c"
                )
                # init sh with input images
                sh[..., 0] = sh[..., 0] + RGB2SH(imgs)

            # Create world-space covariance matrices.
            covariances = build_covariance(scales, rotations)
            c2w_rotations = extrinsics[..., :3, :3]
            covariances = (
                c2w_rotations @ covariances @ c2w_rotations.transpose(-1, -2)
            )

            # Compute Gaussian means.
            origins, directions = get_world_rays(coordinates, extrinsics, intrinsics)
            means = origins + directions * depths[..., None]
            harmonics = (
                sh
                if self.cfg.no_rotate_sh
                else rotate_sh(sh, c2w_rotations[..., None, :, :])
            )

        return Gaussians(
            means=means,
            covariances=covariances,
            harmonics=harmonics,
            opacities=opacities,
            scales=scales,
            rotations=rotations.broadcast_to((*scales.shape[:-1], 4)),
            rotations_unnorm=rotations_unnorm.broadcast_to((*scales.shape[:-1], 4)),
        )

    @property
    def d_sh(self) -> int:
        return (self.cfg.sh_degree + 1) ** 2

    @property
    def d_in(self) -> int:
        return 7 + 3 * self.d_sh


def RGB2SH(rgb):
    C0 = 0.28209479177387814
    return (rgb - 0.5) / C0