# VGG-16 perceptual loss (Phase 3.2/3.3): L1 feature distance at relu1_2, relu2_2, relu3_3.
# Expects images in [-1, 1]; converts internally to [0, 1] then ImageNet-normalises before VGG. Weights frozen.
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as tv_models


# L1 feature-matching loss at relu1_2 ([:4]), relu2_2 ([4:9]), relu3_3 ([9:16]) of VGG-16
class PerceptualLoss(nn.Module):
    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)
        )

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

    # L1 feature distance; real gradients are stopped — only fake trains
    def forward(self, fake: torch.Tensor, real: torch.Tensor) -> torch.Tensor:
        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
