import math
from typing import List, Optional, Tuple, Union

import torch
from diffusers import DiffusionPipeline, ImagePipelineOutput, UNet2DModel
from diffusers.utils import randn_tensor


class ConsistencyPipeline(DiffusionPipeline):
    unet: UNet2DModel

    def __init__(
        self,
        unet: UNet2DModel,
    ) -> None:
        super().__init__()
        self.register_modules(unet=unet)

    @torch.no_grad()
    def __call__(
        self,
        steps: int = 1,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        time_min: float = 0.002,
        time_max: float = 80.0,
        data_std: float = 0.5,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        **kwargs,
    ) -> Union[Tuple, ImagePipelineOutput]:
        img_size = self.unet.config.sample_size
        shape = (1, 3, img_size, img_size)

        model = self.unet

        time: float = time_max

        sample = randn_tensor(shape, generator=generator, device=self.device) * time

        for step in self.progress_bar(range(steps)):
            if step > 0:
                time = self.search_previous_time(time)
                sigma = math.sqrt(time**2 - time_min**2 + 1e-6)
                sample = sample + sigma * randn_tensor(sample.shape, device=self.device, generator=generator)

            out = model(sample, torch.tensor([time], device=self.device)).sample

            skip_coef = data_std**2 / ((time - time_min) ** 2 + data_std**2)
            out_coef = data_std * time / (time**2 + data_std**2) ** (0.5)

            sample = (sample * skip_coef + out * out_coef).clamp(-1.0, 1.0)

        sample = (sample / 2 + 0.5).clamp(0, 1)
        image = sample.cpu().permute(0, 2, 3, 1).numpy()

        if output_type == "pil":
            image = self.numpy_to_pil(image)

        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)

    def search_previous_time(self, time, time_min: float = 0.002, time_max: float = 80.0):
        return (2 * time + time_min) / 3

    def cuda(self):
        self.to("cuda")