"""Scaling layer for input normalization."""

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


class ScalingLayer(nnx.Module):
    """Normalize inputs from [-1, 1] to ImageNet statistics.

    This layer converts inputs from a [-1, 1] range to the normalized
    range expected by ImageNet-pretrained VGG networks.
    """

    def __init__(self):
        # ImageNet normalization values, adjusted for [-1, 1] input range
        # Original ImageNet: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        # For input in [-1, 1], we need to first convert to [0, 1] then normalize
        # shift = -(mean - 0.5) / std = (0.5 - mean) / std
        # scale = 0.5 / std
        # These values are from the official LPIPS implementation
        self.shift = jnp.array([-0.030, -0.088, -0.188]).reshape(1, 1, 1, 3)
        self.scale = jnp.array([0.458, 0.448, 0.450]).reshape(1, 1, 1, 3)

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """Normalize input tensor.

        Args:
            x: Input tensor in NHWC format with values in [-1, 1].

        Returns:
            Normalized tensor ready for VGG.
        """
        return (x - self.shift) / self.scale
