import torch
import torch.nn as nn
import torch.nn.functional as F

class LPIPS(nn.Module):
    def __init__(self, net='alex'):
        super().__init__()
        import torchvision.models as models
        if net == 'alex':
            backbone = models.alexnet(pretrained=True).features
            channels = [64, 192, 384, 256, 256]
        self.slice1 = nn.Sequential(*list(backbone.children())[:2])
        self.slice2 = nn.Sequential(*list(backbone.children())[2:5])
        self.slice3 = nn.Sequential(*list(backbone.children())[5:8])
        self.slice4 = nn.Sequential(*list(backbone.children())[8:10])
        self.slice5 = nn.Sequential(*list(backbone.children())[10:12])
        self.weights = nn.ParameterList([
            nn.Parameter(torch.ones(c) / c) for c in channels])

    def forward(self, x, y):
        x_slices = [self.slice1(x), self.slice2(x),
                     self.slice3(x), self.slice4(x), self.slice5(x)]
        y_slices = [self.slice1(y), self.slice2(y),
                     self.slice3(y), self.slice4(y), self.slice5(y)]
        diffs = [(xs - ys) ** 2 for xs, ys in zip(x_slices, y_slices)]
        lpips = sum(
            (d.mean(dim=[2, 3]) * w).sum(dim=1)
            for d, w in zip(diffs, self.weights))
        return lpips.mean()

class QualityPredictor(nn.Module):
    def __init__(self, backbone_channels=2048):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Sequential(
            nn.Linear(backbone_channels, 512),
            nn.ReLU(inplace=True), nn.Dropout(0.5),
            nn.Linear(512, 128), nn.ReLU(inplace=True),
            nn.Linear(128, 1),
        )

    def forward(self, features):
        pooled = self.pool(features).flatten(1)
        return self.head(pooled)