from typing import Any

import torch

from .scheduler_utils import SNR_to_betas, compute_snr


class ShiftSNRScheduler:
    def __init__(
        self,
        noise_scheduler: Any,
        timesteps: Any,
        shift_scale: float,
        scheduler_class: Any,
    ):
        self.noise_scheduler = noise_scheduler
        self.timesteps = timesteps
        self.shift_scale = shift_scale
        self.scheduler_class = scheduler_class

    def _get_shift_scheduler(self):
        """
        Prepare scheduler for shifted betas.

        :return: A scheduler object configured with shifted betas
        """
        snr = compute_snr(self.timesteps, self.noise_scheduler)
        shifted_betas = SNR_to_betas(snr / self.shift_scale)

        return self.scheduler_class.from_config(
            self.noise_scheduler.config, trained_betas=shifted_betas.numpy()
        )

    def _get_interpolated_shift_scheduler(self):
        """
        Prepare scheduler for shifted betas and interpolate with the original betas in log space.

        :return: A scheduler object configured with interpolated shifted betas
        """
        snr = compute_snr(self.timesteps, self.noise_scheduler)
        shifted_snr = snr / self.shift_scale

        weighting = self.timesteps.float() / (
            self.noise_scheduler.config.num_train_timesteps - 1
        )
        interpolated_snr = torch.exp(
            torch.log(snr) * (1 - weighting) + torch.log(shifted_snr) * weighting
        )

        shifted_betas = SNR_to_betas(interpolated_snr)

        return self.scheduler_class.from_config(
            self.noise_scheduler.config, trained_betas=shifted_betas.numpy()
        )

    @classmethod
    def from_scheduler(
        cls,
        noise_scheduler: Any,
        shift_mode: str = "default",
        timesteps: Any = None,
        shift_scale: float = 1.0,
        scheduler_class: Any = None,
    ):
        # Check input
        if timesteps is None:
            timesteps = torch.arange(0, noise_scheduler.config.num_train_timesteps)
        if scheduler_class is None:
            scheduler_class = noise_scheduler.__class__

        # Create scheduler
        shift_scheduler = cls(
            noise_scheduler=noise_scheduler,
            timesteps=timesteps,
            shift_scale=shift_scale,
            scheduler_class=scheduler_class,
        )

        if shift_mode == "default":
            return shift_scheduler._get_shift_scheduler()
        elif shift_mode == "interpolated":
            return shift_scheduler._get_interpolated_shift_scheduler()
        else:
            raise ValueError(f"Unknown shift_mode: {shift_mode}")


if __name__ == "__main__":
    """
    Compare the alpha values for different noise schedulers.
    """
    import matplotlib.pyplot as plt
    from diffusers import DDPMScheduler

    from .scheduler_utils import compute_alpha

    # Base
    timesteps = torch.arange(0, 1000)
    noise_scheduler_base = DDPMScheduler.from_pretrained(
        "runwayml/stable-diffusion-v1-5", subfolder="scheduler"
    )
    alpha = compute_alpha(timesteps, noise_scheduler_base)
    plt.plot(timesteps.numpy(), alpha.numpy(), label="Base")

    # Kolors
    num_train_timesteps_ = 1100
    timesteps_ = torch.arange(0, num_train_timesteps_)
    noise_kwargs = {"beta_end": 0.014, "num_train_timesteps": num_train_timesteps_}
    noise_scheduler_kolors = DDPMScheduler.from_config(
        noise_scheduler_base.config, **noise_kwargs
    )
    alpha = compute_alpha(timesteps_, noise_scheduler_kolors)
    plt.plot(timesteps_.numpy(), alpha.numpy(), label="Kolors")

    # Shift betas
    shift_scale = 8.0
    noise_scheduler_shift = ShiftSNRScheduler.from_scheduler(
        noise_scheduler_base, shift_mode="default", shift_scale=shift_scale
    )
    alpha = compute_alpha(timesteps, noise_scheduler_shift)
    plt.plot(timesteps.numpy(), alpha.numpy(), label="Shift Noise (scale 8.0)")

    # Shift betas (interpolated)
    noise_scheduler_inter = ShiftSNRScheduler.from_scheduler(
        noise_scheduler_base, shift_mode="interpolated", shift_scale=shift_scale
    )
    alpha = compute_alpha(timesteps, noise_scheduler_inter)
    plt.plot(timesteps.numpy(), alpha.numpy(), label="Interpolated (scale 8.0)")

    # ZeroSNR
    noise_scheduler = DDPMScheduler.from_config(
        noise_scheduler_base.config, rescale_betas_zero_snr=True
    )
    alpha = compute_alpha(timesteps, noise_scheduler)
    plt.plot(timesteps.numpy(), alpha.numpy(), label="ZeroSNR")

    plt.legend()
    plt.grid()
    plt.savefig("check_alpha.png")