"""NetLinLayer for weighting perceptual differences."""

import flax.nnx as nnx
import jax.numpy as jnp


class NetLinLayer(nnx.Module):
    """1x1 convolution layer for weighting perceptual feature differences.

    This layer applies a learned 1x1 convolution to reduce the channel
    dimension to 1, effectively learning importance weights for each
    feature channel.
    """

    def __init__(
        self,
        in_channels: int,
        *,
        use_dropout: bool = False,
        rngs: nnx.Rngs,
    ):
        """Initialize NetLinLayer.

        Args:
            in_channels: Number of input channels.
            use_dropout: Whether to apply dropout (default: False).
            rngs: Random number generators.
        """
        self.use_dropout = use_dropout
        if use_dropout:
            self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)

        # 1x1 conv with no bias
        self.conv = nnx.Conv(
            in_channels,
            1,
            kernel_size=(1, 1),
            use_bias=False,
            rngs=rngs,
        )

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """Apply the linear weighting layer.

        Args:
            x: Input tensor in NHWC format.

        Returns:
            Weighted output with 1 channel.
        """
        if self.use_dropout:
            x = self.dropout(x)
        return self.conv(x)
