"""
GAN 系列算法
案例: example_train_simple_gan.py
代码实现经过验证, 大概率没问题
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Literal,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from zyplib.nn.modules import BaseModule
from zyplib.utils.print import warning

if TYPE_CHECKING:
    from zyplib.train.common import FabricTrainingContext


def _validate_generator_discriminator(
    generator: nn.Module,
    discriminator: nn.Module,
    latent_dim: int,
    real_data_shape: Tuple[int, ...],
    gan_type: str = 'standard',
) -> None:
    """验证生成器和判别器的输入输出维度是否匹配

    Parameters
    ----------
    - `generator` : `nn.Module`
        - 生成器网络
    - `discriminator` : `nn.Module`
        - 判别器网络
    - `latent_dim` : `int`
        - 潜在空间维度
    - `real_data_shape` : `Tuple[int, ...]`
        - 真实数据的形状 (不包括batch维度)
    - `gan_type` : `str`, optional
        - GAN类型，用于激活函数检查, by default "standard"

    Raises
    ----------
    - `ValueError`
        - 当模型维度不匹配时抛出
    """
    device = next(generator.parameters()).device
    batch_size = 2  # 使用小批量进行测试

    if gan_type == StandardGAN.gan_type:
        # 检查判别器是否有 Sigmoid 输出
        last_layer = list(discriminator.children())[-1]
        if not isinstance(last_layer, nn.Sigmoid):
            warning(
                '传入训练标准GAN的判别器模型时，最后一层未使用Sigmoid函数，建议检查确认该设置是否正确'
            )

    with torch.no_grad():
        # 测试生成器
        try:
            z = torch.randn(batch_size, latent_dim, device=device)
            fake_data = generator(z)
            expected_shape = (batch_size,) + real_data_shape
            if fake_data.shape != expected_shape:
                raise ValueError(
                    f'生成器输出形状 {fake_data.shape} 与期望形状 {expected_shape} 不匹配'
                )
        except Exception as e:
            raise ValueError(f'生成器无法处理 {latent_dim} 维潜在向量: {e}')

        # 测试判别器
        try:
            real_data = torch.randn(batch_size, *real_data_shape, device=device)
            real_output = discriminator(real_data)
            discriminator(fake_data)  # 测试判别器是否能处理生成的数据

            # 判别器输出应该是标量或单维度
            if real_output.dim() > 2 or (
                real_output.dim() == 2 and real_output.shape[1] != 1
            ):
                raise ValueError(
                    f'判别器输出形状 {real_output.shape} 应为 (batch_size,) 或 (batch_size, 1)'
                )
        except Exception as e:
            raise ValueError(f'判别器无法处理形状为 {real_data_shape} 的输入: {e}')


class SpectralNormMixin:
    """谱归一化混入类

    为 GAN 模型提供谱归一化功能，可与任意 GAN 变种组合使用。
    """

    @staticmethod
    def apply_spectral_norm(discriminator: nn.Module, name: str = 'weight'):
        """对判别器的所有卷积层和线性层应用谱归一化"""
        for module in discriminator.modules():
            if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Linear)):
                nn.utils.spectral_norm(module, name=name)


@dataclass
class GradientRegularizationConfig:
    """梯度正则化配置"""

    enable_r1: bool = False  # 通常的做法是启用 R1 正则
    enable_r2: bool = False
    lambda_r1: float = 10.0
    lambda_r2: float = 10.0
    r1_every_n_steps: int = 16  # R1正则化频率
    r2_every_n_steps: int = 16  # R2正则化频率


class GradientRegularization:
    """梯度正则化混入类"""

    def __init__(self, config: GradientRegularizationConfig | None = None):
        self.grad_reg_config = config or GradientRegularizationConfig()
        self._step_counter = 0

    @staticmethod
    def r1_regularization(real_data: Tensor, real_output: Tensor) -> Tensor:
        """R1 正则化：在真实数据上的梯度惩罚

        更简单、更稳定的替代 WGAN-GP 的方法
        """
        gradients = torch.autograd.grad(
            outputs=real_output.sum(),
            inputs=real_data,
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
        )[0]

        r1_penalty = gradients.pow(2).view(gradients.size(0), -1).sum(1).mean()
        return r1_penalty

    @staticmethod
    def r2_regularization(fake_data: Tensor, fake_output: Tensor) -> Tensor:
        """R2 正则化：在生成数据上的梯度惩罚"""
        gradients = torch.autograd.grad(
            outputs=fake_output.sum(),
            inputs=fake_data,
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
        )[0]

        r2_penalty = gradients.pow(2).view(gradients.size(0), -1).sum(1).mean()
        return r2_penalty

    def compute_gradient_regularization(
        self,
        real_data: Tensor,
        fake_data: Tensor,
        discriminator: nn.Module,
        **discriminator_kwargs,
    ) -> Tensor:
        """计算梯度正则化项"""
        total_reg = torch.tensor(0.0, device=real_data.device)

        # R1 正则化（在真实数据上）
        if (
            self.grad_reg_config.enable_r1
            and self._step_counter % self.grad_reg_config.r1_every_n_steps == 0
        ):
            with torch.enable_grad():
                # 需要重新计算以获得梯度
                real_data_reg = real_data.detach().clone().requires_grad_(True)
                real_output_reg = discriminator(real_data_reg, **discriminator_kwargs)

                r1_penalty = self.r1_regularization(real_data_reg, real_output_reg)
                total_reg = total_reg + self.grad_reg_config.lambda_r1 * r1_penalty

        # R2 正则化（在生成数据上）
        if (
            self.grad_reg_config.enable_r2
            and self._step_counter % self.grad_reg_config.r2_every_n_steps == 0
        ):
            with torch.enable_grad():
                # 需要重新计算以获得梯度
                fake_data_reg = fake_data.detach().clone().requires_grad_(True)
                fake_output_reg = discriminator(fake_data_reg, **discriminator_kwargs)

                r2_penalty = self.r2_regularization(fake_data_reg, fake_output_reg)
                total_reg = total_reg + self.grad_reg_config.lambda_r2 * r2_penalty

        return total_reg

    def step(self, reset_step_counter: bool = False):
        self._step_counter += 1
        if reset_step_counter:
            self._step_counter = 0

    def need_reg(self):
        enabled = self.grad_reg_config.enable_r1 or self.grad_reg_config.enable_r2
        match_step = (
            self._step_counter % self.grad_reg_config.r1_every_n_steps == 0
            or self._step_counter % self.grad_reg_config.r2_every_n_steps == 0
        )
        return enabled and match_step


class BasicGANFabricProcedure:
    def __init__(self, latent_dim: int | None = None):
        """默认内置的一个GAN训练步骤"""
        self.latent_dim = latent_dim

    def sample_latent(self, batch_size: int, device: torch.device):
        return torch.randn(batch_size, cast(Any, self.latent_dim), device=device)

    @torch.no_grad()
    def gan_val_step(self, ctx: 'FabricTrainingContext', batch, **kwargs):
        """自定义的GAN验证步骤，只计算损失"""
        gan_model = cast(GANBase, ctx.model)
        fabric = ctx.fabric

        real_data = batch[0]
        batch_size = real_data.size(0)

        # 计算损失，但不进行优化
        z = self.sample_latent(batch_size, cast(Any, fabric).device)
        if gan_model.gan_type == 'WGAN-GP':
            d_loss = gan_model.compute_discriminator_loss(
                z, real_data=real_data, use_gp=False
            )  # 在 Evaluate 模式下，手动设置为 False
        else:
            d_loss = gan_model.compute_discriminator_loss(z, real_data=real_data)

        if gan_model.gan_type == 'RaGAN':
            g_loss = gan_model.compute_generator_loss(z, real_data=real_data)
        else:
            g_loss = gan_model.compute_generator_loss(z)

        return None, {'g_loss': g_loss.item(), 'd_loss': d_loss.item()}

    def gan_train_step(self, ctx: 'FabricTrainingContext', batch, **kwargs):
        """自定义的GAN训练步骤，完全复用GANBase的逻辑"""
        gan_model = cast(GANBase, ctx.model)
        optimizer_g, optimizer_d = cast(Any, ctx.optimizer)
        fabric = ctx.fabric

        real_data = batch[0]
        batch_size = real_data.size(0)

        # --- 训练判别器 ---
        optimizer_d.zero_grad()
        z = self.sample_latent(batch_size, cast(Any, fabric).device)
        d_loss = gan_model.compute_discriminator_loss(z, real_data)
        cast(Any, fabric).backward(d_loss)
        optimizer_d.step()

        if gan_model.gan_type == 'WGAN' and (gan_model.clip_value is not None):
            cast(Any, gan_model).clip_discriminator_weights()

        # --- 训练生成器 ---
        optimizer_g.zero_grad()
        z = self.sample_latent(batch_size, cast(Any, fabric).device)
        if gan_model.gan_type == 'RaGAN':
            g_loss = gan_model.compute_generator_loss(z, real_data=real_data)
        else:
            g_loss = gan_model.compute_generator_loss(z)
        cast(Any, fabric).backward(g_loss)
        optimizer_g.step()

        return None, {'d_loss': d_loss.item(), 'g_loss': g_loss.item()}


class GANBase(BaseModule, ABC):
    """生成对抗网络基类

    提供 GAN 的基本接口和通用功能。

    Parameters
    ----------
    - `generator` : `nn.Module`
        - 生成器网络
        - **输入维度**: 应接受形状为 `(batch_size, latent_dim)` 的潜在向量
        - **输出维度**: 应输出形状为 `(batch_size, *real_data_shape)` 的生成数据
    - `discriminator` : `nn.Module`
        - 判别器网络
        - **输入维度**: 应接受形状为 `(batch_size, *real_data_shape)` 的数据
        - **输出维度**: 应输出形状为 `(batch_size,)` 或 `(batch_size, 1)` 的判别结果
    """

    def __init__(self, generator: nn.Module, discriminator: nn.Module):
        super().__init__()
        self.generator = generator
        self.discriminator = discriminator

    gan_type = 'GANBase'

    @abstractmethod
    def generator_loss(
        self, fake_output: Tensor, real_output: Optional[Tensor] = None
    ) -> Tensor:
        """计算生成器损失

        Parameters
        ----------
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出
        - `real_output` : `Tensor`, optional
            - 判别器对真实数据的输出; RaGAN 需要用到

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        """
        pass

    @abstractmethod
    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失

        Parameters
        ----------
        - `real_output` : `Tensor`
            - 判别器对真实数据的输出
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 判别器损失
        """
        pass

    def generate(self, z: Tensor) -> Tensor:
        """生成数据

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在向量

        Returns
        ----------
        - `Tensor`
            - 生成的数据
        """
        return self.generator(z)

    def discriminate(self, x: Tensor) -> Tensor:
        """判别数据

        Parameters
        ----------
        - `x` : `Tensor`
            - 输入数据

        Returns
        ----------
        - `Tensor`
            - 判别结果
        """
        return self.discriminator(x)

    def compute_generator_loss(self, z: Tensor, **kwargs) -> Tensor:
        """计算生成器损失

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在向量
        - `**kwargs`
            - 其他参数

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        - `Tensor`
            - 生成数据
        """
        fake_data = self.generate(z)
        fake_output = self.discriminate(fake_data)
        g_loss = self.generator_loss(fake_output)
        return g_loss

    def compute_discriminator_loss(
        self, z: Tensor, real_data: Tensor, **kwargs
    ) -> Tensor:
        """计算判别器损失

        Parameters
        ----------
        - `real_data` : `Tensor`
            - 真实数据
        - `z` : `Tensor`
            - 潜在向量
        """
        fake_data = self.generate(z).detach()  # 阻止梯度回传到生成器
        real_output = self.discriminate(real_data)
        fake_output = self.discriminate(fake_data)
        d_loss = self.discriminator_loss(real_output, fake_output)

        return d_loss

    @torch.no_grad()
    def sample(
        self, num_samples: int, latent_dim: int, device: Optional[torch.device] = None
    ) -> Tensor:
        """从先验分布采样生成新数据

        Parameters
        ----------
        - `num_samples` : `int`
            - 采样数量
        - `latent_dim` : `int`
            - 潜在空间维度
        - `device` : `torch.device`, optional
            - 设备, by default None

        Returns
        ----------
        - `Tensor`
            - 生成的数据
        """
        if device is None:
            device = next(self.generator.parameters()).device

        z = torch.randn(num_samples, latent_dim, device=device)
        return self.generate(z)


