"""LPIPS (Learned Perceptual Image Patch Similarity) module."""

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

from lpips_nnx.netlin import NetLinLayer
from lpips_nnx.scaling import ScalingLayer
from lpips_nnx.vgg import VGG16Features
from lpips_nnx.weights import load_lpips_weights


def normalize_tensor(x: jnp.ndarray, eps: float = 1e-10) -> jnp.ndarray:
    """L2 normalize tensor along channel dimension (last axis for NHWC)."""
    norm = jnp.sqrt(jnp.sum(x**2, axis=-1, keepdims=True))
    return x / (norm + eps)


def spatial_average(x: jnp.ndarray, keepdims: bool = True) -> jnp.ndarray:
    """Compute spatial average over H and W dimensions."""
    return jnp.mean(x, axis=(1, 2), keepdims=keepdims)


class LPIPS(nnx.Module):
    """LPIPS (Learned Perceptual Image Patch Similarity) metric.

    Computes perceptual similarity between two images using VGG16 features
    with learned linear weights.

    Example:
        >>> lpips = LPIPS(rngs=nnx.Rngs(0))
        >>> # Images should be in NHWC format with values in [-1, 1]
        >>> x = jnp.zeros((1, 64, 64, 3))
        >>> y = jnp.ones((1, 64, 64, 3))
        >>> distance = lpips(x, y)
    """

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize LPIPS module.

        Args:
            rngs: Random number generators for initialization.
        """
        self.scaling = ScalingLayer()
        self.vgg = VGG16Features(rngs=rngs)

        # Linear layers for each VGG feature level
        # Channel counts: [64, 128, 256, 512, 512]
        channels = [64, 128, 256, 512, 512]
        self.lins = nnx.List([NetLinLayer(c, rngs=rngs) for c in channels])

        # Load pretrained LPIPS weights
        self._load_weights()

    def _load_weights(self):
        """Load pretrained LPIPS linear layer weights."""
        weights = load_lpips_weights()
        for i, lin in enumerate(self.lins):
            lin.conv.kernel[...] = weights[f"lin{i}"]

    def __call__(
        self,
        x: jnp.ndarray,
        y: jnp.ndarray,
        normalize: bool = False,
    ) -> jnp.ndarray:
        """Compute LPIPS distance between two images.

        Args:
            x: First image tensor in NHWC format.
            y: Second image tensor in NHWC format.
            normalize: If True, scale inputs from [0, 1] to [-1, 1].

        Returns:
            LPIPS distance tensor of shape (N, 1, 1, 1).
        """
        if normalize:
            x = 2 * x - 1
            y = 2 * y - 1

        # Scale inputs for VGG
        x_scaled = self.scaling(x)
        y_scaled = self.scaling(y)

        # Extract VGG features
        x_feats = self.vgg(x_scaled)
        y_feats = self.vgg(y_scaled)

        # Compute LPIPS distance
        diffs = []
        for i, (xf, yf) in enumerate(zip(x_feats, y_feats)):
            # L2 normalize features
            xf_norm = normalize_tensor(xf)
            yf_norm = normalize_tensor(yf)

            # Squared difference
            diff = (xf_norm - yf_norm) ** 2

            # Apply linear weighting
            weighted = self.lins[i](diff)

            # Spatial average
            avg = spatial_average(weighted, keepdims=True)
            diffs.append(avg)

        # Sum across all layers
        return sum(diffs)
