#CycleGAN详解：循环一致性对抗网络原理与PyTorch实现

import torch
import torch.nn as nn

class ResidualBlock(nn.Module):
    """残差块：保持内容一致性的关键"""
    def __init__(self, channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(channels, channels, 3),
            nn.InstanceNorm2d(channels),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(channels, channels, 3),
            nn.InstanceNorm2d(channels)
        )

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

class Generator(nn.Module):
    """CycleGAN 生成器"""
    def __init__(self, input_nc=3, output_nc=3, n_residual_blocks=9):
        super().__init__()

        # 1. 初始卷积块 (c7s1-64)
        model = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(input_nc, 64, 7),
            nn.InstanceNorm2d(64),
            nn.ReLU(inplace=True)
        ]

        # 2. 下采样 (d128, d256)
        in_features = 64
        out_features = in_features * 2
        for _ in range(2):
            model += [
                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                nn.InstanceNorm2d(out_features),
                nn.ReLU(inplace=True)
            ]
            in_features = out_features
            out_features = in_features * 2

        # 3. 残差块 (R256 * 9)
        for _ in range(n_residual_blocks):
            model += [ResidualBlock(in_features)]

        # 4. 上采样 (u128, u64)
        out_features = in_features // 2
        for _ in range(2):
            model += [
                nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
                nn.InstanceNorm2d(out_features),
                nn.ReLU(inplace=True)
            ]
            in_features = out_features
            out_features = in_features // 2

        # 5. 输出层
        model += [
            nn.ReflectionPad2d(3),
            nn.Conv2d(64, output_nc, 7),
            nn.Tanh()
        ]

        self.model = nn.Sequential(*model)

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

class Discriminator(nn.Module):
    """PatchGAN 判别器，输出是一个 30×30 的特征图，每个值代表对应 Patch 的真假"""
    def __init__(self, input_nc=3):
        super().__init__()

        def discriminator_block(in_filters, out_filters, norm=True):
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
            if norm:
                layers.append(nn.InstanceNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *discriminator_block(input_nc, 64, norm=False),   # C64
            *discriminator_block(64, 128),                   # C128
            *discriminator_block(128, 256),                  # C256
            *discriminator_block(256, 512),                  # C512
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, 4, padding=1)                  # 最后一层不做 Norm，输出 1 通道
        )

    def forward(self, img):
        return self.model(img)

import torch.nn.functional as F

def compute_loss(real_A, real_B, G, F, D_A, D_B, lambda_cycle=10.0):
    # --- 1. 对抗损失 (Adversarial Loss) ---
    fake_B = G(real_A)
    pred_fake = D_B(fake_B)
    loss_GAN_G = F.mse_loss(pred_fake, torch.ones_like(pred_fake))  # 让 D 认为 fake_B 是真的

    # --- 2. 循环一致性损失 (Cycle Loss) ---
    # 前向循环：A -> B -> A
    fake_B = G(real_A)
    rec_A = F(fake_B)
    loss_cycle_A = F.l1_loss(rec_A, real_A)

    # 反向循环：B -> A -> B
    fake_A = F(real_B)
    rec_B = G(fake_A)
    loss_cycle_B = F.l1_loss(rec_B, real_B)

    # 总生成器损失：对抗损失 + λ * 循环一致性损失
    loss_G = loss_GAN_G + lambda_cycle * (loss_cycle_A + loss_cycle_B)
    return loss_G, fake_A, fake_B

# 初始化
G = Generator()   # A -> B
F = Generator()   # B -> A
D_A = Discriminator()
D_B = Discriminator()

# 优化器（生成器和判别器使用独立的优化器）
opt_G = torch.optim.Adam(list(G.parameters()) + list(F.parameters()), lr=0.0002, betas=(0.5, 0.999))
opt_D = torch.optim.Adam(list(D_A.parameters()) + list(D_B.parameters()), lr=0.0002, betas=(0.5, 0.999))

# 训练循环
for epoch in range(epochs):
    for real_A, real_B in zip(dataloader_A, dataloader_B):
        # ======= 训练生成器 G & F =======
        opt_G.zero_grad()
        # ... 调用上面的 compute_loss ...
        loss_G.backward()
        opt_G.step()

        # ======= 训练判别器 D_A & D_B =======
        opt_D.zero_grad()

        # 训练 D_A：判断真实图像
        pred_real = D_A(real_A)
        loss_D_real = F.mse_loss(pred_real, torch.ones_like(pred_real))

        # 训练 D_A：判断生成图像
        fake_A = F(real_B).detach()          # detach 防止梯度回传至生成器
        pred_fake = D_A(fake_A)
        loss_D_fake = F.mse_loss(pred_fake, torch.zeros_like(pred_fake))

        loss_D_A = (loss_D_real + loss_D_fake) * 0.5
        loss_D_A.backward()

        # 同理训练 D_B（此处省略...）

        opt_D.step()