def _compute_log(tensor: Tensor, eps: float = 1e-8) -> Tensor:
    """计算对数，防止数值不稳定"""
    return torch.log(torch.clamp(tensor, min=eps))


class StandardGAN(GANBase):
    gan_type = 'StandardGAN'

    def __init__(
        self,
        generator: nn.Module,
        discriminator: nn.Module,
        real_data_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        generator_loss_type: Literal['saturating', 'non-saturating'] = 'non-saturating',
        validate_model: bool = True,
        apply_spectral_norm: bool = False,
    ):
        """标准生成对抗网络

        实现原始 GAN 的 minimax 损失函数。

        References
        ----------
        Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., ... & Bengio, Y. (2014).
        Generative adversarial nets.
        Advances in neural information processing systems, 27, 2672-2680.

        Parameters
        ----------
        - `generator` : `nn.Module`
            - 生成器网络
            - **输入维度**: 应接受形状为 `(batch_size, latent_dim)` 的潜在向量
            - **输出维度**: 应输出形状为 `(batch_size, *real_data_shape)` 的生成数据
        - `discriminator` : `nn.Module`
            - 判别器网络
            - **输入维度**: 应接受形状为 `(batch_size, *real_data_shape)` 的数据
            - **输出维度**: 应输出形状为 `(batch_size,)` 或 `(batch_size, 1)` 的判别结果
        - `real_data_shape` : `Tuple[int, ...]`, optional
            - 真实数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `latent_dim` : `int`, optional
            - 潜在空间维度，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True
        - `generator_loss_type` : `str`, optional
            - 生成器损失类型, 可选 'saturating' 或 'non-saturating', by default 'non-saturating'

        Examples
        ----------
        ```python
        # 创建生成器和判别器 (MNIST 28x28)
        generator = nn.Sequential(
            # ...
            nn.Unflatten(1, (28, 28))
        )
        discriminator = nn.Sequential(
            # ...
            nn.Sigmoid()
        )
        ```

        Raises
        ----------
        - `ValueError`
            - 当模型维度不匹配且 validate_model=True 时抛出
            - 当 `generator_loss_type` 不支持时抛出
        """
        super().__init__(generator, discriminator)

        if generator_loss_type not in ['saturating', 'non-saturating']:
            raise ValueError(f'Unsupported generator_loss_type: {generator_loss_type}')
        self.generator_loss_type = generator_loss_type

        # 模型验证
        if validate_model and real_data_shape is not None and latent_dim is not None:
            _validate_generator_discriminator(
                generator,
                discriminator,
                latent_dim,
                real_data_shape,
                gan_type=self.gan_type,
            )

        # 应用谱归一化
        if apply_spectral_norm:
            SpectralNormMixin.apply_spectral_norm(self.discriminator)

    def generator_loss(
        self, fake_output: Tensor, real_output: Optional[Tensor] = None
    ) -> Tensor:
        """计算生成器损失

        支持两种损失类型:
        - `saturating`: `min log(1 - D(G(z)))` (原始GAN论文中的损失函数)
            - 在 D(G(z)) 趋近于 0 时梯度较弱, 容易饱和。
        - `non-saturating`: `min -log(D(G(z)))` <==> `max log(D(G(z)))` (启发式修改)
            - 在 D(G(z)) 趋近于 0 时能提供很强的梯度信号, 有助于训练早期的稳定性。

        Parameters
        ----------
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出 (经过 sigmoid)

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        """
        if self.generator_loss_type == 'non-saturating':
            # 核心逻辑修改: 实现非饱和损失函数
            # return -torch.mean(torch.log(fake_output + 1e-8))
            return -torch.mean(_compute_log(fake_output, eps=1e-8))
        elif self.generator_loss_type == 'saturating':
            # 核心逻辑修改: 实现原始的饱和损失函数
            return torch.mean(_compute_log(1 - fake_output, eps=1e-8))
        raise AssertionError(
            f'Unsupported generator_loss_type: {self.generator_loss_type}'
        )

    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失 (最大化 log(D(x)) + log(1 - D(G(z))))

        Parameters
        ----------
        - `real_output` : `Tensor`
            - 判别器对真实数据的输出
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 判别器损失
        """
        real_loss = -torch.mean(_compute_log(real_output, eps=1e-8))
        fake_loss = -torch.mean(_compute_log(1 - fake_output, eps=1e-8))
        return real_loss + fake_loss


class WGAN(GANBase):
    gan_type = 'WGAN'

    def __init__(
        self,
        generator: nn.Module,
        discriminator: nn.Module,
        clip_value: float = 0.01,
        real_data_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
        apply_spectral_norm: bool = False,
    ):
        """Wasserstein GAN

        实现基于 Wasserstein 距离的 GAN，使用权重裁剪确保 Lipschitz 约束。

        References
        ----------
        Arjovsky, M., Chintala, S., & Bottou, L. (2017).
        Wasserstein generative adversarial networks.
        In International conference on machine learning (pp. 214-223). PMLR.

        Parameters
        ----------
        - `generator` : `nn.Module`
            - 生成器网络
        - `discriminator` : `nn.Module`
            - 判别器网络 (在 WGAN 中称为 critic)
            - **判别器最后一层不使用 sigmoid**
        - `clip_value` : `float`, optional
            - 权重裁剪的范围 [-clip_value, clip_value]
        - `real_data_shape` : `Tuple[int, ...]`, optional
            - 真实数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `latent_dim` : `int`, optional
            - 潜在空间维度，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True

        Examples
        ----------
        ```python
        # 创建 WGAN (注意判别器最后一层不使用 sigmoid)
        generator = nn.Sequential(
            ...
        )
        discriminator = nn.Sequential(
            ...
            nn.Linear(512, 1)  # 无 sigmoid
        )

        wgan = WGAN(generator, discriminator, clip_value=0.01)

        # 训练判别器
        d_loss = wgan.compute_discriminator_loss(z, real_data)
        d_loss.backward()
        optimizer_d.step()
        wgan.clip_discriminator_weights() # 在训练判别器后裁剪权重
        # 训练生成器
        ```
        """
        super().__init__(generator, discriminator)
        self.clip_value = clip_value

        # 模型验证
        if validate_model and real_data_shape is not None and latent_dim is not None:
            _validate_generator_discriminator(
                generator,
                discriminator,
                latent_dim,
                real_data_shape,
                gan_type=self.gan_type,
            )

        # 应用谱归一化
        if apply_spectral_norm:
            SpectralNormMixin.apply_spectral_norm(self.discriminator)

    def generator_loss(self, fake_output: Tensor) -> Tensor:
        """计算生成器损失 (最大化 D(G(z)))

        Parameters
        ----------
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        """
        return -torch.mean(fake_output)

    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失 (最大化 D(x) - D(G(z)))

        Parameters
        ----------
        - `real_output` : `Tensor`
            - 判别器对真实数据的输出
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 判别器损失
        """
        return torch.mean(fake_output) - torch.mean(real_output)

    def clip_discriminator_weights(self):
        """裁剪判别器权重以满足 Lipschitz 约束"""
        for param in self.discriminator.parameters():
            param.data.clamp_(-self.clip_value, self.clip_value)


