# Copyright 2026 Sperid Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.attention import AttentionModuleMixin
from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils import BaseOutput


@dataclass
class IdolGaussianDecoderOutput(BaseOutput):
    """
    Activated Gaussians, `num_points * gaussians_per_point` per batch element. The object is z-up, inside
    [-0.5, 0.5]^3.

    Args:
        xyz (`torch.Tensor` of shape `(batch_size, num_gaussians, 3)`): positions.
        features_dc (`torch.Tensor` of shape `(batch_size, num_gaussians, 1, 3)`): degree-0 spherical harmonics.
        scaling (`torch.Tensor` of shape `(batch_size, num_gaussians, 3)`): scales.
        rotation (`torch.Tensor` of shape `(batch_size, num_gaussians, 4)`): unit quaternions (w, x, y, z).
        opacity (`torch.Tensor` of shape `(batch_size, num_gaussians, 1)`): opacities.
    """

    xyz: torch.Tensor
    features_dc: torch.Tensor
    scaling: torch.Tensor
    rotation: torch.Tensor
    opacity: torch.Tensor


class IdolLayerNorm(nn.LayerNorm):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return super().forward(hidden_states.float()).type(hidden_states.dtype)


class IdolMultiHeadRMSNorm(nn.Module):
    def __init__(self, head_dim: int, num_heads: int):
        super().__init__()
        self.scale = head_dim**0.5
        self.weight = nn.Parameter(torch.ones(num_heads, head_dim))

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        normed = F.normalize(hidden_states.float(), dim=-1) * self.weight * self.scale
        return normed.to(hidden_states.dtype)


class IdolFeedForward(nn.Module):
    def __init__(self, dim: int, mlp_ratio: float = 4.0):
        super().__init__()
        inner_dim = int(dim * mlp_ratio)
        self.net = nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU(approximate="tanh"), nn.Linear(inner_dim, dim))

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.net(hidden_states)


class IdolPointEmbedding(nn.Module):
    """Sinusoidal embedding of 3D query points in [0, 1]."""

    def __init__(self, channels: int, in_channels: int = 3, max_res: int = 10):
        super().__init__()
        self.channels = channels
        freq_dim = channels // in_channels // 2
        self.freqs = torch.pow(2.0, torch.linspace(0.0, float(max_res), steps=freq_dim, dtype=torch.float32))

    def forward(self, points: torch.Tensor) -> torch.Tensor:
        num_points = points.shape[0]
        angles = points.unsqueeze(-1) * self.freqs.to(points.device, dtype=points.dtype) * torch.pi
        embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1).reshape(num_points, -1)
        if embedding.shape[1] < self.channels:
            padding = torch.zeros(
                num_points, self.channels - embedding.shape[1], device=embedding.device, dtype=embedding.dtype
            )
            embedding = torch.cat([embedding, padding], dim=-1)
        return embedding


class IdolAttnProcessor:
    _attention_backend = None
    _parallel_config = None

    def __call__(
        self, attn: "IdolAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None
    ) -> torch.Tensor:
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        query = attn.norm_q(attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1)))
        key = attn.norm_k(attn.to_k(encoder_hidden_states).unflatten(-1, (attn.heads, -1)))
        value = attn.to_v(encoder_hidden_states).unflatten(-1, (attn.heads, -1))

        hidden_states = dispatch_attention_fn(
            query, key, value, backend=self._attention_backend, parallel_config=self._parallel_config
        )
        hidden_states = hidden_states.flatten(2, 3).type_as(query)
        return attn.to_out[0](hidden_states)


class IdolAttention(nn.Module, AttentionModuleMixin):
    _default_processor_cls = IdolAttnProcessor
    _available_processors = [IdolAttnProcessor]

    def __init__(self, dim: int, heads: int, cross_attention_dim: int | None = None):
        super().__init__()
        self.heads = heads
        head_dim = dim // heads
        cross_attention_dim = cross_attention_dim or dim
        self.to_q = nn.Linear(dim, dim)
        self.to_k = nn.Linear(cross_attention_dim, dim)
        self.to_v = nn.Linear(cross_attention_dim, dim)
        self.norm_q = IdolMultiHeadRMSNorm(head_dim, heads)
        self.norm_k = IdolMultiHeadRMSNorm(head_dim, heads)
        self.to_out = nn.ModuleList([nn.Linear(dim, dim)])
        self.set_processor(IdolAttnProcessor())

    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> torch.Tensor:
        return self.processor(self, hidden_states, encoder_hidden_states)


