import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm
from ldm.util import default
from ldm.modules.diffusionmodules.util import  extract_into_tensor
from .ddpm import DDPM



class LatentDiffusion(DDPM):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # hardcoded 
        self.clip_denoised = False
        


    def q_sample(self, x_start, t, noise=None):
        noise = default(noise, lambda: torch.randn_like(x_start))
        return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
                extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)


    "Does not support DDPM sampling anymore. Only do DDIM or PLMS"

    # = = = = = = = = = = = = Below is for sampling = = = = = = = = = = = = # 

    # def predict_start_from_noise(self, x_t, t, noise):
    #     return ( extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
    #              extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise )

    # def q_posterior(self, x_start, x_t, t):
    #     posterior_mean = (
    #             extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
    #             extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
    #     )
    #     posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
    #     posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
    #     return posterior_mean, posterior_variance, posterior_log_variance_clipped


    # def p_mean_variance(self, model, x, c, t):

    #     model_out = model(x, t, c)
    #     x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)

    #     if self.clip_denoised:
    #         x_recon.clamp_(-1., 1.)

    #     model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
    #     return model_mean, posterior_variance, posterior_log_variance, x_recon


    # @torch.no_grad()
    # def p_sample(self, model, x, c, t):
    #     b, *_, device = *x.shape, x.device
    #     model_mean, _, model_log_variance, x0 = self.p_mean_variance(model, x=x, c=c, t=t, )
    #     noise = torch.randn_like(x) 

    #     # no noise when t == 0
    #     nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))

    #     return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0


    # @torch.no_grad()
    # def p_sample_loop(self, model, shape, c):
    #     device = self.betas.device
    #     b = shape[0]
    #     img = torch.randn(shape, device=device)

    #     iterator = tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps) 
    #     for i in iterator:
    #         ts = torch.full((b,), i, device=device, dtype=torch.long)
    #         img, x0 = self.p_sample(model, img, c, ts)

    #     return img


    # @torch.no_grad()
    # def sample(self, model, shape, c, uc=None, guidance_scale=None):
    #     return self.p_sample_loop(model, shape, c)