import typing
from abc import ABC, abstractmethod
from typing import (
    Any,
    Callable,
    Dict,
    Literal,
    Optional,
    Tuple,
    Union,
    TYPE_CHECKING,
    cast,
)

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

from zyplib.nn.modules import BaseModule


if TYPE_CHECKING:
    from zyplib.train.common import FabricTrainingContext


LossFn = Callable[..., Tensor]
AutoEncoderLoss = Literal['mse', 'l1', 'cos'] | LossFn
VariationalLoss = Literal['mse', 'bce'] | LossFn
BatchTransform = Callable[[Tuple[Tensor, ...]], Any]


def _validate_encoder_decoder(
    encoder: nn.Module,
    decoder: nn.Module,
    input_shape: Tuple[int, ...],
    latent_dim: int,
    is_vae: bool = False,
) -> None:
    """验证编码器和解码器的输入输出维度是否匹配

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `decoder` : `nn.Module`
        - 解码器网络
    - `input_shape` : `Tuple[int, ...]`
        - 输入数据的形状 (不包括batch维度)
    - `latent_dim` : `int`
        - 潜在空间维度
    - `is_vae` : `bool`, optional
        - 是否为VAE模型，VAE编码器输出应为latent_dim*2, by default False

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

    with torch.no_grad():
        # 创建测试输入
        test_input = torch.randn(batch_size, *input_shape, device=device)

        # 测试编码器
        try:
            encoder_output = encoder(test_input)
        except Exception as e:
            raise ValueError(f'编码器无法处理形状为 {input_shape} 的输入: {e}')

        # 检查编码器输出维度
        if is_vae:
            expected_encoder_dim = latent_dim * 2  # mu + logvar
            if encoder_output.shape[-1] != expected_encoder_dim:
                raise ValueError(
                    f'VAE编码器输出维度应为 {expected_encoder_dim} (latent_dim*2), '
                    f'但实际为 {encoder_output.shape[-1]}'
                )
            # 为解码器测试创建潜在向量
            test_latent = torch.randn(batch_size, latent_dim, device=device)
        else:
            if encoder_output.shape[-1] != latent_dim:
                raise ValueError(
                    f'编码器输出维度应为 {latent_dim}, 但实际为 {encoder_output.shape[-1]}'
                )
            test_latent = encoder_output

        # 测试解码器
        try:
            decoder_output = decoder(test_latent)
        except Exception as e:
            raise ValueError(f'解码器无法处理形状为 {test_latent.shape} 的输入: {e}')

        # 检查解码器输出维度
        expected_output_shape = (batch_size,) + input_shape
        if decoder_output.shape != expected_output_shape:
            raise ValueError(
                f'解码器输出形状应为 {expected_output_shape}, '
                f'但实际为 {decoder_output.shape}'
            )


def _validate_conditional_encoder_decoder(
    encoder: nn.Module,
    decoder: nn.Module,
    latent_dim: int,
    example_input_x: Tensor,
    example_input_c: Tensor,
) -> None:
    """验证条件VAE的编码器和解码器维度

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `decoder` : `nn.Module`
        - 解码器网络
    - `latent_dim` : `int`
        - 潜在空间维度
    - `example_input_x` : `Tensor`
        - 示例输入数据
    - `example_input_c` : `Tensor`
        - 示例条件向量

    Raises
    ----------
    - `ValueError`
        - 当模型维度不匹配时抛出
    """
    device = next(encoder.parameters()).device
    batch_size = 2

    # 调整示例输入的 batch size 以匹配测试
    if example_input_x.shape[0] != batch_size:
        shape_x = example_input_x.shape[1:]
        example_input_x = torch.randn(batch_size, *shape_x, device=device)
    if example_input_c.shape[0] != batch_size:
        shape_c = example_input_c.shape[1:]
        example_input_c = torch.randn(batch_size, *shape_c, device=device)

    with torch.no_grad():
        # --- 编码器验证 ---
        try:
            # 直接传递 x 和 c
            encoder_output = encoder(example_input_x, example_input_c)
        except Exception as e:
            raise ValueError(
                f'编码器无法处理 (x, c) 输入。请确保 encoder.forward(self, x, c) 定义正确。错误: {e}'
            )

        expected_encoder_dim = latent_dim * 2
        if encoder_output.shape[-1] != expected_encoder_dim:
            raise ValueError(
                f'编码器输出维度应为 {expected_encoder_dim}, 但实际为 {encoder_output.shape[-1]}'
            )

        # --- 解码器验证 ---
        test_latent = torch.randn(batch_size, latent_dim, device=device)
        try:
            # 直接传递 z 和 c
            decoder_output = decoder(test_latent, example_input_c)
        except Exception as e:
            raise ValueError(
                f'解码器无法处理 (z, c) 输入。请确保 decoder.forward(self, z, c) 定义正确。错误: {e}'
            )

        expected_output_shape = example_input_x.shape
        if decoder_output.shape != expected_output_shape:
            raise ValueError(
                f'解码器输出形状应为 {expected_output_shape}, 但实际为 {decoder_output.shape}'
            )


class AutoEncoderBase(BaseModule, ABC):
    """自动编码器基类

    提供自动编码器的基本接口和通用功能。

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `decoder` : `nn.Module`
        - 解码器网络
    """

    def __init__(self, encoder: nn.Module, decoder: nn.Module):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    @abstractmethod
    def encode(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, ...]]:
        """编码输入数据

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

        Returns
        ----------
        - `Union[Tensor, Tuple[Tensor, ...]]`
            - 编码后的潜在表示
        """
        pass

    @abstractmethod
    def decode(self, z: Tensor) -> Tensor:
        """解码潜在表示

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在表示

        Returns
        ----------
        - `Tensor`
            - 重构的数据
        """
        pass

    @abstractmethod
    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        """前向传播

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

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含重构结果和其他相关信息的字典
        """
        pass

    @abstractmethod
    def compute_loss(
        self, x: Tensor, output: typing.Mapping[str, Tensor], **kwargs
    ) -> Dict[str, Tensor]:
        """计算损失函数

        Parameters
        ----------
        - `x` : `Tensor`
            - 原始输入数据
        - `output` : `Dict[str, Tensor]`
            - 模型前向传播的输出

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含各种损失项的字典
        """
        pass


def cos_loss(X_recon: Tensor, X: Tensor, dim: int = -1) -> Tensor:
    cos = F.cosine_similarity(X_recon, X, dim=dim)
    return 1 - cos.mean()


class AutoEncoder(AutoEncoderBase):
    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        reconstruction_loss: AutoEncoderLoss = 'mse',
        input_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
    ):
        """基本自动编码器

        实现标准的自动编码器，使用重构损失进行训练。

        References
        ----------
        Hinton, G. E., & Salakhutdinov, R. R. (2006).
        Reducing the dimensionality of data with neural networks.
        Science, 313(5786), 504-507.

        Parameters
        ----------
        - `encoder` : `nn.Module`
            - 编码器网络
            - **输入维度**: 应接受形状为 `(batch_size, *input_shape)` 的输入
            - **输出维度**: 应输出形状为 `(batch_size, latent_dim)` 的潜在表示
        - `decoder` : `nn.Module`
            - 解码器网络
            - **输入维度**: 应接受形状为 `(batch_size, latent_dim)` 的潜在表示
            - **输出维度**: 应输出形状为 `(batch_size, *input_shape)` 的重构数据
        - `reconstruction_loss` : `str`, optional
            - 重构损失类型，支持 'mse', 'l1', by default 'mse'
        - `input_shape` : `Tuple[int, ...]`, optional
            - 输入数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `latent_dim` : `int`, optional
            - 潜在空间维度，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True

        Raises
        ----------
        - `ValueError`
            - 当模型维度不匹配且 validate_model=True 时抛出
        """
        super().__init__(encoder, decoder)
        self.reconstruction_loss = reconstruction_loss
        self.loss_fn: LossFn

        if callable(reconstruction_loss):
            self.loss_fn = reconstruction_loss
        elif reconstruction_loss == 'mse':
            self.loss_fn = F.mse_loss
        elif reconstruction_loss == 'l1':
            self.loss_fn = F.l1_loss
        elif reconstruction_loss == 'cos':
            self.loss_fn = cos_loss
        else:
            raise ValueError(f'Unsupported reconstruction loss: {reconstruction_loss}')

        # 模型验证
        if validate_model and input_shape is not None and latent_dim is not None:
            _validate_encoder_decoder(
                encoder, decoder, input_shape, latent_dim, is_vae=False
            )

    def encode(self, x: Tensor) -> Tensor:
        """编码输入数据

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

        Returns
        ----------
        - `Tensor`
            - 编码后的潜在表示
        """
        return self.encoder(x)

    def decode(self, z: Tensor) -> Tensor:
        """解码潜在表示

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在表示

        Returns
        ----------
        - `Tensor`
            - 重构的数据
        """
        return self.decoder(z)

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        """前向传播

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

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 'z' (潜在表示) 和 'x_recon' (重构数据) 的字典
        """
        z = self.encode(x)
        x_recon = self.decode(z)

        return {'z': z, 'x_recon': x_recon}

    def compute_loss(
        self, x: Tensor, output: typing.Mapping[str, Tensor]
    ) -> Dict[str, Tensor]:
        """计算重构损失

        Parameters
        ----------
        - `x` : `Tensor`
            - 原始输入数据
        - `output` : `Dict[str, Tensor]`
            - 模型前向传播的输出
        - `reduction` : `str`, optional
            - 损失函数的归约方式, by default 'mean'

        Returns
        ----------
        - `loss`: `Dict[str, Tensor]`, 重构损失; 唯一 key `total_loss`
        """
        x_recon = output['x_recon']
        recon_loss = self.loss_fn(x_recon, x)

        return {'total_loss': recon_loss}


class AutoEncoderWithLatentReg(AutoEncoder):
    """在 AutoEncoder 基础上, 增加"""

    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        reconstruction_loss: AutoEncoderLoss = 'mse',
        lambda_reg: float = 0.5,
        input_shape: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        validate_model: bool = True,
    ) -> None:
        super().__init__(
            encoder, decoder, reconstruction_loss, input_shape, latent_dim, validate_model
        )
        self.lambda_reg = lambda_reg
        self.reg_loss_fn = F.mse_loss

    def compute_loss(
        self, x: Tensor, output: typing.Mapping[str, Tensor]
    ) -> Dict[str, Tensor]:
        """
        计算总损失，包括重构损失和潜在空间一致性损失。

        Parameters
        ----------
        - `x` : `Tensor`
            - 原始输入数据
        - `output` : `Dict[str, Tensor]`
            - 模型前向传播的输出，应包含 'z' 和 'x_recon'

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 'total_loss', 'recon_loss', 'latent_reg_loss' 的字典
        """
        # 1. 计算基础的重构损失 (调用父类的方法)
        base_loss_dict = super().compute_loss(x, output)
        recon_loss = base_loss_dict['total_loss']

        # 2. 提取计算正则化损失所需的张量
        z = output['z']
        x_recon = output['x_recon']

        # 3. 计算潜在空间一致性损失 (Latent Consistency Loss)
        #    将重构数据 x_recon 再次送入编码器
        z_recon = self.encode(x_recon)

        # 计算 z 和 z_recon 之间的差异。
        # 使用 z.detach() 将 z 视为一个固定的目标，
        # 这样梯度就不会通过 z 回传，避免训练不稳定
        latent_reg_loss = self.reg_loss_fn(z_recon, z.detach())

        # 4. 组合总损失
        total_loss = recon_loss + self.lambda_reg * latent_reg_loss

        return {
            'total_loss': total_loss,
            'recon_loss': recon_loss,
            'latent_reg_loss': latent_reg_loss,
        }


class VariationalAutoEncoder(AutoEncoderBase):
    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        latent_dim: int,
        beta: float = 1.0,
        reconstruction_loss: VariationalLoss = 'mse',
        input_shape: Optional[Tuple[int, ...]] = None,
        validate_model: bool = True,
        skip_validation: bool = False,
    ):
        """变分自动编码器 (VAE)

        实现变分自动编码器，使用重构损失和KL散度损失进行训练。

        References
        ----------
        Kingma, D. P., & Welling, M. (2013).
        Auto-encoding variational bayes.
        arXiv preprint arXiv:1312.6114.

        Parameters
        ----------
        - `encoder` : `nn.Module`
            - 编码器网络
            - **输入维度**: 应接受形状为 `(batch_size, *input_shape)` 的输入
            - **输出维度**: 两种
                1. 输出形状为 `(batch_size, latent_dim*2)` 的向量（前半部分为 mu，后半部分为 logvar）
                2. 输出两个张量 `(mu: Tensor, logvar: Tensor)`，其中 `mu` 和 `logvar` 的形状均为 `(batch_size, latent_dim)`。
        - `decoder` : `nn.Module`
            - 解码器网络
            - **输入维度**: 应接受形状为 `(batch_size, latent_dim)` 的潜在表示
            - **输出维度**: 应输出形状为 `(batch_size, *input_shape)` 的重构数据
        - `latent_dim` : `int`
            - 潜在空间维度
        - `beta` : `float`, optional
            - KL散度损失的权重系数, by default 1.0
        - `reconstruction_loss` : `str`, optional
            - 重构损失类型，支持 'mse', 'bce', by default 'mse'
        - `input_shape` : `Tuple[int, ...]`, optional
            - 输入数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True
        - `skip_validation` : `bool`, optional
            - 是否跳过模型验证, by default False

        Raises
        ----------
        - `ValueError`
            - 当模型维度不匹配且 validate_model=True 时抛出
        """
        super().__init__(encoder, decoder)
        self.latent_dim = latent_dim
        self.beta = beta
        self.reconstruction_loss = reconstruction_loss
        self.loss_fn: LossFn

        if callable(reconstruction_loss):
            self.loss_fn = reconstruction_loss
        elif reconstruction_loss == 'mse':
            self.loss_fn = F.mse_loss
        elif reconstruction_loss == 'bce':
            self.loss_fn = F.binary_cross_entropy
        else:
            raise ValueError(f'Unsupported reconstruction loss: {reconstruction_loss}')

        # 模型验证
        if validate_model and input_shape is not None and not skip_validation:
            _validate_encoder_decoder(
                encoder, decoder, input_shape, latent_dim, is_vae=True
            )

    def encode(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        """编码输入数据为均值和对数方差

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

        Returns
        ----------
        - `Tuple[Tensor, Tensor]`
            - (mu, logvar) 均值和对数方差
        """
        h = self.encoder(x)
        if isinstance(h, tuple) and len(h) == 2:
            mu, logvar = h
        else:
            h = cast(Tensor, h)
            mu = h[:, : self.latent_dim]
            logvar = h[:, self.latent_dim :]
        return mu, logvar

    def reparameterize(self, mu: Tensor, logvar: Tensor) -> Tensor:
        """重参数化技巧

        Parameters
        ----------
        - `mu` : `Tensor`
            - 均值
        - `logvar` : `Tensor`
            - 对数方差

        Returns
        ----------
        - `Tensor`
            - 采样的潜在变量
        """
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z: Tensor) -> Tensor:
        """解码潜在表示

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在表示

        Returns
        ----------
        - `Tensor`
            - 重构的数据
        """
        return self.decoder(z)

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        """前向传播

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

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 'mu', 'logvar', 'z', 'x_recon' 的字典
        """
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z)

        return {'mu': mu, 'logvar': logvar, 'z': z, 'x_recon': x_recon}

    def compute_loss(
        self, x: Tensor, output: typing.Mapping[str, Tensor]
    ) -> Dict[str, Tensor]:
        """计算VAE损失 (重构损失 + KL散度损失)

        Parameters
        ----------
        - `x` : `Tensor`
            - 原始输入数据
        - `output` : `Dict[str, Tensor]`
            - 模型前向传播的输出

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 'reconstruction_loss', 'kl_loss', 'total_loss' 的字典
        """
        x_recon = output['x_recon']
        mu = output['mu']  # [batch_size, latent_dim]
        logvar = output['logvar']  # [batch_size, latent_dim]

        # 重构损失
        recon_loss = self.loss_fn(x_recon, x, reduction='mean')

        # KL散度损失
        # kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
        kl_per_sample = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
        kl_loss = torch.mean(kl_per_sample)

        # 总损失
        total_loss = recon_loss + self.beta * kl_loss

        return {
            'reconstruction_loss': recon_loss,
            'kl_loss': kl_loss,
            'total_loss': total_loss,
        }

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

        Parameters
        ----------
        - `num_samples` : `int`
            - 采样数量
        - `device` : `torch.device`, optional
            - 设备, by default None

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

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


class BetaVAE(VariationalAutoEncoder):
    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        latent_dim: int,
        beta: float = 4.0,
        reconstruction_loss: VariationalLoss = 'mse',
        input_shape: Optional[Tuple[int, ...]] = None,
        validate_model: bool = True,
    ):
        """Beta-VAE

        Beta-VAE 是 VAE 的变种，通过调整 beta 参数来控制解耦表示的学习。

        References
        ----------
        Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., ... & Lerchner, A. (2017).
        beta-VAE: Learning basic visual concepts with a constrained variational framework.
        In International Conference on Learning Representations.

        Parameters
        ----------
        - `encoder` : `nn.Module`
            - 编码器网络
        - `decoder` : `nn.Module`
            - 解码器网络
        - `latent_dim` : `int`
            - 潜在空间维度
        - `beta` : `float`, optional
            - KL散度损失的权重系数，通常 > 1.0 用于解耦, by default 4.0
        - `reconstruction_loss` : `str`, optional
            - 重构损失类型, by default 'mse'
        - `input_shape` : `Tuple[int, ...]`, optional
            - 输入数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True

        Examples
        ----------
        ```python
        # 创建 Beta-VAE，使用更高的 beta 值促进解耦
        beta_vae = BetaVAE(encoder, decoder, latent_dim=64, beta=4.0)
        ```
        """
        super().__init__(
            encoder,
            decoder,
            latent_dim,
            beta,
            reconstruction_loss,
            input_shape,
            validate_model,
        )


class ConditionalVAE(VariationalAutoEncoder):
    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        latent_dim: int,
        condition_dim: int,
        beta: float = 1.0,
        reconstruction_loss: VariationalLoss = 'mse',
        input_shape: Optional[Tuple[int, ...]] = None,
        validate_model: bool = True,
    ):
        """条件变分自动编码器 (CVAE)

        在VAE基础上添加条件信息，可以根据给定条件生成特定类型的数据。

        References
        ----------
        Sohn, K., Lee, H., & Yan, X. (2015).
        Learning structured output representation using deep conditional generative models.
        Advances in neural information processing systems, 28, 3483-3491.

        Parameters
        ----------
        - `encoder` : `nn.Module`
            - 编码器网络
            - **forward 方法签名**: 应接受 `(x: Tensor, c: Tensor)` 两个参数。
            - **输出维度**: 两种
                1. 输出形状为 `(batch_size, latent_dim*2)` 的向量（前半部分为 mu，后半部分为 logvar）
                2. 输出两个张量 `(mu: Tensor, logvar: Tensor)`，其中 `mu` 和 `logvar` 的形状均为 `(batch_size, latent_dim)`。
            - **注意**: 编码器内部应自行处理 `x` 和 `c` 的组合方式（例如拼接、FiLM 等）。
        - `decoder` : `nn.Module`
            - 解码器网络
            - **forward 方法签名**: 应接受 `(z: Tensor, c: Tensor)` 两个参数。
            - **输出维度**: 应输出形状为 `(batch_size, *input_shape)` 的重构数据。
            - **注意**: 解码器内部应自行处理 `z` 和 `c` 的组合方式。
        - `latent_dim` : `int`
            - 潜在空间维度
        - `condition_dim` : `int`
            - 条件向量维度
        - `beta` : `float`, optional
            - KL散度损失的权重系数, by default 1.0
        - `reconstruction_loss` : `str`, optional
            - 重构损失类型, by default 'mse'
        - `input_shape` : `Tuple[int, ...]`, optional
            - 输入数据的形状 (不包括batch维度)，用于模型验证, by default None
        - `validate_model` : `bool`, optional
            - 是否在初始化时验证模型维度, by default True

        Raises
        ----------
        - `ValueError`
            - 当模型维度不匹配且 validate_model=True 时抛出
        """
        # 跳过父类的标准VAE验证
        super().__init__(
            encoder,
            decoder,
            latent_dim,
            beta,
            reconstruction_loss,
            input_shape,
            validate_model,
            skip_validation=True,
        )
        self.condition_dim = condition_dim

        # 条件VAE的模型验证
        if validate_model and input_shape is not None:
            example_input_x = torch.randn(1, *input_shape)
            example_input_c = torch.randn(1, condition_dim)
            _validate_conditional_encoder_decoder(
                encoder, decoder, latent_dim, example_input_x, example_input_c
            )

    @override
    def encode(self, x: Tensor, c: Tensor) -> Tuple[Tensor, Tensor]:
        """编码输入数据和条件

        Parameters
        ----------
        - `x` : `Tensor`
            - 输入数据
        - `c` : `Tensor`
            - 条件向量

        Returns
        ----------
        - `Tuple[Tensor, Tensor]`
            - (mu, logvar) 均值和对数方差
        """
        h = self.encoder(x, c)
        if isinstance(h, tuple) and len(h) == 2:
            mu, logvar = h
        else:
            h = cast(Tensor, h)
            mu = h[:, : self.latent_dim]
            logvar = h[:, self.latent_dim :]
        return mu, logvar

    @override
    def decode(self, z: Tensor, c: Tensor) -> Tensor:
        """解码潜在表示和条件

        Parameters
        ----------
        - `z` : `Tensor`
            - 潜在表示
        - `c` : `Tensor`
            - 条件向量

        Returns
        ----------
        - `Tensor`
            - 重构的数据
        """
        return self.decoder(z, c)

    def forward(self, x: Tensor, c: Tensor) -> Dict[str, Tensor]:
        """前向传播

        Parameters
        ----------
        - `x` : `Tensor`
            - 输入数据
        - `c` : `Tensor`
            - 条件向量

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 'mu', 'logvar', 'z', 'x_recon' 的字典
        """
        mu, logvar = self.encode(x, c)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z, c)

        return {'mu': mu, 'logvar': logvar, 'z': z, 'x_recon': x_recon}

    def sample(
        self,
        c: Tensor,
        num_samples: Optional[int] = None,
        device: Optional[torch.device] = None,
    ) -> Tensor:
        """根据条件采样生成新数据

        Parameters
        ----------
        - `c` : `Tensor`
            - 条件向量，形状为 (batch_size, condition_dim) 或 (condition_dim,)
        - `num_samples` : `int`, optional
            - 采样数量，如果 c 是单个条件向量则需要指定, by default None
        - `device` : `torch.device`, optional
            - 设备, by default None

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

        if c.dim() == 1:
            # 单个条件向量，需要扩展
            if num_samples is None:
                raise ValueError(
                    'num_samples must be specified when c is a single condition vector'
                )
            c = c.unsqueeze(0).expand(num_samples, -1)

        batch_size = c.size(0)
        z = torch.randn(batch_size, self.latent_dim, device=device)
        return self.decode(z, c)