class IdolGaussianDecoderBlock(nn.Module):
    """Self attention + cross attention to the object latent + FFN."""

    def __init__(self, dim: int, cross_attention_dim: int, num_heads: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = IdolLayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm2 = IdolLayerNorm(dim, elementwise_affine=True, eps=1e-6)
        self.norm3 = IdolLayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn1 = IdolAttention(dim, heads=num_heads)
        self.attn2 = IdolAttention(dim, heads=num_heads, cross_attention_dim=cross_attention_dim)
        self.ff = IdolFeedForward(dim, mlp_ratio=mlp_ratio)

    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = hidden_states + self.attn1(self.norm1(hidden_states))
        hidden_states = hidden_states + self.attn2(self.norm2(hidden_states), encoder_hidden_states)
        return hidden_states + self.ff(self.norm3(hidden_states))


def inverse_sigmoid(x: torch.Tensor) -> torch.Tensor:
    return torch.log(x / (1 - x))


def inverse_softplus(x: torch.Tensor) -> torch.Tensor:
    return x + torch.log(-torch.expm1(-x))


class IdolGaussianDecoderModel(ModelMixin, ConfigMixin):
    r"""
    Turns each anchor point into `gaussians_per_point` Gaussians, conditioned on the object latent, and returns them
    activated (positions, colours, scales, rotations, opacities).
    """

    _no_split_modules = ["IdolGaussianDecoderBlock"]
    _repeated_blocks = ["IdolGaussianDecoderBlock"]
    _keep_in_fp32_modules = [
        "in_proj",
        "input_layer",
        "out_proj",
        "norm2",
        "norm_q",
        "norm_k",
        "points_offset_perturbation",
        "base_offset_scale",
    ]

    # (name, shape per Gaussian) in the order the output channels are laid out
    _layout = [
        ("xyz", (3,)),
        ("features_dc", (1, 3)),
        ("scaling", (3,)),
        ("rotation", (4,)),
        ("opacity", (1,)),
        ("offset_scale", (1,)),
    ]

    @register_to_config
    def __init__(
        self,
        in_channels: int = 3,
        hidden_size: int = 1024,
        latent_channels: int = 16,
        num_layers: int = 16,
        num_attention_heads: int = 16,
        mlp_ratio: float = 4.0,
        gaussians_per_point: int = 32,
        xyz_lr: float = 1.0,
        features_dc_lr: float = 1.0,
        opacity_lr: float = 1.0,
        scaling_lr: float = 1.0,
        rotation_lr: float = 0.1,
        perturb_size: float = 1.5,
        filter_kernel_size: float = 0.0009,
        scaling_bias: float = 0.004,
        opacity_bias: float = 0.1,
    ):
        super().__init__()
        self.ranges = {}
        start = 0
        for name, shape in self._layout:
            size = gaussians_per_point * int(torch.tensor(shape).prod())
            self.ranges[name] = (start, start + size)
            start += size

        self.in_proj = nn.Linear(in_channels, hidden_size)
        self.pos_embedder = IdolPointEmbedding(hidden_size)
        self.input_layer = nn.Linear(hidden_size, hidden_size)
        self.blocks = nn.ModuleList(
            [
                IdolGaussianDecoderBlock(hidden_size, latent_channels, num_attention_heads, mlp_ratio)
                for _ in range(num_layers)
            ]
        )
        self.out_proj = nn.Linear(hidden_size, start)
        self.register_buffer("points_offset_perturbation", torch.zeros(gaussians_per_point, 3))
        self.register_buffer("base_offset_scale", torch.zeros(()))

    def _slice(self, features: torch.Tensor, name: str, shape: tuple) -> torch.Tensor:
        start, end = self.ranges[name]
        return features[:, :, start:end].reshape(features.shape[0], -1, self.config.gaussians_per_point, *shape)

    def forward(
        self, points: torch.Tensor, latents: torch.Tensor, return_dict: bool = True
    ) -> IdolGaussianDecoderOutput | tuple:
        """
        Args:
            points (`torch.Tensor` of shape `(batch_size, num_points, 3)`):
                Anchor points in [0, 1]^3, placed by the octree decoder.
            latents (`torch.Tensor` of shape `(batch_size, num_latent_tokens, latent_channels)`):
                Object latent.
            return_dict (`bool`, defaults to `True`):
                Whether to return an [`IdolGaussianDecoderOutput`] instead of a tuple.

        Returns:
            [`IdolGaussianDecoderOutput`] or `tuple`: the activated Gaussians.
        """
        batch_size, num_points, channels = points.shape
        hidden_states = self.in_proj(points) + self.pos_embedder(points.view(-1, channels)).view(
            batch_size, num_points, -1
        )
        hidden_states = self.input_layer(hidden_states)
        block_states = hidden_states.type(self.dtype)
        context = latents.type(self.dtype)
        for block in self.blocks:
            block_states = block(block_states, context)
        block_states = block_states.type(hidden_states.dtype).type(latents.dtype)
        features = self.out_proj(F.layer_norm(block_states, block_states.shape[-1:]))

        config = self.config
        offset_scale = F.softplus(self._slice(features, "offset_scale", (1,)) + self.base_offset_scale)
        offset = self._slice(features, "xyz", (3,)) * config.xyz_lr + self.points_offset_perturbation
        offset = torch.tanh(offset) * 0.5 * config.perturb_size * offset_scale
        raw_xyz = (offset + points[:, :, None, :]).flatten(1, 2)

        raw_features_dc = self._slice(features, "features_dc", (1, 3)).flatten(1, 2) * config.features_dc_lr
        raw_scaling = self._slice(features, "scaling", (3,)).flatten(1, 2) * config.scaling_lr
        raw_rotation = self._slice(features, "rotation", (4,)).flatten(1, 2) * config.rotation_lr
        raw_opacity = self._slice(features, "opacity", (1,)).flatten(1, 2) * config.opacity_lr

        device = raw_xyz.device
        scaling_offset = inverse_softplus(torch.tensor(config.scaling_bias)).to(device)
        opacity_offset = inverse_sigmoid(torch.tensor(config.opacity_bias)).to(device)
        rotation_offset = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device)

        # the object box is [-0.5, 0.5]^3
        xyz = raw_xyz * 1.0 - 0.5
        scaling = F.softplus(raw_scaling + scaling_offset)
        scaling = torch.sqrt(torch.square(scaling) + config.filter_kernel_size**2)
        rotation = F.normalize(raw_rotation + rotation_offset, dim=-1)
        opacity = torch.sigmoid(raw_opacity + opacity_offset)

        if not return_dict:
            return (xyz, raw_features_dc, scaling, rotation, opacity)
        return IdolGaussianDecoderOutput(
            xyz=xyz, features_dc=raw_features_dc, scaling=scaling, rotation=rotation, opacity=opacity
        )