"""DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models [(DPM-Solver)](https://arxiv.org/abs/2206.00927) [(DPM-Solver++)](https://arxiv.org/abs/2211.01095)."""

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

# %% auto #0
__all__ = ['DPMSchedulerOutput', 'DPMScheduler']

# %% ../../src/scheduler/scheduler_dpm.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from .scheduler_ddpm import DDPMScheduler

# %% ../../src/scheduler/scheduler_dpm.ipynb #c6d3b1e2-18d3-4cab-968c-853cb3469b3a
@dataclass
class DPMSchedulerOutput:
    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None

# %% ../../src/scheduler/scheduler_dpm.ipynb #de0475b1-e0aa-42eb-9f56-12c131ef868f
class DPMScheduler(DDPMScheduler):
    """A `Scheduler` implementing [(DPM-Solver++)](https://arxiv.org/abs/2211.01095)."""
    
    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,
                 solver_order: int = 2,
                 **kwargs
                ) -> None:    
        super().__init__(device, num_train_timesteps, beta_start, beta_end, beta_schedule, input_perturbation, prediction_type, enable_zero_terminal_snr)

        self.solver_order = solver_order
        if self.solver_order != 2:
            raise NotImplementedError(f"{self.solver_order=} is not implemented for {self.__class__}")
    
    @property
    def params_config(self):         
        params_config = super().params_config
        params_config["solver_order"] = self.solver_order
        return params_config
    
    #------------------------------------
    # Inference functions
    
    def step(self,         
             model_output: torch.FloatTensor,
             timesteps: torch.IntTensor,
             sample: torch.FloatTensor,
             uncond_model_output: torch.FloatTensor = None, # for CFG++
            ) -> DPMSchedulerOutput:
        """
        Denoising step of DPM-Solver++(2M) (Lu et al., 2022b), 
        implemeted as CFG++ variant (CFG++, https://arxiv.org/pdf/2406.08070)
        """

        uncond_model_output = default(uncond_model_output, model_output)
        
        assert timesteps.numel() == 1

        # note: here we enforce the sampling to be strictly defined by self.timesteps
        is_warmup_step = (self.timesteps[0]  == timesteps)
        # is_last_step   = (self.timesteps[-1] == timesteps)
        
        alphas_cumprod = self.unsqueeze_vector_to_shape(self.alphas_cumprod[timesteps], sample.shape)
        betas_cumprod  = 1.0 - alphas_cumprod

        prev_timesteps = timesteps - self.num_train_timesteps // self.num_inference_steps
        prev_timesteps = prev_timesteps.clamp(0, self.num_train_timesteps-1)

        alphas_cumprod_tm1 = self.unsqueeze_vector_to_shape(self.alphas_cumprod[prev_timesteps], sample.shape)
        betas_cumprod_tm1  = 1.0 - alphas_cumprod_tm1

        # ---------
        if self.prediction_type == "v-type":
            a = alphas_cumprod.sqrt()
            b = betas_cumprod.sqrt()
            
            x0        = a * sample - b * model_output
            x0_uncond = a * sample - b * uncond_model_output 

        elif self.prediction_type == "x0":
            x0 = model_output
            x0_uncond = uncond_model_output
        
        else:
            raise NotImplementedError(f"{self.prediction_type} is not implemented for {self.__class__}.step()")

        # ---------
        solver_order = self.solver_order        
        # mod here for adyptive adjust, if needed
        if solver_order == 2:
            pass
            
        else:
            raise NotImplementedError(f"{solver_order} is not implemented for {self.__class__}")

        # ---------

        lambda_t   = 0.5 * torch.log(alphas_cumprod     / betas_cumprod)
        lambda_tm1 = 0.5 * torch.log(alphas_cumprod_tm1 / betas_cumprod_tm1)

        h_tm1 = lambda_tm1 - lambda_t
         
        if is_warmup_step:
            x_dir = alphas_cumprod_tm1.sqrt() * (x0 - torch.exp(-h_tm1) * x0_uncond)
        
        else:
            r_tm1 = self.last_h_tm1 / h_tm1
            
            sqrt_alphas_cumprod_tm1 = alphas_cumprod_tm1.sqrt()
            exp_mhtm1 = torch.exp(-h_tm1)
            
            x_dir = sqrt_alphas_cumprod_tm1 * x0 - sqrt_alphas_cumprod_tm1 * exp_mhtm1 * x0_uncond + sqrt_alphas_cumprod_tm1 * (0.5/r_tm1) * (x0_uncond-self.last_x0_uncond) * (1.0-exp_mhtm1)
            
        xtm1 = (betas_cumprod_tm1/betas_cumprod).sqrt() * sample + x_dir

        # is needed for multistesp integration of DPM                
        self.last_x0_uncond = x0_uncond 
        self.last_h_tm1     = h_tm1
                       
        return DPMSchedulerOutput(prev_sample=xtm1, pred_original_sample=x0)    