class WGANGP(GANBase):
    gan_type = 'WGAN-GP'

    def __init__(
        self,
        generator: nn.Module,
        discriminator: nn.Module,
        lambda_gp: float = 10.0,
        gp_target: float = 1.0,
        real_data_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
        apply_spectral_norm: bool = False,
        **kwargs,
    ):
        """WGAN with Gradient Penalty

        使用梯度惩罚替代权重裁剪的 WGAN 改进版本。

        References
        ----------
        Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., & Courville, A. C. (2017).
        Improved training of wasserstein gans.
        Advances in neural information processing systems, 30, 5767-5777.

        Parameters
        ----------
        - `generator` : `nn.Module`
            - 生成器网络
        - `discriminator` : `nn.Module`
            - 判别器网络; **判别器最后一层不使用 sigmoid**
        - `lambda_gp` : `float`, optional
            - 梯度惩罚的权重系数, by default 10.0
        - `gp_target` : `float`, optional
            - 梯度惩罚的目标, 默认 1
        - `real_data_shape` : `Tuple[int, ...]`, optional
            - 真实数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `latent_dim` : `int`, optional
            - 潜在空间维度，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True
        """
        GANBase.__init__(self, generator, discriminator)
        self.lambda_gp = lambda_gp
        self.gp_target = gp_target

        # 模型验证
        if validate_model and real_data_shape is not None and latent_dim is not None:
            _validate_generator_discriminator(
                generator,
                discriminator,
                latent_dim,
                real_data_shape,
                gan_type=self.gan_type,
            )

        # 应用谱归一化
        if apply_spectral_norm:
            SpectralNormMixin.apply_spectral_norm(self.discriminator)

    def gradient_penalty(self, real_data: Tensor, fake_data: Tensor) -> Tensor:
        """计算梯度惩罚

        Parameters
        ----------
        - `real_data` : `Tensor`
            - 真实数据
        - `fake_data` : `Tensor`
            - 生成数据

        Returns
        ----------
        - `Tensor`
            - 梯度惩罚项
        """
        batch_size = real_data.size(0)
        device = real_data.device

        # 随机插值
        alpha_shape = [batch_size] + [1] * (real_data.dim() - 1)
        alpha = torch.rand(alpha_shape, device=device, dtype=real_data.dtype)

        # 扩展到与数据相同的维度
        alpha = alpha.expand_as(real_data)

        interpolated = alpha * real_data + (1 - alpha) * fake_data
        interpolated.requires_grad_(True)

        # 计算判别器对插值数据的输出
        d_interpolated = self.discriminator(interpolated)

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

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

        # 梯度惩罚: (||∇D(x̂)||₂ - 1)²
        penalty = torch.mean((gradient_norm - self.gp_target) ** 2)

        return penalty

    def generator_loss(self, fake_output: Tensor) -> Tensor:
        """计算生成器损失 (最大化 D(G(z)))

        Parameters
        ----------
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        """
        return -torch.mean(fake_output)

    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失 (最大化 D(x) - D(G(z)))

        Parameters
        ----------
        - `real_output` : `Tensor`
            - 判别器对真实数据的输出
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 判别器损失
        """
        return torch.mean(fake_output) - torch.mean(real_output)

    def compute_discriminator_loss(
        self, z: Tensor, real_data: Tensor, use_gp: bool = True, **kwargs
    ) -> Tensor:
        """计算包含梯度惩罚的判别器损失

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在向量
        - `real_data` : `Tensor`
            - 真实数据
        - `use_gp` : `bool`, optional
            - 是否使用梯度惩罚, by default True
            - 在 Evaluate 模式下，应该手动设置为 False
        """
        fake_data = self.generate(z).detach()
        real_output = self.discriminate(real_data)
        fake_output = self.discriminate(fake_data)

        # WGAN 损失
        wgan_loss = self.discriminator_loss(real_output, fake_output)

        # 梯度惩罚
        if use_gp:
            gp = self.gradient_penalty(real_data, fake_data)
            total_loss = wgan_loss + self.lambda_gp * gp
        else:
            gp = torch.tensor(0.0, device=real_data.device)
            total_loss = wgan_loss

        return total_loss


class LSGAN(GANBase):
    gan_type = 'LSGAN'

    def __init__(
        self,
        generator: nn.Module,
        discriminator: nn.Module,
        real_data_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
        apply_spectral_norm: bool = False,
    ):
        """Least Squares GAN

        使用最小二乘损失替代交叉熵损失的 GAN 变种。

        References
        ----------
        Mao, X., Li, Q., Xie, H., Lau, R. Y., Wang, Z., & Paul Smolley, S. (2017).
        Least squares generative adversarial networks.
        In Proceedings of the IEEE international conference on computer vision (pp. 2794-2802).

        Parameters
        ----------
        - `generator` : `nn.Module`
            - 生成器网络
        - `discriminator` : `nn.Module`
            - 判别器网络 (输出不使用 sigmoid)
        - `real_data_shape` : `Tuple[int, ...]`, optional
            - 真实数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `latent_dim` : `int`, optional
            - 潜在空间维度，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True
        - `apply_spectral_norm` : `bool`, optional
            - 是否对判别器应用谱归一化, by default False
        """
        super().__init__(generator, discriminator)

        # 模型验证
        if validate_model and real_data_shape is not None and latent_dim is not None:
            _validate_generator_discriminator(
                generator,
                discriminator,
                latent_dim,
                real_data_shape,
                gan_type=self.gan_type,
            )

        # 应用谱归一化
        if apply_spectral_norm:
            SpectralNormMixin.apply_spectral_norm(self.discriminator)

    def generator_loss(self, fake_output: Tensor) -> Tensor:
        """计算生成器损失 (最小化 (D(G(z)) - 1)²)

        Parameters
        ----------
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        """
        return 0.5 * torch.mean((fake_output - 1) ** 2)

    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失 (最小化 (D(x) - 1)² + D(G(z))²)

        Parameters
        ----------
        - `real_output` : `Tensor`
            - 判别器对真实数据的输出
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 判别器损失
        """
        real_loss = 0.5 * torch.mean((real_output - 1) ** 2)
        fake_loss = 0.5 * torch.mean(fake_output**2)
        return real_loss + fake_loss


