from abc import abstractmethod
from functools import partial

import numpy as np
import torch

from ...modules.diffusionmodules.util import make_beta_schedule
from ...util import append_zero


def generate_roughly_equally_spaced_steps(num_substeps: int,
                                          max_step: int) -> np.ndarray:
    return np.linspace(max_step - 1, 0, num_substeps,
                       endpoint=False).astype(int)[::-1]


class Discretization:

    def __call__(self,
                 n,
                 do_append_zero=True,
                 device='cpu',
                 flip=False,
                 return_idx=False):
        if return_idx:
            sigmas, idx = self.get_sigmas(n,
                                          device=device,
                                          return_idx=return_idx)
        else:
            sigmas = self.get_sigmas(n, device=device, return_idx=return_idx)
        sigmas = append_zero(sigmas) if do_append_zero else sigmas
        if return_idx:
            return sigmas if not flip else torch.flip(sigmas, (0, )), idx
        else:
            return sigmas if not flip else torch.flip(sigmas, (0, ))

    @abstractmethod
    def get_sigmas(self, n, device):
        pass


class EDMDiscretization(Discretization):

    def __init__(self, sigma_min=0.002, sigma_max=80.0, rho=7.0):
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.rho = rho

    def get_sigmas(self, n, device='cpu'):
        ramp = torch.linspace(0, 1, n, device=device)
        min_inv_rho = self.sigma_min**(1 / self.rho)
        max_inv_rho = self.sigma_max**(1 / self.rho)
        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**self.rho
        return sigmas


class LegacyDDPMDiscretization(Discretization):

    def __init__(
        self,
        linear_start=0.00085,
        linear_end=0.0120,
        num_timesteps=1000,
    ):
        super().__init__()
        self.num_timesteps = num_timesteps
        betas = make_beta_schedule('linear',
                                   num_timesteps,
                                   linear_start=linear_start,
                                   linear_end=linear_end)
        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.to_torch = partial(torch.tensor, dtype=torch.float32)

    def get_sigmas(self, n, device='cpu'):
        if n < self.num_timesteps:
            timesteps = generate_roughly_equally_spaced_steps(
                n, self.num_timesteps)
            alphas_cumprod = self.alphas_cumprod[timesteps]
        elif n == self.num_timesteps:
            alphas_cumprod = self.alphas_cumprod
        else:
            raise ValueError

        to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
        sigmas = to_torch((1 - alphas_cumprod) / alphas_cumprod)**0.5
        return torch.flip(sigmas, (0, ))  # sigma_t: 14.4 -> 0.029


class ZeroSNRDDPMDiscretization(Discretization):

    def __init__(
        self,
        linear_start=0.00085,
        linear_end=0.0120,
        num_timesteps=1000,
        shift_scale=1.0,  # noise schedule t_n -> t_m: logSNR(t_m) = logSNR(t_n) - log(shift_scale)
        keep_start=False,
        post_shift=False,
    ):
        super().__init__()
        if keep_start and not post_shift:
            linear_start = linear_start / (shift_scale +
                                           (1 - shift_scale) * linear_start)
        self.num_timesteps = num_timesteps
        betas = make_beta_schedule('linear',
                                   num_timesteps,
                                   linear_start=linear_start,
                                   linear_end=linear_end)
        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.to_torch = partial(torch.tensor, dtype=torch.float32)

        # SNR shift
        if not post_shift:
            self.alphas_cumprod = self.alphas_cumprod / (
                shift_scale + (1 - shift_scale) * self.alphas_cumprod)

        self.post_shift = post_shift
        self.shift_scale = shift_scale

    def get_sigmas(self, n, device='cpu', return_idx=False):
        if n < self.num_timesteps:
            timesteps = generate_roughly_equally_spaced_steps(
                n, self.num_timesteps)
            alphas_cumprod = self.alphas_cumprod[timesteps]
        elif n == self.num_timesteps:
            alphas_cumprod = self.alphas_cumprod
        else:
            raise ValueError

        to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
        alphas_cumprod = to_torch(alphas_cumprod)
        alphas_cumprod_sqrt = alphas_cumprod.sqrt()
        alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone()
        alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone()

        alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T
        alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 -
                                                        alphas_cumprod_sqrt_T)

        if self.post_shift:
            alphas_cumprod_sqrt = (
                alphas_cumprod_sqrt**2 /
                (self.shift_scale +
                 (1 - self.shift_scale) * alphas_cumprod_sqrt**2))**0.5

        if return_idx:
            return torch.flip(alphas_cumprod_sqrt, (0, )), timesteps
        else:
            return torch.flip(alphas_cumprod_sqrt,
                              (0, ))  # sqrt(alpha_t): 0 -> 0.99
