"""Denoising diffusion probabilistic models [(DDPM)](https://arxiv.org/abs/2006.11239): reverse beta is fixed and diagonal."""

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

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

# %% ../../src/scheduler/scheduler_ddpm.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from .scheduler import Scheduler
from ..utils.config_loader import load_tensor

# %% ../../src/scheduler/scheduler_ddpm.ipynb #ce85e7b8-1c87-48fe-b75d-bd0dbc107494
@dataclass
class DDPMSchedulerOutput:
    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None

# %% ../../src/scheduler/scheduler_ddpm.ipynb #de0475b1-e0aa-42eb-9f56-12c131ef868f
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,
                 prediction_type = "epsilon",
                 enable_zero_terminal_snr = True
                ):    
        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

        self.input_perturbation = input_perturbation # Input Perturbation Reduces Exposure Bias in Diffusion Models, https://arxiv.org/pdf/2301.11706.pdf        
        self.prediction_type    = prediction_type   # one of "epsilon", "v-type", "x0", "mu"
        
        if self.prediction_type not in ["epsilon", "v-type", "x0"]: 
            raise NotImplementedError(f"{self.prediction_type} does is not implemented for {self.__class__}")

        #-----------
        
        if beta_schedule == "linear":
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
            
        elif beta_schedule == "linear_sqrt": #LDM
            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          
            #print("[INFO]: using cos_alpha beta-schedule, ignoring beta_start and beta_end!")
            f = lambda t: np.cos((t/self.num_train_timesteps + 0.008)*np.pi/2.016)**2    # is alpha_bar
            _betas = []            
            for i in range(self.num_train_timesteps):            
                b = 1.0-(f(i+1.0)/f(i))
                if not enable_zero_terminal_snr: # v-type allows zero terminal SNR
                    b = min(b, 0.999)                # clipping disables zero terminal SNR
                _betas.append(b)                
            self.betas = torch.tensor(_betas, dtype=torch.float32)  

        elif beta_schedule == "cos_alpha4":  #cosine-based-variance          
            #print("[INFO]: using cos_alpha4 beta-schedule, ignoring beta_start and beta_end!")
            f = lambda t: np.cos((t/self.num_train_timesteps + 0.008)*np.pi/2.016)**4    # is alpha_bar
            _betas = []            
            for i in range(self.num_train_timesteps):            
                b = 1.0-(f(i+1.0)/f(i))
                if not enable_zero_terminal_snr: # v-type allows zero terminal SNR
                    b = min(b, 0.999)                # clipping disables zero terminal SNR
                _betas.append(b)                
            self.betas = torch.tensor(_betas, dtype=torch.float32)  
        
        elif "path:" in beta_schedule:
            _save_path = beta_schedule[len("path:"):]
            self.betas = load_tensor(save_path=_save_path, device=device)["0"]
            
            print(f"[INFO]: Loaded beta_schedule ({beta_schedule}).")
        
        else:
            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
    
        #-----------
        
        if (self.prediction_type in ["v-type", "x0"]) and enable_zero_terminal_snr and (beta_schedule not in ["cos_alpha", "laplace"]): # v-type allows zero terminal SNR
            self.betas = self.enforce_zero_terminal_snr(self.betas)

        #-----------

        self.sigmas = torch.sqrt(self.betas)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)   #only do cumprod witj alphas, as betas will go within precision to zero

        #----------
        
        self.to(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           
        params_config["prediction_type"]     = self.prediction_type    
        return params_config
    
    def to(self, device: Union[str, torch.device], non_blocking=False):
        self.device                        = device
        self.alphas_cumprod                = self.alphas_cumprod.to(device, non_blocking=non_blocking)
        self.sigmas                        = self.sigmas.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)
        return self
    
    #------------------------------------
    
    @property
    def SNR(self):
        alphas_bar = self.alphas_cumprod
        betas_bar  = 1.0 - alphas_bar
        return alphas_bar / betas_bar
    
    #------------------------------------
    # Inference functions

    def enforce_zero_terminal_snr(self, betas):
        # Algorithm 1 in https://arxiv.org/pdf/2305.08891.pdf
        
        # Convert betas to alphas_bar_sqrt
        alphas = 1 - betas
        alphas_bar = alphas.cumprod(0)
        alphas_bar_sqrt = alphas_bar.sqrt()
        
        # Store old values.
        alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
        alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
        # Shift so last timestep is zero.
        alphas_bar_sqrt -= alphas_bar_sqrt_T
        # Scale so first timestep is back to old value.
        alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
        
        # Convert alphas_bar_sqrt to betas
        alphas_bar = alphas_bar_sqrt ** 2
        alphas = alphas_bar[1:] / alphas_bar[:-1]
        alphas = torch.cat([alphas_bar[0:1], alphas])
        betas = 1 - alphas
        return betas
    
    def set_timesteps(self, num_inference_steps: Optional[int] = None, timesteps: Optional[torch.Tensor] = None):     
        if exists(num_inference_steps):
            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)

        elif exists(timesteps):
            self.num_inference_steps = torch.tensor(timesteps.shape[0])
            self.timesteps           = timesteps.clone()

        else:
            raise RuntimeError("provide `num_inference_steps` or `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)
        
        if self.prediction_type == "epsilon":
            #estimate the final img
            x0 = (sample - sqrt_one_minus_alphas_cumprod * model_output) / sqrt_alphas_cumprod   #DDPM eq.15

            xt_coeff = betas / sqrt_one_minus_alphas_cumprod  
            mu_t     = (sample - xt_coeff * model_output) / sqrt_alphas
        
        elif self.prediction_type == "v-type":
            x0   = sqrt_alphas_cumprod * sample - sqrt_one_minus_alphas_cumprod * model_output 

            prev_timesteps     = timesteps - self.num_train_timesteps // self.num_inference_steps
            alphas_cumprod_tm1 = self.unsqueeze_vector_to_shape(self.alphas_cumprod[prev_timesteps], sample.shape)
            
            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.alphas_cumprod[0]
           
            mu_t = (betas * alphas_cumprod_tm1.sqrt() * x0 + sqrt_alphas * (1.0-alphas_cumprod_tm1) * sample) / sqrt_one_minus_alphas_cumprod
                      
        else:
            raise NotImplementedError(f"{self.prediction_type} is not implemented for {self.__class__}.step()")
                
        #less noisy latent     
        non_zero_t = (timesteps>0).float()
        noise = torch.randn(sample.shape, device=self.device) 
        noise = noise * non_zero_t.reshape(-1, 1, 1, 1)
           
        xt = mu_t + sigmas * noise 
        
        return DDPMSchedulerOutput(prev_sample=xt, pred_original_sample=x0)

    #------------------------------------
    # Training functions
        
    def add_noise(self,
                  original_samples: torch.FloatTensor,
                  noise: torch.FloatTensor,
                  timesteps: torch.IntTensor,
                  train: bool=False
                 ) -> 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       # F^2
        
        if exists(self.input_perturbation) and train:
            noisy_latents = noisy_latents + torch.sqrt(1.0 - alphas_cumprod) * torch.randn_like(noise) * self.input_perturbation
        
        return noisy_latents
