#SRGAN详解：超分辨率生成对抗网络原理与PyTorch实现

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

class ResidualBlock(nn.Module):
    """SRGAN残差块：跳跃连接防止梯度消失"""
    def __init__(self, channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.PReLU(),  # PReLU比ReLU更适合超分任务（减少死神经元）
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )

    def forward(self, x):
        return x + self.block(x)   # 局部跳跃连接

class UpsampleBlock(nn.Module):
    """PixelShuffle上采样：避免棋盘伪影"""
    def __init__(self, in_channels, up_scale=2):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, in_channels * up_scale**2, 3, padding=1),
            nn.PixelShuffle(up_scale),
            nn.PReLU(),
        )

    def forward(self, x):
        return self.block(x)

import torchvision.models as models

class VGG19Extractor(nn.Module):
    """冻结预训练VGG19，提取特征用于感知损失计算"""
    def __init__(self, layer_idx=35):  # 取VGG19的第35层（ReLU激活后）
        super().__init__()
        vgg = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1).features
        self.extractor = nn.Sequential(*list(vgg.children())[:layer_idx+1])
        
        for p in self.extractor.parameters():
            p.requires_grad = False
        
        # ImageNet的归一化参数
        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 forward(self, x):
        # 生成器输出在[-1,1]范围，先映射到[0,1]，再用ImageNet均值标准差归一化
        x = (x + 1) / 2
        x = (x - self.mean) / self.std
        return self.extractor(x)

def content_loss(sr, hr, vgg, mse_w=0.01):
    """内容损失 = 小权重像素MSE + 大权重感知损失"""
    pixel_loss = F.mse_loss(sr, hr)
    sr_feats = vgg(sr)
    hr_feats = vgg(hr)
    percept_loss = F.mse_loss(sr_feats, hr_feats)
    return mse_w * pixel_loss + percept_loss

def adv_loss_g(generator_out):
    """生成器损失：希望判别器认为生成的图是真的（标签为1）"""
    return F.mse_loss(generator_out, torch.ones_like(generator_out))

def adv_loss_d(real_out, fake_out):
    """判别器损失：真图标1，假图标0"""
    real_l = F.mse_loss(real_out, torch.ones_like(real_out))
    fake_l = F.mse_loss(fake_out, torch.zeros_like(fake_out))
    return (real_l + fake_l) / 2

def pretrain_gen(generator, dataloader, device, epochs=50, lr=1e-4):
    opt = torch.optim.Adam(generator.parameters(), lr=lr)
    mse = nn.MSELoss()
    generator.train()
    
    for epoch in range(epochs):
        for batch_idx, (hr, lr) in enumerate(dataloader):
            hr, lr = hr.to(device), lr.to(device)
            opt.zero_grad()
            sr = generator(lr)
            loss = mse(sr, hr)
            loss.backward()
            opt.step()
            if batch_idx % 100 == 0:
                print(f"Pretrain E{epoch} B{batch_idx} | MSE Loss: {loss:.4f}")
    torch.save(generator.state_dict(), "srresnet_pretrain.pth")

from PIL import Image
import torchvision.transforms as transforms

def enhance_old_photo(img_path, generator_path, device, scale=4):
    # 加载预训练生成器
    gen = Generator(scale_factor=scale).to(device)
    gen.load_state_dict(torch.load(generator_path, map_location=device))
    gen.eval()
    
    # 预处理：转tensor并归一化到[-1,1]
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
    ])
    
    img = Image.open(img_path).convert('RGB')
    lr = transform(img).unsqueeze(0).to(device)
    
    # 推理
    with torch.no_grad():
        sr = gen(lr)
    
    # 后处理：反归一化回[0,1]，再转PIL图像
    sr = (sr.squeeze(0).cpu() + 1) / 2
    sr = torch.clamp(sr, 0, 1)
    return transforms.ToPILImage()(sr)