import torch
import math
import numpy as np
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, UniPCMultistepScheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from typing import Union, Optional, List, Callable, Dict, Any, Tuple
from momentum_scheduler import (
    GHVBScheduler,
    PLMSWithHBScheduler,
    PLMSWithNTScheduler,
    MomentumDPMSolverMultistepScheduler,
    MomentumUniPCMultistepScheduler,
)

available_solvers = {
    "GHVB": GHVBScheduler,
    "PLMS_HB": PLMSWithHBScheduler,
    "PLMS_NT": PLMSWithNTScheduler,
    "DPM-Solver++": MomentumDPMSolverMultistepScheduler,
    "UniPC": MomentumUniPCMultistepScheduler,
}

def get_momentum_number(order, beta):
    out = order if beta == 1.0 else order - (1 - beta)
    return out

def setup_scheduler(pipe, scheduler, momentum_type="Polyak's heavy ball", order=4.0, beta=1.0, original_config=None):
    assert original_config is not None
    
    if scheduler in ["DPM-Solver++", "UniPC"]:
        if momentum_type in ["Nesterov"]:
            raise NotImplementedError(f"{scheduler} w/ Nesterov is not implemented.")

        pipe.scheduler = available_solvers[scheduler].from_config(original_config)
        pipe.scheduler.initialize_momentum(beta=beta)

    elif scheduler in ["PLMS"]:
        momentum_number = get_momentum_number(order, beta)
        method = "PLMS_HB" if momentum_type == "Polyak's heavy ball" else "PLMS_NT"
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(original_config)
        pipe.init_scheduler(method=method, order=momentum_number)
        pipe.clear_scheduler()

    elif scheduler in ["GHVB"]:
        momentum_number = get_momentum_number(order, beta)
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(original_config)
        pipe.init_scheduler(method="GHVB", order=momentum_number)
        pipe.clear_scheduler()

    return pipe

class CustomPipeline(StableDiffusionPipeline):
    def clear_scheduler(self):
        self.scheduler_uncond.clear()
        self.scheduler_text.clear()
    
    def init_scheduler(self, method, order):
        # equivalent to not applied numerical operator splitting since orders are the same
        self.scheduler_uncond = available_solvers[method](self.scheduler, order)
        self.scheduler_text = available_solvers[method](self.scheduler, order)

    def get_noise(self, latents, prompt_embeds, guidance_scale, t, do_classifier_free_guidance):
        # expand the latents if we are doing classifier free guidance
        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
        latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

        # predict the noise residual
        noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample

        if do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            grads_a = guidance_scale * (noise_pred_text - noise_pred_uncond)

        return noise_pred_uncond, grads_a

    def denoising_step(
            self,
            latents,
            prompt_embeds,
            guidance_scale,
            t,
            do_classifier_free_guidance,
            method,
            extra_step_kwargs,
    ):
        noise_pred_uncond, grads_a = self.get_noise(
            latents, prompt_embeds, guidance_scale, t, do_classifier_free_guidance
        )
        if method in ["dpm", "unipc"]:
            latents = self.scheduler.step(noise_pred_uncond + grads_a, t, latents, **extra_step_kwargs).prev_sample
        
        elif method in ["hb", "ghvb", "nt"]:
            latents = self.scheduler_uncond.step(noise_pred_uncond, t, latents, output_mode="scale")
            latents = self.scheduler_text.step(grads_a, t, latents, output_mode='back')
        else:
            raise NotImplementedError

        return latents

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        method="ghvb",
    ):
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps
        # print(timesteps)

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                latents = self.denoising_step(
                    latents,
                    prompt_embeds,
                    guidance_scale,
                    t,
                    do_classifier_free_guidance,
                    method,
                    extra_step_kwargs,
                )

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
            has_nsfw_concept = None
        elif output_type == "pil":
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            # image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
            has_nsfw_concept = False

            # 10. Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            # image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
            has_nsfw_concept = False

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

    def generate(self, params):
        params["output_type"] = "latent"
        ori_latents = self.__call__(**params)["images"]

        with torch.no_grad():
            latents = torch.clone(ori_latents)
            image = self.decode_latents(latents)
            image = self.numpy_to_pil(image)[0]

        return image, ori_latents