import torch
from tqdm.notebook import tqdm 
from . import scheduler
from . import share

from lib.utils.iimage import IImage

class DDIM:
    def __init__(self, config, vae, encoder, unet):
        self.vae = vae
        self.encoder = encoder
        self.unet = unet
        self.config = config
        self.schedule = scheduler.linear(1000, config.linear_start, config.linear_end)
    
    def __call__(
            self, prompt = '', dt = 50, shape = (1,4,64,64), seed = None, negative_prompt = '', unet_condition = None, 
            context = None, verbose = True):
        if seed is not None: torch.manual_seed(seed)
        if unet_condition is not None:
            zT = torch.randn((1,4) + unet_condition.shape[2:]).cuda()
        else:
            zT = torch.randn(shape).cuda()

        with torch.autocast('cuda'), torch.no_grad():
            if context is None: context = self.encoder.encode([negative_prompt, prompt])

            zt = zT
            pbar = tqdm(range(999, 0, -dt)) if verbose else range(999, 0, -dt)
            for timestep in share.DDIMIterator(pbar):
                _zt = zt if unet_condition is None else torch.cat([zt, unet_condition], 1)
                eps_uncond, eps = self.unet(
                    torch.cat([_zt, _zt]), 
                    timesteps = torch.tensor([timestep, timestep]).cuda(), 
                    context = context
                ).chunk(2)
                
                eps = (eps_uncond + 7.5 * (eps - eps_uncond))
                
                z0 = (zt - self.schedule.sqrt_one_minus_alphas[timestep] * eps) / self.schedule.sqrt_alphas[timestep]
                zt = self.schedule.sqrt_alphas[timestep - dt] * z0 + self.schedule.sqrt_one_minus_alphas[timestep - dt] * eps
        return IImage(self.vae.decode(z0 / self.config.scale_factor))

    def get_inpainting_condition(self, image, mask):
        latent_size = [x//8 for x in image.size]
        dtype = self.vae.encoder.conv_in.weight.dtype
        with torch.no_grad():
            masked_image = image.torch().cuda() * ~mask.torch(0).bool().cuda()
            masked_image = masked_image.to(dtype)
            condition_x0 = self.vae.encode(masked_image).mean * self.config.scale_factor

        condition_mask = mask.resize(latent_size[::-1]).cuda().torch(0).bool().to(dtype)
        condition_x0 += 0.01 * condition_mask * torch.randn_like(condition_mask)
        return torch.cat([condition_mask, condition_x0], 1)

    inpainting_condition = get_inpainting_condition