class HingeGAN(GANBase):
    gan_type = 'HingeGAN'

    def __init__(
        self,
        generator: nn.Module,
        discriminator: nn.Module,
        real_data_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
        apply_spectral_norm: bool = True,
    ):
        """Hinge Loss GAN

        使用 Hinge 损失的 GAN 变种，常用于 Self-Attention GAN 等模型。

        References
        ----------
        Zhang, H., Goodfellow, I., Metaxas, D., & Odena, A. (2019).
        Self-attention generative adversarial networks.
        In International conference on machine learning (pp. 7354-7363). PMLR.

        Parameters
        ----------
        - `generator` : `nn.Module`
            - 生成器网络
        - `discriminator` : `nn.Module`
            - 判别器网络 (输出不使用 sigmoid)
        - `real_data_shape` : `Tuple[int, ...]`, optional
            - 真实数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `latent_dim` : `int`, optional
            - 潜在空间维度，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True
        - `apply_spectral_norm` : `bool`, optional
            - 是否对判别器应用谱归一化, by default True

        Examples
        ----------
        ```python
        hinge_gan = HingeGAN(generator, discriminator)
        ```
        """
        super().__init__(generator, discriminator)

        # 模型验证
        if validate_model and real_data_shape is not None and latent_dim is not None:
            _validate_generator_discriminator(
                generator,
                discriminator,
                latent_dim,
                real_data_shape,
                gan_type=self.gan_type,
            )

        # 应用谱归一化
        if apply_spectral_norm:
            SpectralNormMixin.apply_spectral_norm(self.discriminator)

    def generator_loss(self, fake_output: Tensor) -> Tensor:
        """计算生成器损失 (-D(G(z)))

        Parameters
        ----------
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 生成器损失
        """
        return -torch.mean(fake_output)

    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失 (ReLU(1 - D(x)) + ReLU(1 + D(G(z))))

        Parameters
        ----------
        - `real_output` : `Tensor`
            - 判别器对真实数据的输出
        - `fake_output` : `Tensor`
            - 判别器对生成数据的输出

        Returns
        ----------
        - `Tensor`
            - 判别器损失
        """
        real_loss = torch.mean(F.relu(1 - real_output))
        fake_loss = torch.mean(F.relu(1 + fake_output))
        return real_loss + fake_loss


class RaGAN(GANBase):
    """Relativistic average GAN (RaGAN)

    RaGAN 的判别器输出应为原始 logits (无 Sigmoid 激活)。

    References
    ----------
    Jolicoeur-Martineau, A. (2018).
    The relativistic discriminator: a key element missing from standard GAN.
    arXiv preprint arXiv:1807.00734.
    """

    gan_type = 'RaGAN'

    def __init__(
        self,
        generator: nn.Module,
        discriminator: nn.Module,
        real_data_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
        apply_spectral_norm: bool = False,
    ):
        super().__init__(generator, discriminator)
        self.bce_with_logits = nn.BCEWithLogitsLoss()

        # 模型验证
        if validate_model and real_data_shape is not None and latent_dim is not None:
            _validate_generator_discriminator(
                generator,
                discriminator,
                latent_dim,
                real_data_shape,
                gan_type=self.gan_type,
            )

        # 应用谱归一化
        if apply_spectral_norm:
            SpectralNormMixin.apply_spectral_norm(self.discriminator)

    def generator_loss(
        self, fake_output: Tensor, real_output: Optional[Tensor] = None
    ) -> Tensor:
        """计算生成器损失

        生成器的目标是让判别器认为“虚假数据比真实数据更真实”。
        """
        if real_output is None:
            raise ValueError('RaGAN requires real_output for generator loss computation')
        # real_output 和 fake_output 是判别器的原始 logits
        real_rel_logits = real_output - torch.mean(fake_output)
        fake_rel_logits = fake_output - torch.mean(real_output)

        # 生成器想让 D(G(z)) 相对于 D(x) 更真实
        g_loss_real = self.bce_with_logits(
            fake_rel_logits, torch.ones_like(fake_rel_logits)
        )
        # 生成器也想让 D(x) 相对于 D(G(z)) 更不真实
        g_loss_fake = self.bce_with_logits(
            real_rel_logits, torch.zeros_like(real_rel_logits)
        )

        return (g_loss_real + g_loss_fake) / 2

    def discriminator_loss(self, real_output: Tensor, fake_output: Tensor) -> Tensor:
        """计算判别器损失

        判别器的目标是让判别器认为“真实数据比虚假数据更真实”。
        """
        # real_output 和 fake_output 是判别器的原始 logits
        real_rel_logits = real_output - torch.mean(fake_output.detach())
        fake_rel_logits = fake_output - torch.mean(real_output.detach())

        # 判别器想让 D(x) 相对于 D(G(z)) 更真实
        d_loss_real = self.bce_with_logits(
            real_rel_logits, torch.ones_like(real_rel_logits)
        )
        # 判别器也想让 D(G(z)) 相对于 D(x) 更不真实
        d_loss_fake = self.bce_with_logits(
            fake_rel_logits, torch.zeros_like(fake_rel_logits)
        )

        return (d_loss_real + d_loss_fake) / 2

    def compute_generator_loss(
        self, z: Tensor, *, real_data: Tensor | None = None, **kwargs
    ) -> Tensor:
        """计算生成器损失 (RaGAN 特有实现)

        RaGAN 的生成器损失同时需要真实数据和虚假数据的判别结果。
        """
        fake_data = self.generate(z)

        real_output = self.discriminate(cast(Tensor, real_data)).detach()
        fake_output = self.discriminate(fake_data)

        # 在 RaGAN 中，G 的损失函数需要 real_output 和 fake_output
        g_loss = self.generator_loss(fake_output, real_output)
        return g_loss

    def compute_discriminator_loss(
        self, z: Tensor, real_data: Tensor, **kwargs
    ) -> Tensor:
        """计算判别器损失 (RaGAN 特有实现)"""
        fake_data = self.generate(z).detach()
        real_output = self.discriminate(real_data)
        fake_output = self.discriminate(fake_data)

        d_loss = self.discriminator_loss(real_output, fake_output)
        return d_loss
