"""Stochastic Decomposition Layers to inject controllable noise into feature maps

In the original paper the Stochastic Decomposition Layer (SDL) is described as

The SDL decomposes the intermediate feature map into a deterministic component (the input)
and a stochastic component (the noise). This decomposition allows the model to separate
the signal processing from the uncertainty quantification.
The stochastic component is generated by modulating random Gaussian noise with a learned
style vector derived from a latent control variable. This architecture enables the
ensemble generation process to be explicitly controlled by the latent variable, rather
than relying on implicit randomness.

The SDL operation is defined as:
Output = x + (alpha * Style(z) * epsilon)

Where x is the deterministic input, z is the latent control vector, and alpha is a
learnable channel-wise scaling factor that determines the magnitude of the stochastic
perturbation.

"""

import torch
import torch.nn as nn


class StochasticDecompositionLayer(nn.Module):
    """Stochastic Decomposition Layer for controllable probabilistic outputs."""

    def __init__(self, input_dim: int, latent_dim: int):
        """Initialize the Stochastic Decomposition Layer.

        Args:
            input_dim: Number of channels in the input feature map
            latent_dim: Dimension of the latent control vector
        """
        super().__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim

        self.alpha = nn.Parameter(torch.zeros(1, input_dim, 1))

        self.style_net = nn.Linear(latent_dim, input_dim)

    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        """Apply stochastic decomposition to input features.

        Args:
            x: Input features [Batch, Channels, *Spatial]
            z: Latent control vector [Batch, Latent_Dim]

        Returns:
            Output features with injected stochasticity
        """
        if x.size(1) != self.input_dim:
            raise ValueError(f"Expected {self.input_dim} channels, got {x.size(1)}")
        epsilon = torch.randn_like(x)

        style = self.style_net(z)  # [B, C]

        spatial_dims = x.dim() - 2
        for _ in range(spatial_dims):
            style = style.unsqueeze(-1)

        alpha_broadcast = self.alpha
        while alpha_broadcast.dim() < x.dim():
            alpha_broadcast = alpha_broadcast.unsqueeze(-1)

        return x + (alpha_broadcast * style * epsilon)
