import torch
import torch.nn as nn
import numpy as np
from typing import Optional, Tuple, Union
import math


class BetaScheduler:
    """Beta调度器"""
    @staticmethod
    def linear(num_timesteps: int, beta_start: float = 0.0001, beta_end: float = 0.02) -> np.ndarray:
        """线性调度"""
        return np.linspace(beta_start, beta_end, num_timesteps, dtype=np.float32)
    
    @staticmethod
    def cosine(num_timesteps: int, s: float = 0.008) -> np.ndarray:
        """余弦调度"""
        steps = num_timesteps + 1
        x = np.linspace(0, num_timesteps, steps)
        alphas_cumprod = np.cos(((x / num_timesteps) + s) / (1 + s) * np.pi * 0.5) ** 2
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
        return np.clip(betas, 0, 0.999)
    
    @staticmethod
    def scaled_linear(num_timesteps: int) -> np.ndarray:
        """缩放线性调度（Stable Diffusion默认）"""
        beta_start = 0.00085
        beta_end = 0.012
        return np.linspace(beta_start**0.5, beta_end**0.5, num_timesteps) ** 2


class DiffusionProcess:
    """扩散过程管理"""
    def __init__(self, config: dict):
        self.config = config
        diff_config = config.get('diffusion', {})
        
        self.num_train_timesteps = diff_config.get('num_train_timesteps', 1000)
        self.num_inference_timesteps = diff_config.get('num_inference_timesteps', 50)
        self.beta_start = diff_config.get('beta_start', 0.00085)
        self.beta_end = diff_config.get('beta_end', 0.012)
        self.beta_schedule = diff_config.get('beta_schedule', 'scaled_linear')
        self.prediction_type = diff_config.get('prediction_type', 'epsilon')
        
        # 初始化调度参数
        self._init_schedule()
    
    def _init_schedule(self):
        """初始化扩散调度参数"""
        # 计算betas
        if self.beta_schedule == "linear":
            betas = BetaScheduler.linear(
                self.num_train_timesteps, 
                self.beta_start, 
                self.beta_end
            )
        elif self.beta_schedule == "cosine":
            betas = BetaScheduler.cosine(self.num_train_timesteps)
        elif self.beta_schedule == "scaled_linear":
            betas = BetaScheduler.scaled_linear(self.num_train_timesteps)
        else:
            raise ValueError(f"Unknown beta schedule: {self.beta_schedule}")
        
        self.betas = torch.from_numpy(betas).float()
        
        # 计算alphas
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
        
        # 计算扩散后验方差
        self.variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        
        # 注册为buffer
        self.register_buffer = lambda name, tensor: setattr(self, name, tensor)
        self.register_buffer('betas', self.betas)
        self.register_buffer('alphas', self.alphas)
        self.register_buffer('alphas_cumprod', self.alphas_cumprod)
        self.register_buffer('alphas_cumprod_prev', self.alphas_cumprod_prev)
        self.register_buffer('variance', self.variance)
        
        # 计算采样系数
        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(self.alphas_cumprod))
        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1.0 - self.alphas_cumprod))
        self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1.0 - self.alphas_cumprod))
        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1.0 / self.alphas_cumprod))
        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1.0 / self.alphas_cumprod - 1))
    
    def q_sample(self, x_start: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
        """前向扩散过程：加噪"""
        if noise is None:
            noise = torch.randn_like(x_start)
        
        sqrt_alphas_cumprod_t = self.extract(self.sqrt_alphas_cumprod, t, x_start.shape)
        sqrt_one_minus_alphas_cumprod_t = self.extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
        
        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
    
    def extract(self, a: torch.Tensor, t: torch.Tensor, x_shape: Tuple[int, ...]) -> torch.Tensor:
        """从张量a中提取索引t处的值"""
        batch_size = t.shape[0]
        out = a.gather(-1, t.cpu())
        return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
    
    def get_loss_weight(self, snr: torch.Tensor, gamma: float = 5.0) -> torch.Tensor:
        """根据SNR计算损失权重"""
        if gamma is None:
            return torch.ones_like(snr)
        
        snr = torch.clamp(snr, min=1e-8)
        min_snr = torch.tensor(gamma, device=snr.device)
        weight = torch.minimum(snr, min_snr) / snr
        return weight
    
    def compute_snr(self, timesteps: torch.Tensor) -> torch.Tensor:
        """计算信噪比(SNR)"""
        alphas_cumprod = self.extract(self.alphas_cumprod, timesteps, timesteps.shape)
        snr = alphas_cumprod / (1 - alphas_cumprod)
        return snr


class DDIMScheduler:
    """DDIM采样器"""
    def __init__(self, diffusion: DiffusionProcess):
        self.diffusion = diffusion
        self.num_train_timesteps = diffusion.num_train_timesteps
        self.num_inference_timesteps = diffusion.num_inference_timesteps
        
        # 设置时间步
        self.set_timesteps(self.num_inference_timesteps)
    
    def set_timesteps(self, num_inference_timesteps: int):
        """设置推理时间步"""
        self.num_inference_timesteps = num_inference_timesteps
        
        # 选择时间步
        if self.num_train_timesteps == self.num_inference_timesteps:
            self.timesteps = torch.arange(0, self.num_train_timesteps).long()
        else:
            step_ratio = self.num_train_timesteps // self.num_inference_timesteps
            self.timesteps = torch.arange(0, self.num_train_timesteps, step_ratio).long()
        
        self.timesteps = self.timesteps.flip(0)  # 从T到0
    
    @torch.no_grad()
    def step(self, model_output: torch.Tensor, timestep: int, sample: torch.Tensor, eta: float = 0.0) -> torch.Tensor:
        """DDIM单步采样"""
        # 获取当前时间步的参数
        prev_timestep = timestep - self.num_train_timesteps // self.num_inference_timesteps
        
        # 提取alpha参数
        alpha_prod_t = self.diffusion.extract(self.diffusion.alphas_cumprod, timestep, sample.shape)
        alpha_prod_t_prev = self.diffusion.extract(
            self.diffusion.alphas_cumprod, 
            prev_timestep, 
            sample.shape
        ) if prev_timestep >= 0 else torch.ones_like(alpha_prod_t)
        
        # 根据预测类型处理模型输出
        if self.diffusion.prediction_type == "epsilon":
            pred_original_sample = (sample - (1 - alpha_prod_t) ** 0.5 * model_output) / alpha_prod_t ** 0.5
            pred_epsilon = model_output
        elif self.diffusion.prediction_type == "sample":
            pred_original_sample = model_output
            pred_epsilon = (sample - alpha_prod_t ** 0.5 * pred_original_sample) / (1 - alpha_prod_t) ** 0.5
        elif self.diffusion.prediction_type == "v_prediction":
            pred_original_sample = (alpha_prod_t ** 0.5) * sample - (1 - alpha_prod_t) ** 0.5 * model_output
            pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (1 - alpha_prod_t) ** 0.5 * sample
        else:
            raise ValueError(f"Unsupported prediction type: {self.diffusion.prediction_type}")
        
        # 计算x_t-1的方差
        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
        std_dev_t = eta * variance ** 0.5
        
        # 计算x_t-1的均值
        pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon
        prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
        
        # 添加噪声
        if eta > 0:
            noise = torch.randn_like(model_output)
            prev_sample = prev_sample + std_dev_t * noise
        
        return prev_sample


class DiffusionModel(nn.Module):
    """扩散模型封装"""
    def __init__(self, unet: nn.Module, diffusion: DiffusionProcess):
        super().__init__()
        self.unet = unet
        self.diffusion = diffusion
        self.scheduler = DDIMScheduler(diffusion)
        
    def forward(self, x: torch.Tensor, timesteps: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
        """前向传播：预测噪声"""
        return self.unet(x, timesteps, context)
    
    def compute_loss(self, x_start: torch.Tensor, context: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
        """计算扩散损失"""
        if noise is None:
            noise = torch.randn_like(x_start)
        
        # 随机采样时间步
        batch_size = x_start.shape[0]
        timesteps = torch.randint(
            0, self.diffusion.num_train_timesteps, 
            (batch_size,), device=x_start.device
        ).long()
        
        # 前向扩散
        x_noisy = self.diffusion.q_sample(x_start, timesteps, noise)
        
        # 预测噪声
        predicted_noise = self.unet(x_noisy, timesteps, context)
        
        # 计算损失
        loss = F.mse_loss(predicted_noise, noise)
        
        return loss
    
    @torch.no_grad()
    def generate(
        self, 
        context: torch.Tensor, 
        num_samples: int = 1, 
        height: int = 512, 
        width: int = 512,
        guidance_scale: float = 7.5
    ) -> torch.Tensor:
        """生成图像"""
        # 初始化噪声
        latents = torch.randn(
            (num_samples, self.unet.in_channels, height // 8, width // 8),
            device=next(self.unet.parameters()).device
        )
        
        # DDIM采样
        self.scheduler.set_timesteps(self.diffusion.num_inference_timesteps)
        
        for t in self.scheduler.timesteps:
            # 扩展latents以匹配批大小
            latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents
            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
            
            # 预测噪声
            timesteps = torch.full((num_samples,), t, device=latents.device).long()
            if guidance_scale > 1.0:
                timesteps = torch.cat([timesteps] * 2)
            
            noise_pred = self.unet(latent_model_input, timesteps, context)
            
            # 应用分类器自由引导
            if guidance_scale > 1.0:
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            
            # 计算上一个样本
            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
        
        return latents