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


class DDIMSampler:
    """DDIM采样器"""
    def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 50):
        self.model = model
        self.diffusion = diffusion
        self.num_inference_steps = num_inference_steps
        
        # 设置时间步
        self.set_timesteps(num_inference_steps)
    
    def set_timesteps(self, num_inference_steps: int):
        """设置推理时间步"""
        self.num_inference_steps = num_inference_steps
        
        # 选择时间步
        if self.diffusion.num_train_timesteps == num_inference_steps:
            timesteps = np.arange(0, self.diffusion.num_train_timesteps)
        else:
            step_ratio = self.diffusion.num_train_timesteps // num_inference_steps
            timesteps = np.arange(0, self.diffusion.num_train_timesteps, step_ratio)
        
        self.timesteps = torch.from_numpy(timesteps).long().flip(0)
    
    @torch.no_grad()
    def step(
        self,
        model_output: torch.Tensor,
        timestep: int,
        sample: torch.Tensor,
        eta: float = 0.0,
        use_clipped_model_output: bool = False
    ) -> torch.Tensor:
        """DDIM单步采样"""
        # 获取当前和上一个时间步
        prev_timestep = timestep - self.diffusion.num_train_timesteps // self.num_inference_steps
        
        # 提取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}")
        
        # 裁剪预测的原始样本
        if use_clipped_model_output:
            pred_original_sample = torch.clamp(pred_original_sample, -1, 1)
        
        # 计算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
        
        # 当eta > 0时，使用随机采样
        if eta > 0:
            noise = torch.randn_like(model_output)
            variance = std_dev_t ** 2
        else:
            noise = 0
            variance = 0
        
        # 计算x_t-1的均值
        pred_sample_direction = (1 - alpha_prod_t_prev - variance) ** 0.5 * pred_epsilon
        prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
        
        # 添加噪声
        if eta > 0:
            prev_sample = prev_sample + std_dev_t * noise
        
        return prev_sample
    
    @torch.no_grad()
    def sample(
        self,
        prompt_embeds: torch.Tensor,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        height: int = 512,
        width: int = 512,
        num_images_per_prompt: int = 1,
        guidance_scale: float = 7.5,
        eta: float = 0.0,
        generator: Optional[torch.Generator] = None,
        progress_bar: bool = True
    ) -> torch.Tensor:
        """生成样本"""
        # 设置模型为评估模式
        self.model.eval()
        
        # 批次大小
        batch_size = prompt_embeds.shape[0]
        
        # 初始化潜在表示
        latents = torch.randn(
            (batch_size * num_images_per_prompt, self.model.in_channels, height // 8, width // 8),
            device=prompt_embeds.device,
            generator=generator
        )
        
        # 准备额外的条件
        if negative_prompt_embeds is not None:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
        
        # 分类器自由引导的缩放因子
        do_classifier_free_guidance = guidance_scale > 1.0
        if do_classifier_free_guidance:
            latents = torch.cat([latents] * 2)
        
        # 采样循环
        timesteps = self.timesteps.to(latents.device)
        
        if progress_bar:
            timesteps = tqdm(timesteps, desc="DDIM Sampling")
        
        for i, t in enumerate(timesteps):
            # 扩展潜在表示以匹配引导的批次大小
            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
            latent_model_input = self.diffusion.scale_model_input(latent_model_input, t)
            
            # 预测噪声
            noise_pred = self.model(latent_model_input, t, prompt_embeds)
            
            # 执行分类器自由引导
            if do_classifier_free_guidance:
                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.step(noise_pred, t, latents, eta)
        
        return latents


class DPMSampler:
    """DPM采样器（更快）"""
    def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 20):
        self.model = model
        self.diffusion = diffusion
        self.num_inference_steps = num_inference_steps
    
    @torch.no_grad()
    def sample(
        self,
        prompt_embeds: torch.Tensor,
        height: int = 512,
        width: int = 512,
        guidance_scale: float = 7.5,
        progress_bar: bool = True
    ) -> torch.Tensor:
        """DPM采样"""
        self.model.eval()
        
        # 初始化潜在表示
        latents = torch.randn(
            (1, self.model.in_channels, height // 8, width // 8),
            device=prompt_embeds.device
        )
        
        # 简化的DPM采样
        timesteps = torch.linspace(1, 0, self.num_inference_steps + 1, device=latents.device)
        
        if progress_bar:
            timesteps_iter = tqdm(enumerate(timesteps[:-1]), total=len(timesteps)-1, desc="DPM Sampling")
        else:
            timesteps_iter = enumerate(timesteps[:-1])
        
        for i, t in timesteps_iter:
            # 预测噪声
            noise_pred = self.model(latents, t.unsqueeze(0) * 999, prompt_embeds)
            
            # 应用引导
            if guidance_scale > 1.0:
                # 简单引导
                noise_pred = noise_pred * guidance_scale
            
            # DPM更新步骤
            dt = timesteps[i + 1] - t
            latents = latents + dt * noise_pred
        
        return latents


class LCMSampler:
    """LCM（潜在一致性模型）采样器，极快"""
    def __init__(self, model: nn.Module, diffusion, num_inference_steps: int = 4):
        self.model = model
        self.diffusion = diffusion
        self.num_inference_steps = num_inference_steps
        
        # LCM特定的参数
        self.c_skip = 1.0
        self.c_out = 1.0
        self.c_in = 1.0
        self.c_noise = 1.0
    
    @torch.no_grad()
    def sample(
        self,
        prompt_embeds: torch.Tensor,
        height: int = 512,
        width: int = 512,
        guidance_scale: float = 7.5,
        progress_bar: bool = True
    ) -> torch.Tensor:
        """LCM采样（极快，只需要4-8步）"""
        self.model.eval()
        
        # 初始化潜在表示
        latents = torch.randn(
            (1, self.model.in_channels, height // 8, width // 8),
            device=prompt_embeds.device
        )
        
        # LCM采样循环
        timesteps = torch.linspace(1, 0, self.num_inference_steps + 1, device=latents.device)
        
        if progress_bar:
            timesteps_iter = tqdm(enumerate(timesteps[:-1]), total=len(timesteps)-1, desc="LCM Sampling")
        else:
            timesteps_iter = enumerate(timesteps[:-1])
        
        for i, t in timesteps_iter:
            # LCM特定的缩放
            c_skip = self.c_skip
            c_out = self.c_out
            c_in = self.c_in
            c_noise = self.c_noise
            
            # 缩放输入
            scaled_latents = c_in * latents
            
            # 预测
            noise_pred = self.model(scaled_latents, c_noise * t.unsqueeze(0), prompt_embeds)
            
            # LCM更新规则
            denoised = c_skip * latents + c_out * noise_pred
            
            # 更新潜在表示
            dt = timesteps[i + 1] - t
            latents = denoised + dt * noise_pred
        
        return latents


class SamplerFactory:
    """采样器工厂"""
    @staticmethod
    def create_sampler(
        sampler_type: str,
        model: nn.Module,
        diffusion,
        num_inference_steps: int = 50
    ):
        """创建采样器"""
        if sampler_type == "ddim":
            return DDIMSampler(model, diffusion, num_inference_steps)
        elif sampler_type == "dpm":
            return DPMSampler(model, diffusion, num_inference_steps)
        elif sampler_type == "lcm":
            return LCMSampler(model, diffusion, num_inference_steps)
        else:
            raise ValueError(f"未知的采样器类型: {sampler_type}")


class TextToImagePipeline:
    """文本到图像管道"""
    def __init__(
        self,
        model: nn.Module,
        diffusion,
        text_encoder,
        vae_decoder,
        sampler_type: str = "ddim",
        device: str = "cuda"
    ):
        self.model = model.to(device)
        self.diffusion = diffusion
        self.text_encoder = text_encoder
        self.vae_decoder = vae_decoder
        self.sampler_type = sampler_type
        self.device = device
        
        # 创建采样器
        self.sampler = SamplerFactory.create_sampler(
            sampler_type, model, diffusion
        )
        
        # 设置为评估模式
        self.model.eval()
        if self.vae_decoder is not None:
            self.vae_decoder.eval()
    
    @torch.no_grad()
    def __call__(
        self,
        prompt: str,
        negative_prompt: str = "",
        height: int = 512,
        width: int = 512,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        num_images: int = 1,
        seed: Optional[int] = None,
        progress_bar: bool = True
    ) -> List:
        """生成图像"""
        # 设置随机种子
        if seed is not None:
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(seed)
        
        # 编码提示
        prompt_embeds = self.text_encoder.encode([prompt]).to(self.device)
        negative_prompt_embeds = None
        
        if negative_prompt:
            negative_prompt_embeds = self.text_encoder.encode([negative_prompt]).to(self.device)
        
        # 生成潜在表示
        latents = self.sampler.sample(
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            height=height,
            width=width,
            num_images_per_prompt=num_images,
            guidance_scale=guidance_scale,
            progress_bar=progress_bar
        )
        
        # 解码为图像
        images = []
        for i in range(num_images):
            latent = latents[i:i+1]
            
            if self.vae_decoder is not None:
                image = self.vae_decoder(latent)
            else:
                # 如果没有VAE解码器，返回潜在表示
                image = latent
            
            images.append(image.cpu())
        
        return images
    
    def generate_grid(
        self,
        prompts: List[str],
        grid_size: Tuple[int, int] = (2, 2),
        **kwargs
    ) -> torch.Tensor:
        """生成图像网格"""
        images = []
        
        for prompt in prompts[:grid_size[0] * grid_size[1]]:
            image = self(prompt, **kwargs)[0]
            images.append(image)
        
        # 创建网格
        from torchvision.utils import make_grid
        grid = make_grid(torch.cat(images, dim=0), nrow=grid_size[1])
        
        return grid


def test_sampler():
    """测试采样器"""
    import torch.nn as nn
    
    # 创建模拟模型
    class MockModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.in_channels = 4
        
        def forward(self, x, t, context):
            # 返回随机噪声
            return torch.randn_like(x)
    
    # 创建模拟扩散过程
    class MockDiffusion:
        def __init__(self):
            self.num_train_timesteps = 1000
            self.alphas_cumprod = torch.ones(1000)
            self.prediction_type = "epsilon"
        
        def extract(self, a, t, x_shape):
            return torch.ones(x_shape[0], 1, 1, 1)
        
        def scale_model_input(self, x, t):
            return x
    
    model = MockModel()
    diffusion = MockDiffusion()
    
    # 测试DDIM采样器
    sampler = DDIMSampler(model, diffusion, num_inference_steps=10)
    
    # 测试采样
    prompt_embeds = torch.randn(1, 77, 768)
    latents = sampler.sample(prompt_embeds, height=64, width=64, progress_bar=False)
    
    print(f"DDIM采样完成，潜在表示形状: {latents.shape}")
    
    return sampler, latents


if __name__ == '__main__':
    test_sampler()