import torch
from torch import nn


class GaussianLatentDecoder(nn.Module):
    def __init__(
        self,
        latent_channels: int,
        hidden_channels: int,
        num_gaussian_parameters: int,
        init_sh_input_img: bool = True,
    ) -> None:
        super().__init__()
        in_channels = latent_channels + 1  # out_hat + depth
        self.decoder = nn.Sequential(
            nn.Conv2d(in_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"),
            nn.GELU(),
            nn.Conv2d(
                hidden_channels,
                num_gaussian_parameters,
                3,
                1,
                1,
                padding_mode="replicate",
            ),
            nn.GELU(),
            nn.Conv2d(
                num_gaussian_parameters,
                num_gaussian_parameters,
                3,
                1,
                1,
                padding_mode="replicate",
            ),
        )
        self._init_output_layer(init_sh_input_img)

    def _init_output_layer(self, init_sh_input_img: bool) -> None:
        # Keep the Gaussian head initially close to identity-like residuals.
        # Channel order: opacity, offset_xy, scales, quaternion, SH residuals.
        output_layer = self.decoder[-1]
        if init_sh_input_img:
            nn.init.zeros_(output_layer.weight[10:])
            nn.init.zeros_(output_layer.bias[10:])

        nn.init.zeros_(output_layer.weight[3:6])
        nn.init.zeros_(output_layer.bias[3:6])

    def forward(
        self,
        latent: torch.Tensor,
        depth: torch.Tensor,
    ) -> torch.Tensor:
        x = torch.cat((latent, depth), dim=1)
        return self.decoder(x)


class GaussianLatentDepthDecoder(nn.Module):
    def __init__(
        self,
        latent_channels: int,
        hidden_channels: int,
    ) -> None:
        super().__init__()
        self.inv_depth_logits = nn.Sequential(
            nn.Conv2d(latent_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"),
            nn.GELU(),
            nn.Conv2d(hidden_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"),
            nn.GELU(),
            nn.Conv2d(hidden_channels, 1, 3, 1, 1, padding_mode="replicate"),
        )
        # Predict the midpoint of the valid inverse-depth interval at initialization.
        nn.init.zeros_(self.inv_depth_logits[-1].weight)
        nn.init.zeros_(self.inv_depth_logits[-1].bias)

    def forward(
        self,
        latent: torch.Tensor,
        near: torch.Tensor,
        far: torch.Tensor,
        return_inverse: bool = False,
        return_normalized_inverse: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        with torch.autocast(device_type=latent.device.type, enabled=False):
            inv_depth_logits = self.inv_depth_logits(latent.float())
            near = near.float().clamp_min(torch.finfo(torch.float32).eps)
            far = far.float().clamp_min(torch.finfo(torch.float32).eps)
            inv_min = far.reciprocal()
            inv_max = near.reciprocal()
            inv_range = (inv_max - inv_min).clamp_min(torch.finfo(torch.float32).eps)

            normalized_inv_depth = torch.sigmoid(inv_depth_logits.float())
            inv_depth = inv_min + normalized_inv_depth * inv_range
            depth = inv_depth.clamp_min(torch.finfo(torch.float32).eps).reciprocal()

        outputs: list[torch.Tensor] = [depth]
        if return_inverse:
            outputs.append(inv_depth)
        if return_normalized_inverse:
            outputs.append(normalized_inv_depth)
        if len(outputs) == 1:
            return outputs[0]
        return tuple(outputs)