"""x0-parameterised diffusion: cosine schedule, forward noising, DDIM sampling with CFG, EMA."""
import copy
import math

import torch


class Schedule:
    def __init__(self, num_train_timesteps=100):
        self.T = num_train_timesteps
        s = 0.008
        f = lambda t: math.cos((t / self.T + s) / (1 + s) * math.pi / 2) ** 2
        betas = [min(1 - f(t + 1) / f(t), 0.999) for t in range(self.T)]
        self.alpha_bar = torch.cumprod(1 - torch.tensor(betas, dtype=torch.float64), 0).float()

    def ab(self, t, device):
        return self.alpha_bar.to(device)[t].view(-1, 1, 1, 1)

    def q_sample(self, x0, t, noise):
        ab = self.ab(t, x0.device)
        return ab.sqrt() * x0 + (1 - ab).sqrt() * noise

    def sample_t(self, n, device):
        return torch.randint(0, self.T, (n,), device=device)

    def timesteps(self, steps):
        """Evenly spaced descending timesteps, e.g. 99, 94, ..., 0 for steps=20."""
        return torch.linspace(self.T - 1, 0, steps).round().long().tolist()

    @torch.no_grad()
    def ddim(self, model_fn, x, steps=20, start=None, callback=None):
        """Deterministic DDIM. model_fn(x_t, t_tensor, i) -> x0_hat. Returns x_0.

        `start` lets sampling begin from an intermediate step index (used for re-noising diagnostics).
        `callback(i, t, x0_hat, x_t)` is called at every step with the prediction and the current noisy sample.
        """
        ts = self.timesteps(steps)
        for i, t in enumerate(ts):
            if start is not None and i < start:
                continue
            tt = torch.full((x.shape[0],), t, device=x.device, dtype=torch.long)
            x0 = model_fn(x, tt, i).clamp(-1, 1)
            if callback is not None:
                callback(i, t, x0, x)
            ab = self.ab(tt, x.device)
            ab_prev = self.ab(torch.full_like(tt, ts[i + 1]), x.device) if i + 1 < len(ts) else torch.ones_like(ab)
            eps = (x - ab.sqrt() * x0) / (1 - ab).sqrt()
            x = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps
        return x


def cfg(cond_out, uncond_out, scale):
    return uncond_out + scale * (cond_out - uncond_out)


class EMA:
    def __init__(self, model, rate=0.999):
        self.rate = rate
        self.model = copy.deepcopy(model).eval().requires_grad_(False)

    @torch.no_grad()
    def update(self, model):
        torch._foreach_lerp_(list(self.model.parameters()), list(model.parameters()), 1 - self.rate)
        for e, b in zip(self.model.buffers(), model.buffers()):
            e.copy_(b)