# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/scheduler/scheduler_ddpm.ipynb.

# %% auto 0
__all__ = ['DDPMSchedulerOutput', 'DDPMScheduler']

# %% ../../src/scheduler/scheduler_ddpm.ipynb 3
from ..imports import *
from .scheduler import Scheduler

# %% ../../src/scheduler/scheduler_ddpm.ipynb 4
@dataclass
class DDPMSchedulerOutput:
    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None
    
class DDPMScheduler(Scheduler):
    """A `Scheduler` implementing [(DDPM)](https://arxiv.org/abs/2006.11239)"""
    
    non_blocking = False
    
    def __init__(self, 
                 device: Union[str, torch.device],     
                 num_train_timesteps: int = 1000,
                 beta_start: float = 0.0001,
                 beta_end: float = 0.02,
                 beta_schedule: str = "linear",
                 input_perturbation = 0.1
                ):    
        super().__init__()
        self.device = device        
        self.num_train_timesteps = torch.tensor(num_train_timesteps)
        self.num_inference_steps = torch.tensor(num_train_timesteps)
        
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.beta_schedule = beta_schedule

        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) #careful is defined reversed for easy denoising looping
            
        if beta_schedule == "linear":
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
            
        elif beta_schedule == "linear_sqrt":
            self.betas = torch.linspace(beta_start ** 0.5, beta_end ** 0.5, num_train_timesteps, dtype=torch.float32) ** 2
                      
        elif beta_schedule == "cos_alpha":  #cosine-based-variance                      
            f = lambda t: np.cos((t/self.num_train_timesteps + 0.008)*np.pi/2.016)**2 
            _betas = []            
            for i in range(self.num_train_timesteps):                           
                _betas.append(min(1.0-(f(i+1.0)/f(i)),0.999))                
            self.betas = torch.tensor(_betas, dtype=torch.float32)                       
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
              
        self.sigmas = torch.sqrt(self.betas)
                 
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)  
        
        ## Is this used anymore?
        self.sqrt_alphas                   = torch.sqrt(self.alphas)  
        self.sqrt_alphas_cumprod           = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1 - self.alphas_cumprod)
        ##
        
        #----------
        # Input Perturbation Reduces Exposure Bias in Diffusion Models
        # https://arxiv.org/pdf/2301.11706.pdf
        
        self.input_perturbation = input_perturbation
        
        #----------
        
        self.to_device(self.device)
    
    @property
    def params_config(self):         
        params_config = {}           
        params_config["device"]              = str(self.device)
        params_config["num_train_timesteps"] = self.num_train_timesteps.item()
        params_config["beta_start"]          = self.beta_start
        params_config["beta_end"]            = self.beta_end
        params_config["beta_schedule"]       = self.beta_schedule
        params_config["input_perturbation"]  = self.input_perturbation            
        return params_config
    
    def to_device(self, device: Union[str, torch.device], non_blocking=False):
        #non_blocking = self.non_blocking
        
        self.device                        = device
        self.alphas_cumprod                = self.alphas_cumprod.to(device, non_blocking=non_blocking)
        self.sqrt_alphas_cumprod           = self.sqrt_alphas_cumprod.to(device, non_blocking=non_blocking)
        self.sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod.to(device, non_blocking=non_blocking)
        self.sigmas                        = self.sigmas.to(device, non_blocking=non_blocking)
        self.sqrt_alphas                   = self.sqrt_alphas.to(device, non_blocking=non_blocking)
        self.betas                         = self.betas.to(device, non_blocking=non_blocking)
        self.num_train_timesteps           = self.num_train_timesteps.to(device, non_blocking=non_blocking)
        self.num_inference_steps           = self.num_inference_steps.to(device, non_blocking=non_blocking)
        
    #------------------------------------
    # Inference functions
    
    def set_timesteps(self, num_inference_steps: int):        
        if num_inference_steps >= self.num_train_timesteps: raise ValueError("num_inference_steps >= self.num_train_timesteps")     
        self.num_inference_steps = torch.tensor(num_inference_steps)
        step_ratio = self.num_train_timesteps // self.num_inference_steps
        timesteps = (np.arange(0, num_inference_steps) * step_ratio.item()).round()[::-1].copy().astype(np.int64)
        self.timesteps = torch.from_numpy(timesteps)
               
    def step(self,
             model_output: torch.FloatTensor,
             timesteps: Union[int, torch.IntTensor],
             sample: torch.FloatTensor
            ) -> DDPMSchedulerOutput:
        """Denoising step"""
        
        sqrt_alphas_cumprod = self.unsqueeze_vector_to_shape(self.sqrt_alphas_cumprod[timesteps], sample.shape)
        sqrt_one_minus_alphas_cumprod = self.unsqueeze_vector_to_shape(self.sqrt_one_minus_alphas_cumprod[timesteps], sample.shape)
           
        sigmas      = self.unsqueeze_vector_to_shape(self.sigmas[timesteps], sample.shape)
        sqrt_alphas = self.unsqueeze_vector_to_shape(self.sqrt_alphas[timesteps], sample.shape)
        betas       = self.unsqueeze_vector_to_shape(self.betas[timesteps], sample.shape)
      
        non_zero_t = (timesteps!=0).float()
    
        #estimate the final img
        x0 = (sample - sqrt_one_minus_alphas_cumprod * model_output) / sqrt_alphas_cumprod   #DDPM eq.15
               
        #less noisy latent     
        noise = torch.randn(sample.shape, device=self.device) 
        noise = noise * non_zero_t.reshape(-1, 1, 1, 1)
        
        xt_coeff = betas / sqrt_one_minus_alphas_cumprod        
        xt = (sample - xt_coeff * model_output) / sqrt_alphas + sigmas * noise 
        
        return DDPMSchedulerOutput(prev_sample=xt, pred_original_sample=x0)

    def add_noise_LEdit(self, original_samples: torch.FloatTensor):
        # LEDITS: Real Image Editing with DDPM Inversion and Semantic Guidance; Note: SEGA (Semantic Guidance) is just multiple negative promts with a pixel based weight
        # https://arxiv.org/pdf/2307.00522.pdf

        noisy_latents = []
        noises        = []

        noisy_latent_t = original_samples
        
        for t in self.timesteps[::-1]: #start from no noise and diffuse in non analytic fashion
            noise_t        = torch.randn_like(noise)   
            alpha_t        = self.unsqueeze_vector_to_shape(self.alphas[t], original_samples.shape)    
            noisy_latent_t = torch.sqrt(alpha_t) * noisy_latent_t + torch.sqrt(1.0 - alpha_t) * noise_t       

            noises.append(noise_t)
            noisy_latents.append(noisy_latent_t)
        
        return noisy_latents[::-1], noises[::-1] # invert to self.timestep definition
    
    #------------------------------------
    # Training functions
        
    def add_noise(self,
                  original_samples: torch.FloatTensor,
                  noise: torch.FloatTensor,
                  timesteps: torch.IntTensor,
                 ) -> torch.FloatTensor:
            
        alphas_cumprod = self.unsqueeze_vector_to_shape(self.alphas_cumprod[timesteps], original_samples.shape)        
        noisy_latents = torch.sqrt(alphas_cumprod) * original_samples + torch.sqrt(1.0 - alphas_cumprod) * noise       
        
        if self.input_perturbation is not None:
            noisy_latents = noisy_latents + torch.sqrt(1.0 - alphas_cumprod) * torch.randn_like(noise) * self.input_perturbation
        
        return noisy_latents  
