import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

# 设置随机种子
torch.manual_seed(42)

# ── 生成器网络 ──────────────────────────────────────
class Generator(nn.Module):
    """
    生成器：从随机噪声生成数据
    输入：噪声向量 (batch_size, noise_dim)
    输出：生成数据 (batch_size, data_dim)
    """
    def __init__(self, noise_dim, data_dim, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(noise_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, data_dim),
            # 输出不激活，GAN 会学习合适的分布
        )

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


# ── 判别器网络 ──────────────────────────────────────
class Discriminator(nn.Module):
    """
    判别器：区分真实数据与生成数据
    输入：数据点 (batch_size, data_dim)
    输出：真实数据的概率 (batch_size, 1)
    """
    def __init__(self, data_dim, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(data_dim, hidden_dim),
            nn.LeakyReLU(0.2),  # LeakyReLU 防止梯度消失
            nn.Linear(hidden_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()  # 输出概率
        )

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


# 超参数
NOISE_DIM = 16
DATA_DIM = 2
HIDDEN_DIM = 64
BATCH_SIZE = 128

# 创建网络
generator = Generator(NOISE_DIM, DATA_DIM, HIDDEN_DIM)
discriminator = Discriminator(DATA_DIM, HIDDEN_DIM)

print(f"生成器参数量: {sum(p.numel() for p in generator.parameters()):,}")
print(f"判别器参数量: {sum(p.numel() for p in discriminator.parameters()):,}")

# ── 优化器 ──────────────────────────────────────
lr = 0.001
g_optimizer = optim.Adam(generator.parameters(), lr=lr)
d_optimizer = optim.Adam(discriminator.parameters(), lr=lr)

# 损失函数：二分类交叉熵
criterion = nn.BCELoss()

# ── 训练数据：环形分布 ──────────────────────────
def generate_real_data(batch_size):
    """生成环形分布的真实数据"""
    angles = torch.rand(batch_size) * 2 * torch.pi
    radius = 1.0 + torch.randn(batch_size) * 0.1  # 半径约为 1
    x = radius * torch.cos(angles)
    y = radius * torch.sin(angles)
    return torch.stack([x, y], dim=1)


# ── 训练循环 ──────────────────────────────────────
NUM_EPOCHS = 1000
d_losses = []
g_losses = []

for epoch in range(NUM_EPOCHS):
    # 1. 训练判别器
    # 生成假数据
    noise = torch.randn(BATCH_SIZE, NOISE_DIM)
    fake_data = generator(noise).detach()  # detach 避免计算生成器梯度

    # 生成真实数据
    real_data = generate_real_data(BATCH_SIZE)

    # 判别器损失
    real_pred = discriminator(real_data)
    fake_pred = discriminator(fake_data)
    d_loss = criterion(real_pred, torch.ones_like(real_pred)) + \
             criterion(fake_pred, torch.zeros_like(fake_pred))

    # 更新判别器
    d_optimizer.zero_grad()
    d_loss.backward()
    d_optimizer.step()

    # 2. 训练生成器
    # 生成新一批假数据
    noise = torch.randn(BATCH_SIZE, NOISE_DIM)
    fake_data = generator(noise)

    # 生成器损失：让判别器认为生成的数据是真实的
    fake_pred = discriminator(fake_data)
    g_loss = criterion(fake_pred, torch.ones_like(fake_pred))

    # 更新生成器
    g_optimizer.zero_grad()
    g_loss.backward()
    g_optimizer.step()

    # 记录损失
    d_losses.append(d_loss.item())
    g_losses.append(g_loss.item())

    if (epoch + 1) % 100 == 0:
        print(f"Epoch {epoch+1:4d} | D_loss: {d_loss:.4f} | G_loss: {g_loss:.4f}")

print("训练完成！")

# 生成数据并可视化
def visualize_results(generator, num_samples=1000):
    noise = torch.randn(num_samples, NOISE_DIM)
    generated_data = generator(noise).detach().numpy()

    plt.figure(figsize=(6, 6))
    plt.scatter(generated_data[:, 0], generated_data[:, 1],
                alpha=0.5, s=10, c='blue', label='Generated')
    plt.xlim(-2, 2)
    plt.ylim(-2, 2)
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('GAN Generated Data')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()


# 查看生成效果
visualize_results(generator)

import torch
import torch.nn as nn

# ── DCGAN 生成器 ─────────────────────────────────
class DCGenerator(nn.Module):
    """
    DCGAN 生成器：使用转置卷积上采样
    """
    def __init__(self, noise_dim=100, channels=3, features_g=64):
        super().__init__()
        self.noise_dim = noise_dim

        # 输入: noise_dim x 1 x 1
        self.net = nn.Sequential(
            # 转置卷积: (batch, features_g*8, 4, 4)
            nn.ConvTranspose2d(noise_dim, features_g * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(features_g * 8),
            nn.ReLU(True),

            # (batch, features_g*4, 8, 8)
            nn.ConvTranspose2d(features_g * 8, features_g * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_g * 4),
            nn.ReLU(True),

            # (batch, features_g*2, 16, 16)
            nn.ConvTranspose2d(features_g * 4, features_g * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_g * 2),
            nn.ReLU(True),

            # (batch, features_g, 32, 32)
            nn.ConvTranspose2d(features_g * 2, features_g, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_g),
            nn.ReLU(True),

            # 输出: (batch, channels, 64, 64)
            nn.ConvTranspose2d(features_g, channels, 4, 2, 1, bias=False),
            nn.Tanh()  # 输出范围 [-1, 1]
        )

    def forward(self, x):
        # x: (batch, noise_dim) -> (batch, noise_dim, 1, 1)
        x = x.view(x.size(0), x.size(1), 1, 1)
        return self.net(x)


# ── DCGAN 判别器 ─────────────────────────────────
class DCDiscriminator(nn.Module):
    """
    DCGAN 判别器：使用卷积下采样
    """
    def __init__(self, channels=3, features_d=64):
        super().__init__()

        # 输入: (batch, channels, 64, 64)
        self.net = nn.Sequential(
            # (batch, features_d, 32, 32)
            nn.Conv2d(channels, features_d, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            # (batch, features_d*2, 16, 16)
            nn.Conv2d(features_d, features_d * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_d * 2),
            nn.LeakyReLU(0.2, inplace=True),

            # (batch, features_d*4, 8, 8)
            nn.Conv2d(features_d * 2, features_d * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_d * 4),
            nn.LeakyReLU(0.2, inplace=True),

            # (batch, features_d*8, 4, 4)
            nn.Conv2d(features_d * 4, features_d * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_d * 8),
            nn.LeakyReLU(0.2, inplace=True),

            # 输出: (batch, 1, 1, 1)
            nn.Conv2d(features_d * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.net(x).view(x.size(0), -1)


# 测试网络
noise_dim = 100
generator = DCGenerator(noise_dim=noise_dim, channels=3, features_g=64)
discriminator = DCDiscriminator(channels=3, features_d=64)

# 测试前向传播
noise = torch.randn(2, noise_dim)
fake_images = generator(noise)
print(f"生成图像形状: {fake_images.shape}")  # torch.Size([2, 3, 64, 64])

decision = discriminator(fake_images)
print(f"判别结果形状: {decision.shape}")     # torch.Size([2, 1])

# ── 训练配置 ─────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

NOISE_DIM = 100
LEARNING_RATE = 0.0002
BETA1 = 0.5  # Adam 参数

# 创建网络
generator = DCGenerator(noise_dim=NOISE_DIM).to(device)
discriminator = DCDiscriminator().to(device)

# 优化器
g_optimizer = optim.Adam(generator.parameters(), lr=LEARNING_RATE, betas=(BETA1, 0.999))
d_optimizer = optim.Adam(discriminator.parameters(), lr=LEARNING_RATE, betas=(BETA1, 0.999))

criterion = nn.BCELoss()

# ── 训练循环 ─────────────────────────────────────
fixed_noise = torch.randn(64, NOISE_DIM, device=device)  # 用于可视化

def train_dcgan(generator, discriminator, g_optimizer, d_optimizer, criterion,
                num_epochs, device, fixed_noise):
    G_losses = []
    D_losses = []

    for epoch in range(num_epochs):
        for batch_idx in range(100):  # 假设每个 epoch 有 100 个 batch
            # 训练判别器
            discriminator.zero_grad()

            # 真实图像（假设已有）
            # real_images = ...
            # 这里用随机噪声模拟
            real_images = torch.randn(32, 3, 64, 64).to(device)

            batch_size = real_images.size(0)
            labels_real = torch.ones(batch_size, 1).to(device)
            labels_fake = torch.zeros(batch_size, 1).to(device)

            # 真实图像损失
            output = discriminator(real_images)
            d_loss_real = criterion(output, labels_real)

            # 生成图像损失
            noise = torch.randn(batch_size, NOISE_DIM).to(device)
            fake_images = generator(noise)
            output = discriminator(fake_images.detach())
            d_loss_fake = criterion(output, labels_fake)

            # 总损失
            d_loss = d_loss_real + d_loss_fake
            d_loss.backward()
            d_optimizer.step()

            # 训练生成器
            generator.zero_grad()

            noise = torch.randn(batch_size, NOISE_DIM).to(device)
            fake_images = generator(noise)
            output = discriminator(fake_images)
            g_loss = criterion(output, labels_real)  # 希望生成图像被判定为真

            g_loss.backward()
            g_optimizer.step()

            # 记录损失
            if batch_idx % 50 == 0:
                G_losses.append(g_loss.item())
                D_losses.append(d_loss.item())
                print(f"[{epoch}/{num_epochs}][{batch_idx}/100] "
                      f"D_loss: {d_loss:.4f} | G_loss: {g_loss:.4f}")

    return G_losses, D_losses


# 开始训练
# G_losses, D_losses = train_dcgan(generator, discriminator, g_optimizer,
#                                   d_optimizer, criterion, 5, device, fixed_noise)
print("DCGAN 架构已定义完成，可以开始训练！")

# WGAN 损失函数（替代 BCE）
def wgan_d_loss(real_pred, fake_pred):
    """判别器损失：真实样本得分高，生成样本得分低"""
    return -(torch.mean(real_pred) - torch.mean(fake_pred))

def wgan_g_loss(fake_pred):
    """生成器损失：让生成样本得分高"""
    return -torch.mean(fake_pred)

# 梯度惩罚（Gradient Penalty）- WGAN-GP
def gradient_penalty(discriminator, real_images, fake_images, device):
    """WGAN-GP 梯度惩罚项"""
    batch_size = real_images.size(0)
    alpha = torch.rand(batch_size, 1, 1, 1).to(device)

    # 在真实和生成图像之间插值
    interpolated = alpha * real_images + (1 - alpha) * fake_images
    interpolated.requires_grad = True

    # 计算插值图像的判别器输出
    pred = discriminator(interpolated)

    # 计算梯度
    gradients = torch.autograd.grad(
        outputs=pred,
        inputs=interpolated,
        grad_outputs=torch.ones_like(pred),
        create_graph=True,
        retain_graph=True,
        only_inputs=True
    )[0]

    # 计算梯度范数
    gradients = gradients.view(batch_size, -1)
    gradient_norm = gradients.norm(2, dim=1)
    penalty = ((gradient_norm - 1) ** 2).mean()

    return penalty

import torch.nn.utils.spectral_norm as spectral_norm

# 使用谱归一化的判别器
class SNDiscriminator(nn.Module):
    def __init__(self, channels=3, features_d=64):
        super().__init__()
        self.net = nn.Sequential(
            spectral_norm(nn.Conv2d(channels, features_d, 4, 2, 1)),
            nn.LeakyReLU(0.2, inplace=True),

            spectral_norm(nn.Conv2d(features_d, features_d * 2, 4, 2, 1)),
            nn.LeakyReLU(0.2, inplace=True),

            spectral_norm(nn.Conv2d(features_d * 2, features_d * 4, 4, 2, 1)),
            nn.LeakyReLU(0.2, inplace=True),

            spectral_norm(nn.Conv2d(features_d * 4, 1, 4, 1, 0)),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.net(x).view(x.size(0), -1)

import torch
import torch.nn as nn

class ConditionalGenerator(nn.Module):
    """条件生成器：同时接收噪声和类别标签"""
    def __init__(self, noise_dim, num_classes, embed_dim, img_channels, features_g=64):
        super().__init__()
        self.label_emb = nn.Embedding(num_classes, embed_dim)

        # 将噪声和标签嵌入拼接
        self.net = nn.Sequential(
            nn.Linear(noise_dim + embed_dim, features_g * 8 * 4 * 4),
            nn.BatchNorm1d(features_g * 8 * 4 * 4),
            nn.ReLU(),
            # 然后接转置卷积（类似 DCGAN）
            # ...
        )

    def forward(self, noise, labels):
        # 将类别标签嵌入到与噪声相同的维度
        label_embedding = self.label_emb(labels)
        # 拼接噪声和标签嵌入
        x = torch.cat([noise, label_embedding], dim=1)
        return self.net(x)


class ConditionalDiscriminator(nn.Module):
    """条件判别器：同时接收图像和类别标签"""
    def __init__(self, img_channels, num_classes, embed_dim, features_d=64):
        super().__init__()
        self.label_emb = nn.Embedding(num_classes, embed_dim)

        # 将图像和标签嵌入拼接
        self.net = nn.Sequential(
            nn.Conv2d(img_channels + embed_dim, features_d, 4, 2, 1),
            nn.LeakyReLU(0.2),
            # ...
        )

    def forward(self, img, labels):
        # 将标签嵌入调整到与图像相同的空间尺寸
        label_embedding = self.label_emb(labels)
        # 调整形状以便拼接
        label_embedding = label_embedding.unsqueeze(2).unsqueeze(3)
        label_embedding = label_embedding.expand(-1, -1, img.size(2), img.size(3))
        # 拼接图像和标签
        x = torch.cat([img, label_embedding], dim=1)
        return self.net(x)

import numpy as np
from scipy import linalg

def calculate_fid(real_activations, fake_activations):
    """
    计算 Fréchet Inception Distance
    real_activations: 真实图像的特征向量 (N, dim)
    fake_activations: 生成图像的特征向量 (N, dim)
    """
    # 计算均值和协方差
    mu1, sigma1 = real_activations.mean(axis=0), np.cov(real_activations, rowvar=False)
    mu2, sigma2 = fake_activations.mean(axis=0), np.cov(fake_activations, rowvar=False)

    # 计算 FID
    diff = mu1 - mu2
    # 计算协方差矩阵的和
    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)

    # 避免复数
    if np.iscomplexobj(covmean):
        covmean = covmean.real

    fid = diff.dot(diff) + np.trace(sigma1 + sigma2 - 2 * covmean)
    return fid


# 简化示例：使用随机数据
np.random.seed(42)
real_acts = np.random.randn(1000, 2048)  # Inception v3 输出维度
fake_acts = np.random.randn(1000, 2048)

fid_score = calculate_fid(real_acts, fake_acts)
print(f"FID Score: {fid_score:.2f}")
# FID 越低越好，通常小于 50 表示较好的生成质量