class BasicAutoEncoderProcedure:
    """自动编码器基础训练流程

    提供标准的自动编码器训练和验证步骤，支持 AE、VAE 等各种变体。
    """

    def __init__(
        self,
        loss_key: Optional[str] = 'total_loss',
        fn_batch2data: Optional[BatchTransform] = None,
        fn_batch2forward: Optional[BatchTransform] = None,
    ) -> None:
        """初始化训练流程

        Parameters
        ----------
        - `loss_key` : `str`, optional
            - 用于反向传播的损失键名，如果为 None 则使用损失字典中的第一个值, by default 'total_loss'
        - `fn_batch2data` : `Callable`, optional
            - 批次数据转换函数, 默认为 `lambda batch: batch`
        - `fn_batch2forward` : `Callable`, optional
            - 批次数据到模型 forward 方法的转换函数，默认为 `lambda batch: batch`
        """
        self.loss_key = loss_key
        self.fn_batch2data: BatchTransform = (
            fn_batch2data if fn_batch2data is not None else lambda batch: batch
        )
        self.fn_batch2forward: BatchTransform = (
            fn_batch2forward if fn_batch2forward is not None else lambda batch: batch
        )

    def autoencoder_train_step(
        self, ctx: 'FabricTrainingContext', batch: Tuple[Tensor, ...], **kwargs
    ) -> Tuple[Dict[str, Tensor], Dict[str, float]]:
        """自动编码器训练步骤

        Parameters
        ----------
        - `ctx` : `FabricTrainingContext`
            - Fabric 训练上下文
        - `batch` : `Tuple[Tensor, ...]`
            - 批次数据，对于 AE/VAE 通常为 (x,) 或 (x, y)，对于 CVAE 为 (x, c)
        - `epoch_idx` : `int`
            - 当前 epoch 索引
        - `**kwargs`
            - 其他参数

        Returns
        ----------
        - `Tuple[Dict[str, Tensor], Dict[str, float]]`
            - (模型输出字典, 损失字典)
        """
        ae_model = cast(AutoEncoderBase, ctx.model)
        optimizer = cast(Any, ctx.optimizer)
        fabric = ctx.fabric

        ae_model.train()
        optimizer.zero_grad()

        # 提取输入数据
        x = cast(Tensor, self.fn_batch2data(batch))
        forward_args = self.fn_batch2forward(batch)
        output = cast(Dict[str, Tensor], ae_model.forward(cast(Any, forward_args)))
        loss_dict = ae_model.compute_loss(x, output)

        # 选择用于反向传播的损失
        if self.loss_key is not None and self.loss_key in loss_dict:
            train_loss = loss_dict[self.loss_key]
        else:
            # 使用第一个损失值
            train_loss = next(iter(loss_dict.values()))

        # 反向传播
        if fabric:
            fabric.backward(train_loss)
        else:
            train_loss.backward()

        optimizer.step()

        # 转换为标量用于记录
        loss_scalars = {
            k: v.item() if isinstance(v, Tensor) else v for k, v in loss_dict.items()
        }

        return output, loss_scalars

    @torch.no_grad()
    def autoencoder_val_step(
        self,
        ctx: 'FabricTrainingContext',
        batch: Tuple[Tensor, ...],
        epoch_idx: Optional[int] = None,
        **kwargs,
    ) -> Tuple[Dict[str, Tensor], Dict[str, float]]:
        """自动编码器验证步骤

        Parameters
        ----------
        - `ctx` : `FabricTrainingContext`
            - Fabric 训练上下文
        - `batch` : `Tuple[Tensor, ...]`
            - 批次数据
        - `epoch_idx` : `int`, optional
            - 当前 epoch 索引, by default None
        - `**kwargs`
            - 其他参数

        Returns
        ----------
        - `Tuple[Dict[str, Tensor], Dict[str, float]]`
            - (模型输出字典, 损失字典)
        """
        ae_model = cast(AutoEncoderBase, ctx.model)

        ae_model.eval()

        # 提取输入数据
        x = cast(Tensor, self.fn_batch2data(batch))

        forward_args = self.fn_batch2forward(batch)
        output = cast(Dict[str, Tensor], ae_model.forward(cast(Any, forward_args)))
        loss_dict = ae_model.compute_loss(x, output)

        # 转换为标量用于记录
        loss_scalars = {
            k: v.item() if isinstance(v, Tensor) else v for k, v in loss_dict.items()
        }

        return output, loss_scalars
