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

# %% auto 0
__all__ = ['DDIMSchedulerOutput', 'DDIMScheduler']

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

# %% ../../src/scheduler/scheduler_ddim.ipynb 4
@dataclass
class DDIMSchedulerOutput:
    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None
    
class DDIMScheduler(DDPMScheduler):
    """A `Scheduler` implementing [(DDIM)](https://arxiv.org/abs/2010.02502)."""
    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,
                 eta: float = 0
                ):    
        super().__init__(device, num_train_timesteps, beta_start, beta_end, beta_schedule, input_perturbation)
        self.eta = eta
        
        #for stable diff ddim
        set_alpha_to_one = True  # False 
        steps_offset     = 0     # 1
            
        self.steps_offset        = steps_offset
        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
     
    @property
    def params_config(self):         
        params_config = super().params_config
        params_config["eta"] = self.eta
        return params_config
    
    #------------------------------------
    # Inference functions
    
    def set_timesteps(self, num_inference_steps: int):        
        super().set_timesteps(num_inference_steps)   
        self.timesteps += self.steps_offset
    
    clamp_style = None # one of: None, "static", "dynamic"
    
    def step(self,         
             model_output: torch.FloatTensor,
             timesteps: Union[int, torch.IntTensor],
             sample: torch.FloatTensor
            ) -> DDIMSchedulerOutput:
        """Denoising step"""
        
        prev_timesteps = timesteps - self.num_train_timesteps // self.num_inference_steps
        
        #get variance sched
        alphas_cumprod     = self.unsqueeze_vector_to_shape(self.alphas_cumprod[timesteps], sample.shape)
        alphas_cumprod_tm1 = self.unsqueeze_vector_to_shape(self.alphas_cumprod[prev_timesteps], sample.shape)
                       
        #fix negative timesteps to self.final_alpha_cumprod        
        non_zero_tm1 = (prev_timesteps>=0.0).float()      
        non_zero_tm1 = self.unsqueeze_vector_to_shape(non_zero_tm1, sample.shape)
        alphas_cumprod_tm1 = alphas_cumprod_tm1 * non_zero_tm1 + (1.0 - non_zero_tm1) *  self.final_alpha_cumprod  
                              
        #calc vars
        betas_cumprod     = 1.0 - alphas_cumprod
        betas_cumprod_tm1 = 1.0 - alphas_cumprod_tm1
            
        #estimate predicted sample
        x0 = (sample - betas_cumprod.sqrt() * model_output) / alphas_cumprod.sqrt()
        
        if   self.clamp_style == None: pass
        elif self.clamp_style == "static":  x0 = torch.clamp(x0, -1, 1)
        elif self.clamp_style == "dynamic": raise NotImplementedError("clamp_style == 'dynamic'")
        else: raise NotImplementedError("self.clamp_style has to be one of: None, 'static', 'dynamic'")       
        #if self.num_train_timesteps // self.num_inference_steps > 1: x0 = torch.clamp(x0, -1, 1)
    
        #variance
        variance = (betas_cumprod_tm1 / betas_cumprod) * (1.0 - alphas_cumprod / alphas_cumprod_tm1)
        std      = self.eta * variance.sqrt()
        
        #direction to xt
        dir_xt = ( betas_cumprod_tm1 - std.square() ).sqrt() * model_output
    
        #sample noise
        noise = torch.randn(model_output.shape, device=self.device)
    
        #estimate the prev sample      
        xtm1 = alphas_cumprod_tm1.sqrt() * x0 + dir_xt + std * noise
        
        # print(f"{timesteps=} {prev_timesteps=} ;;;   x0:  {x0.mean()}+-{x0.std()}      xtm1:  {xtm1.mean()}+-{xtm1.std()}")
        
        return DDIMSchedulerOutput(prev_sample=xtm1, pred_original_sample=x0)    
