from abc import ABC, abstractmethod
from typing import Dict, Optional, Tuple, Union

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

from zyplib.nn.modules import BaseModule


def _validate_diffusion_model(
    model: nn.Module, input_shape: Tuple[int, ...], timesteps: int
) -> None:
    """验证扩散模型的输入输出维度是否匹配

    Parameters
    ----------
    - `model` : `nn.Module`
        - 扩散模型网络
    - `input_shape` : `Tuple[int, ...]`
        - 输入数据的形状 (不包括batch维度)
    - `timesteps` : `int`
        - 时间步数

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

    with torch.no_grad():
        try:
            # 测试模型
            test_input = torch.randn(batch_size, *input_shape, device=device)
            test_timesteps = torch.randint(0, timesteps, (batch_size,), device=device)

            output = model(test_input, test_timesteps)

            expected_shape = (batch_size, *input_shape)
            if output.shape != expected_shape:
                raise ValueError(
                    f'模型输出形状 {output.shape} 与期望形状 {expected_shape} 不匹配'
                )
        except Exception as e:
            raise ValueError(f'扩散模型无法处理形状为 {input_shape} 的输入: {e}')


class DiffusionBase(BaseModule, ABC):
    """扩散模型基类

    提供扩散模型的基本接口和通用功能。

    Parameters
    ----------
    - `model` : `nn.Module`
        - 扩散模型网络
        - **输入**: 应接受 (x, t) 其中 x 形状为 (batch_size, *input_shape)，t 为时间步
        - **输出**: 应输出形状为 (batch_size, *input_shape) 的预测结果
    - `timesteps` : `int`
        - 扩散过程的时间步数
    - `beta_schedule` : `str`, optional
        - 噪声调度类型，支持 'linear', 'cosine', by default 'linear'
    """

    def __init__(
        self, model: nn.Module, timesteps: int = 1000, beta_schedule: str = 'linear'
    ):
        super().__init__()
        self.model = model
        self.timesteps = timesteps
        self.beta_schedule = beta_schedule

        # 初始化噪声调度
        self._init_noise_schedule()

    def _init_noise_schedule(self):
        """初始化噪声调度

        使用 register_buffer 确保调度参数自动跟随模型设备变化
        """
        if self.beta_schedule == 'linear':
            # 线性调度
            beta_start = 0.0001
            beta_end = 0.02
            betas = torch.linspace(beta_start, beta_end, self.timesteps)
        elif self.beta_schedule == 'cosine':
            # 余弦调度（更平滑）
            def cosine_beta_schedule(timesteps, s=0.008):
                steps = timesteps + 1
                x = torch.linspace(0, timesteps, steps)
                alphas_cumprod = (
                    torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
                )
                alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
                betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
                return torch.clamp(betas, 0.0001, 0.9999)

            betas = cosine_beta_schedule(self.timesteps)
        else:
            raise ValueError(f'Unsupported beta schedule: {self.beta_schedule}')

        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)

        # 计算采样所需的系数
        sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
        sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)

        # 注册为 buffer，自动处理设备传输
        self.register_buffer('betas', betas)
        self.register_buffer('alphas', alphas)
        self.register_buffer('alphas_cumprod', alphas_cumprod)
        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
        self.register_buffer('sqrt_alphas_cumprod', sqrt_alphas_cumprod)
        self.register_buffer(
            'sqrt_one_minus_alphas_cumprod', sqrt_one_minus_alphas_cumprod
        )

    def add_noise(self, x_start: Tensor, noise: Tensor, timesteps: Tensor) -> Tensor:
        """向原始数据添加噪声

        Parameters
        ----------
        - `x_start` : `Tensor`
            - 原始数据
        - `noise` : `Tensor`
            - 噪声
        - `timesteps` : `Tensor`
            - 时间步

        Returns
        ----------
        - `Tensor`
            - 加噪后的数据
        """
        sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[timesteps]
        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[timesteps]

        # 调整形状以便广播
        while len(sqrt_alphas_cumprod_t.shape) < len(x_start.shape):
            sqrt_alphas_cumprod_t = sqrt_alphas_cumprod_t.unsqueeze(-1)
            sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod_t.unsqueeze(
                -1
            )

        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise

    @abstractmethod
    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算损失函数

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 批次数据

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

    @abstractmethod
    def sample(self, shape: Tuple[int, ...], device: torch.device) -> Tensor:
        """从模型采样

        Parameters
        ----------
        - `shape` : `Tuple[int, ...]`
            - 采样形状
        - `device` : `torch.device`
            - 设备

        Returns
        ----------
        - `Tensor`
            - 采样结果
        """
        pass


class DDPM(DiffusionBase):
    """Denoising Diffusion Probabilistic Models (DDPM)

    实现 DDPM 扩散模型。

    References
    ----------
    Ho, J., Jain, A., & Abbeel, P. (2020).
    Denoising diffusion probabilistic models.
    Advances in neural information processing systems, 33, 6840-6851.

    Parameters
    ----------
    - `model` : `nn.Module`
        - 扩散模型网络
    - `timesteps` : `int`, optional
        - 扩散过程的时间步数, by default 1000
    - `beta_schedule` : `str`, optional
        - 噪声调度类型，支持 'linear', 'cosine', by default 'linear'
    - `input_shape` : `Tuple[int, ...]`, optional
        - 输入数据的形状, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建 U-Net 模型
    model = UNet(
        in_channels=3,
        out_channels=3,
        time_embedding_dim=128
    )

    # 创建 DDPM
    ddpm = DDPM(
        model,
        timesteps=1000,
        beta_schedule='cosine',
        input_shape=(3, 32, 32)
    )

    # 训练
    batch_data = {
        'x': torch.randn(32, 3, 32, 32)
    }
    loss_dict = ddpm.compute_loss(batch_data)

    # 采样
    samples = ddpm.sample((4, 3, 32, 32), device='cuda')
    ```
    """

    def __init__(
        self,
        model: nn.Module,
        timesteps: int = 1000,
        beta_schedule: str = 'linear',
        input_shape: Optional[Tuple[int, ...]] = None,
        validate_model: bool = True,
    ):
        super().__init__(model, timesteps, beta_schedule)

        # 模型验证
        if validate_model and input_shape is not None:
            _validate_diffusion_model(model, input_shape, timesteps)

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算 DDPM 损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'x' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含扩散损失的字典
        """
        x_start = batch_data['x']
        batch_size = x_start.shape[0]
        device = x_start.device

        # 随机采样时间步
        t = torch.randint(0, self.timesteps, (batch_size,), device=device)

        # 生成噪声
        noise = torch.randn_like(x_start)

        # 添加噪声
        x_noisy = self.add_noise(x_start, noise, t)

        # 预测噪声
        predicted_noise = self.model(x_noisy, t)

        # 计算均方误差损失
        mse_loss = F.mse_loss(predicted_noise, noise)

        return {'mse_loss': mse_loss, 'total_loss': mse_loss}

    @torch.no_grad()
    def sample(self, shape: Tuple[int, ...], device: torch.device) -> Tensor:
        """DDPM 采样

        Parameters
        ----------
        - `shape` : `Tuple[int, ...]`
            - 采样形状
        - `device` : `torch.device`
            - 设备

        Returns
        ----------
        - `Tensor`
            - 采样结果
        """
        # 从纯噪声开始
        x = torch.randn(shape, device=device)

        # 逆向扩散过程
        for i in reversed(range(self.timesteps)):
            t = torch.full((shape[0],), i, device=device, dtype=torch.long)

            # 预测噪声
            predicted_noise = self.model(x, t)

            # 计算去噪系数
            beta_t = self.betas[i]
            alpha_t = self.alphas[i]
            alpha_cumprod_t = self.alphas_cumprod[i]

            # 去噪
            x = (1 / torch.sqrt(alpha_t)) * (
                x - (beta_t / torch.sqrt(1 - alpha_cumprod_t)) * predicted_noise
            )

            # 添加噪声（除了最后一步）
            if i > 0:
                noise = torch.randn_like(x)
                sigma_t = torch.sqrt(beta_t)
                x = x + sigma_t * noise

        return x


class DDIM(DiffusionBase):
    """Denoising Diffusion Implicit Models (DDIM)

    实现 DDIM 扩散模型，支持确定性采样和快速采样。

    References
    ----------
    Song, J., Meng, C., & Ermon, S. (2020).
    Denoising diffusion implicit models.
    arXiv preprint arXiv:2010.02502.

    Parameters
    ----------
    - `model` : `nn.Module`
        - 扩散模型网络
    - `timesteps` : `int`, optional
        - 扩散过程的时间步数, by default 1000
    - `eta` : `float`, optional
        - 随机性参数，0为确定性采样, by default 0.0
    - `beta_schedule` : `str`, optional
        - 噪声调度类型，支持 'linear', 'cosine', by default 'linear'
    - `input_shape` : `Tuple[int, ...]`, optional
        - 输入数据的形状, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建 DDIM
    ddim = DDIM(
        model,
        timesteps=1000,
        eta=0.0,  # 确定性采样
        beta_schedule='linear',
        input_shape=(3, 32, 32)
    )

    # 快速采样（50步）
    samples = ddim.sample(
        (4, 3, 32, 32),
        device='cuda',
        sampling_timesteps=50
    )
    ```
    """

    def __init__(
        self,
        model: nn.Module,
        timesteps: int = 1000,
        eta: float = 0.0,
        beta_schedule: str = 'linear',
        input_shape: Optional[Tuple[int, ...]] = None,
        validate_model: bool = True,
    ):
        super().__init__(model, timesteps, beta_schedule)
        self.eta = eta

        # 模型验证
        if validate_model and input_shape is not None:
            _validate_diffusion_model(model, input_shape, timesteps)

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算 DDIM 损失（与 DDPM 相同）

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'x' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含扩散损失的字典
        """
        x_start = batch_data['x']
        batch_size = x_start.shape[0]
        device = x_start.device

        # 随机采样时间步
        t = torch.randint(0, self.timesteps, (batch_size,), device=device)

        # 生成噪声
        noise = torch.randn_like(x_start)

        # 添加噪声
        x_noisy = self.add_noise(x_start, noise, t)

        # 预测噪声
        predicted_noise = self.model(x_noisy, t)

        # 计算均方误差损失
        mse_loss = F.mse_loss(predicted_noise, noise)

        return {'mse_loss': mse_loss, 'total_loss': mse_loss}

    @torch.no_grad()
    def sample(
        self,
        shape: Tuple[int, ...],
        device: torch.device,
        sampling_timesteps: Optional[int] = None,
    ) -> Tensor:
        """DDIM 采样

        Parameters
        ----------
        - `shape` : `Tuple[int, ...]`
            - 采样形状
        - `device` : `torch.device`
            - 设备
        - `sampling_timesteps` : `int`, optional
            - 采样时间步数，None表示使用全部时间步

        Returns
        ----------
        - `Tensor`
            - 采样结果
        """
        if sampling_timesteps is None:
            sampling_timesteps = self.timesteps

        # 创建采样时间步序列
        step_size = self.timesteps // sampling_timesteps
        timesteps_seq = list(range(0, self.timesteps, step_size))
        timesteps_seq = timesteps_seq[:sampling_timesteps]

        # 从纯噪声开始
        x = torch.randn(shape, device=device)

        # 逆向扩散过程
        for i, t in enumerate(reversed(timesteps_seq)):
            t_tensor = torch.full((shape[0],), t, device=device, dtype=torch.long)

            # 预测噪声
            predicted_noise = self.model(x, t_tensor)

            # 获取 alpha 值
            alpha_cumprod_t = self.alphas_cumprod[t]

            if i < len(timesteps_seq) - 1:
                alpha_cumprod_prev = self.alphas_cumprod[timesteps_seq[-(i + 2)]]
            else:
                alpha_cumprod_prev = torch.tensor(1.0, device=device)

            # 预测原始图像
            pred_x0 = (
                x - torch.sqrt(1 - alpha_cumprod_t) * predicted_noise
            ) / torch.sqrt(alpha_cumprod_t)

            # 计算方向向量（增强数值稳定性）
            direction_variance = (
                1 - alpha_cumprod_prev - self.eta**2 * (1 - alpha_cumprod_t)
            )
            # 确保方差非负，避免 sqrt 计算负数
            direction_variance = torch.clamp(direction_variance, min=0.0)
            direction = torch.sqrt(direction_variance) * predicted_noise

            # 添加随机噪声（仅在非最后一步且 eta > 0 时）
            noise = (
                torch.randn_like(x) if self.eta > 0 and i < len(timesteps_seq) - 1 else 0
            )

            # 更新 x
            x = (
                torch.sqrt(alpha_cumprod_prev) * pred_x0
                + direction
                + self.eta * torch.sqrt(1 - alpha_cumprod_t) * noise
            )

        return x
