# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/08_timesteps.ipynb.

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

# %% ../../nbs/utils/08_timesteps.ipynb 2
from typing import Literal

import torch
from torch import nn

from vision_architectures.layers.embeddings import (  # noqa: F401
    get_all_timestep_embeddings_1d,
    get_timestep_embeddings_1d,
)

# %% ../../nbs/utils/08_timesteps.ipynb 5
class TimestepSampler(nn.Module):
    """Sampel timesteps using a strategy.

    Args:
        total_timesteps: Total number of timesteps.
        strategy: Sampling strategy.
        gamma_alpha: Used when strategy == gamma. Alpha parameter for the gamma distribution.
        gamma_spread: Used when strategy == gamma. Controls how spread out the distribution is.
        signal_to_noise_ratios: Used when strategy == importance. Signal-to-noise ratios of the noise schedule.
        temperature: Used when strategy == importance. Temperature for the softmax distribution.
    """

    def __init__(
        self,
        total_timesteps: int,
        strategy: Literal["uniform", "gamma", "importance"],
        gamma_alpha: float = 1.0,
        signal_to_noise_ratios: torch.Tensor = None,
        temperature: float = 2.0,
    ):
        super().__init__()

        self.T = total_timesteps
        self.strategy = strategy

        if strategy == "uniform":
            # Sample uniformly from [1, T]
            self.distribution = torch.distributions.Uniform(1, self.T + 1)
        elif strategy == "gamma":
            # Sample from a gamma distribution. There is no point in parameterizing beta as the sample is normalized
            self.distribution = torch.distributions.Gamma(gamma_alpha, 1)
        elif strategy == "importance":
            raise NotImplementedError("This implementation still needs to be completed.")
            # Sample from [0, T] where prob(t) is d(log(SNR))/dt and prob(t=0) = 0
            if signal_to_noise_ratios is None:
                raise ValueError("signal_to_noise_ratios must be provided for importance sampling.")

            self.snr = signal_to_noise_ratios
            self.log_snr = torch.log(signal_to_noise_ratios)
            self.log_snr_gradient = torch.abs(self.log_snr[1:] - self.log_snr[:-1])
            self.log_snr_gradient = torch.cat(
                [
                    torch.zeros(
                        1,
                    ),
                    self.log_snr_gradient[0:1],
                    self.log_snr_gradient,
                ],
                dim=0,
            )
            self.weights = self.log_snr_gradient ** (1 / temperature)
            self.distribution = torch.distributions.Categorical(probs=self.weights)
        else:
            raise NotImplementedError(f"Sampling strategy {strategy} not implemented.")

    def forward(self, num_timesteps: int):
        """Sample timesteps based on the specified strategy.

        Args:
            num_timesteps: Number of timesteps to sample.
        """
        timesteps = self.distribution.sample((num_timesteps,))
        if self.strategy == "gamma":
            timesteps = timesteps / timesteps.max()
            timesteps = timesteps * self.T
        timesteps = timesteps.clamp(1, self.T).long()
        return timesteps
