"""Base class for schedulers."""

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

# %% auto #0
__all__ = ['Scheduler']

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

# %% ../../src/scheduler/scheduler.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
class Scheduler(abc.ABC):
    """Base class for all diffusion schedulers"""
    
    def __init__(self): 
        pass
    
    @abc.abstractmethod
    def set_timesteps(self): pass
     
    @abc.abstractmethod
    def step(self): pass

    @abc.abstractmethod
    def add_noise(self): pass

    @abc.abstractmethod
    def to(self): pass
    #---------------------------------------

    @staticmethod
    def from_config(config, device: torch.device, save_path: str=None, verbose=True, silent=False):  
        """Use this if we have a loaded config."""

        _config = copy.deepcopy(config)
        
        if exists(device): _config["device"] = device   # for loading sub-models
        else:              device = _config.pop("device", "cpu")

        if "beta_schedule" in _config["params"]:
            beta_schedule = _config["params"]["beta_schedule"]
            
            if "path:" in beta_schedule:
                _config["params"]["beta_schedule"] = "path:" + save_path + beta_schedule[len("path:"):]
        
        scheduler = instantiate_from_config(_config)
        return scheduler

    #---------------------------------------
    
    def get_config(self, without_metadata=False):
        if not without_metadata:       
            config = {}
            config["target"] = class_to_str(type(self))
            config["params"] = self.params_config  
        else:
            config = self.params_config  

        self.config = config        
        return config
    
    @property
    @abc.abstractmethod
    def params_config(self): return None

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

    def unsqueeze_vector_to_shape(self, vec, shape):
        return vec.view(*vec.shape, *([1] * (len(shape)-len(vec.shape))) )
      
    #---------------------------------------

    @classmethod
    def from_scheduler(cls, scheduler, **kwargs): 
        _kwargs = scheduler.params_config 
        _kwargs = _kwargs | kwargs
        
        new_scheduler = cls(**_kwargs)
        return new_scheduler
