"""VGG16 feature extractor for LPIPS."""

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

from lpips_nnx.weights import load_vgg16_weights


class VGG16Features(nnx.Module):
    """VGG16 feature extractor that outputs intermediate layer activations.

    Extracts features from 5 layers:
    - relu1_2: After block 1 (64 channels)
    - relu2_2: After block 2 (128 channels)
    - relu3_3: After block 3 (256 channels)
    - relu4_3: After block 4 (512 channels)
    - relu5_3: After block 5 (512 channels)
    """

    def __init__(self, *, rngs: nnx.Rngs):
        # Block 1: 2 conv layers, 64 channels
        self.conv_0 = nnx.Conv(3, 64, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_2 = nnx.Conv(64, 64, kernel_size=(3, 3), padding="SAME", rngs=rngs)

        # Block 2: 2 conv layers, 128 channels
        self.conv_5 = nnx.Conv(64, 128, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_7 = nnx.Conv(128, 128, kernel_size=(3, 3), padding="SAME", rngs=rngs)

        # Block 3: 3 conv layers, 256 channels
        self.conv_10 = nnx.Conv(128, 256, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_12 = nnx.Conv(256, 256, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_14 = nnx.Conv(256, 256, kernel_size=(3, 3), padding="SAME", rngs=rngs)

        # Block 4: 3 conv layers, 512 channels
        self.conv_17 = nnx.Conv(256, 512, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_19 = nnx.Conv(512, 512, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_21 = nnx.Conv(512, 512, kernel_size=(3, 3), padding="SAME", rngs=rngs)

        # Block 5: 3 conv layers, 512 channels
        self.conv_24 = nnx.Conv(512, 512, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_26 = nnx.Conv(512, 512, kernel_size=(3, 3), padding="SAME", rngs=rngs)
        self.conv_28 = nnx.Conv(512, 512, kernel_size=(3, 3), padding="SAME", rngs=rngs)

        # Load pretrained weights
        self._load_weights()

    def _load_weights(self):
        """Load pretrained VGG16 weights from torchvision."""
        weights = load_vgg16_weights()

        conv_layers = [
            self.conv_0,
            self.conv_2,
            self.conv_5,
            self.conv_7,
            self.conv_10,
            self.conv_12,
            self.conv_14,
            self.conv_17,
            self.conv_19,
            self.conv_21,
            self.conv_24,
            self.conv_26,
            self.conv_28,
        ]
        conv_indices = [0, 2, 5, 7, 10, 12, 14, 17, 19, 21, 24, 26, 28]

        for layer, idx in zip(conv_layers, conv_indices):
            name = f"conv_{idx}"
            layer.kernel[...] = weights[name]["kernel"]
            layer.bias[...] = weights[name]["bias"]

    def __call__(self, x: jnp.ndarray) -> list[jnp.ndarray]:
        """Extract VGG16 features at 5 intermediate layers.

        Args:
            x: Input tensor in NHWC format, normalized for ImageNet.

        Returns:
            List of 5 feature maps from relu1_2 through relu5_3.
        """
        features = []

        # Block 1
        x = jax.nn.relu(self.conv_0(x))
        x = jax.nn.relu(self.conv_2(x))
        features.append(x)  # relu1_2
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))

        # Block 2
        x = jax.nn.relu(self.conv_5(x))
        x = jax.nn.relu(self.conv_7(x))
        features.append(x)  # relu2_2
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))

        # Block 3
        x = jax.nn.relu(self.conv_10(x))
        x = jax.nn.relu(self.conv_12(x))
        x = jax.nn.relu(self.conv_14(x))
        features.append(x)  # relu3_3
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))

        # Block 4
        x = jax.nn.relu(self.conv_17(x))
        x = jax.nn.relu(self.conv_19(x))
        x = jax.nn.relu(self.conv_21(x))
        features.append(x)  # relu4_3
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))

        # Block 5
        x = jax.nn.relu(self.conv_24(x))
        x = jax.nn.relu(self.conv_26(x))
        x = jax.nn.relu(self.conv_28(x))
        features.append(x)  # relu5_3

        return features
