"""VGG-16 perceptual loss for Phase 3.2 and 3.3.

Extracts features at relu1_2, relu2_2, relu3_3 and returns the
L1 distance in feature space. VGG weights are frozen.

Input convention: images in [-1, 1] — the loss converts internally to
[0, 1] and then applies ImageNet normalisation before passing to VGG.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as tv_models


class PerceptualLoss(nn.Module):
    """L1 feature-matching loss at three VGG-16 layers.

    VGG-16 feature indices:
      relu1_2: features[:4]   (before first maxpool)
      relu2_2: features[4:9]  (before second maxpool)
      relu3_3: features[9:16] (before third maxpool)
    """

    def __init__(self):
        super().__init__()
        vgg = tv_models.vgg16(weights=tv_models.VGG16_Weights.IMAGENET1K_V1)
        feats = vgg.features
        self.slice1 = nn.Sequential(*list(feats[:4]))   # relu1_2
        self.slice2 = nn.Sequential(*list(feats[4:9]))  # relu2_2
        self.slice3 = nn.Sequential(*list(feats[9:16])) # relu3_3

        for p in self.parameters():
            p.requires_grad_(False)

        self.register_buffer(
            "mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        )
        self.register_buffer(
            "std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
        )

    def _normalise(self, x: torch.Tensor) -> torch.Tensor:
        """Convert [-1, 1] → ImageNet-normalised [0, 1]."""
        x = x * 0.5 + 0.5            # → [0, 1]
        return (x - self.mean) / self.std

    def forward(self, fake: torch.Tensor, real: torch.Tensor) -> torch.Tensor:
        """L1 feature distance. real gradients are stopped — only fake trains."""
        f = self._normalise(fake)
        r = self._normalise(real)

        loss = torch.tensor(0.0, device=fake.device)
        for layer in (self.slice1, self.slice2, self.slice3):
            f = layer(f)
            r = layer(r)
            loss = loss + F.l1_loss(f, r.detach())
        return loss
