from __future__ import annotations

import torch
import torch.nn as nn
from monai.utils import optional_import
from torch.cuda.amp import autocast

tqdm, has_tqdm = optional_import("tqdm", name="tqdm")


class Sampler:
    def __init__(self) -> None:
        super().__init__()

    @torch.no_grad()
    def sampling_fn(
        self,
        noise: torch.Tensor,
        autoencoder_model: nn.Module,
        diffusion_model: nn.Module,
        scheduler: nn.Module,
        prompt_embeds: torch.Tensor,
        guidance_scale: float = 7.0,
        scale_factor: float = 0.3,
    ) -> torch.Tensor:
        if has_tqdm:
            progress_bar = tqdm(scheduler.timesteps)
        else:
            progress_bar = iter(scheduler.timesteps)

        for t in progress_bar:
            noise_input = torch.cat([noise] * 2)
            model_output = diffusion_model(
                noise_input, timesteps=torch.Tensor((t,)).to(noise.device).long(), context=prompt_embeds
            )
            noise_pred_uncond, noise_pred_text = model_output.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            noise, _ = scheduler.step(noise_pred, t, noise)

        with autocast():
            sample = autoencoder_model.decode_stage_2_outputs(noise / scale_factor)

        return sample