import itertools
import math
import torch
import numpy as np
import pytorch_lightning as L
import torchmetrics
from dataclasses import dataclass
import dit, ema
import noise_schedule  # Assuming this is part of the MDLM repository

LOG2 = math.log(2)

@dataclass
class Loss:
    loss: torch.FloatTensor
    nlls: torch.FloatTensor
    token_mask: torch.FloatTensor

class NLL(torchmetrics.MeanMetric):
    pass

class BPD(NLL):
    def compute(self) -> torch.Tensor:
        """Computes the bits per dimension.
        Returns:
          bpd
        """
        return self.mean_value / self.weight / LOG2

class Perplexity(NLL):
    def compute(self) -> torch.Tensor:
        """Computes the Perplexity.
        Returns:
         Perplexity
        """
        return torch.exp(self.mean_value / self.weight)

# Based on MDLM repo
class Diffusion(L.LightningModule):
    def __init__(self, config, latent_dim, tokenizer):
        super().__init__()
        self.config = config
        self.latent_dim = latent_dim
        self.tokenizer = tokenizer

        self.backbone = dit.DIT(self.config, vocab_size=self.latent_dim)
        self.T = self.config.T
        self.subs_masking = self.config.SUBS_MASKING
        self.antithetic_sampling = self.config.Training.ANTITHETIC_SAMPLING
        self.mask_index = self.tokenizer.mask_token_id

        self.softplus = torch.nn.Softplus()
        metrics = torchmetrics.MetricCollection({
            'nll': NLL(),
            'bpd': BPD(),
            'ppl': Perplexity(),
        })
        metrics.set_dtype(torch.float64)
        self.train_metrics = metrics.clone(prefix='train/')
        self.valid_metrics = metrics.clone(prefix='val/')
        self.test_metrics = metrics.clone(prefix='test/')

        self.noise = noise_schedule.get_noise(self.config, dtype=self.dtype)
        self.lr = self.config.Optim.LR
        self.sampling_eps = self.config.Training.SAMPLING_EPS
        self.time_conditioning = self.config.TIME_CONDITIONING
        self.neg_infinity = -1000000.0


    ############ FORWARD DIFFUSION #########
    def subs_parameterization(self, logits, noised_latents):
        # log prob at the mask index = - infinity
        logits[:, :, self.mask_index] += self.neg_infinity
        
        # Normalize the logits such that x.exp() is
        # a probability distribution over vocab_size.
        logits = logits - torch.logsumexp(logits, dim=-1,
                                        keepdim=True)

        # Apply updates directly in the logits matrix.
        # For the logits of the unmasked tokens, set all values
        # to -infinity except for the indices corresponding to
        # the unmasked tokens.
        unmasked_indices = (noised_latents != self.mask_index)
        logits[unmasked_indices] = self.neg_infinity
        logits[unmasked_indices, noised_latents[unmasked_indices]] = 0
        return logits

    def forward(self, latents, sigma):
        latents = latents.long()
        with torch.cuda.amp.autocast(dtype=torch.float32):
            logits = self.backbone(latents, sigma)
        print(logits)
        optimized_logits = self.subs_parameterization(logits, latents)
        return optimized_logits
    
    def q_xt(self, latents, move_chance):
        """
        Computes the noisy sample xt.
        Args:
            x: int torch.Tensor with shape (batch_size, diffusion_model_input_length), input. 
            move_chance: float torch.Tensor with shape (batch_size, 1).
        """
        latents = latents.mean(dim=1) # [bsz x seq_len x 1280] --> [bsz x 1280] as per args
        move_indices = torch.rand(* latents.shape, device=latents.device) < move_chance
        noised_latents = torch.where(move_indices, self.mask_index, latents)
        return noised_latents

    def sample_timestep(self, n, device):
        _eps_t = torch.rand(n, device=device)
        if self.antithetic_sampling:
            offset = torch.arange(n, device=device) / n
            _eps_t = (_eps_t / n + offset) % 1
        t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps
        # if self.importance_sampling:
        #     return self.noise.importance_sampling_transformation(t)
        return t


    def d3pm_loss(self, model_output, xt, x0, t):
        """Computes the D3PM loss between noisy latents and the original input at a given time step."""
        dt = 1 / self.T

        if torch.is_tensor(t):
            t = t[:, None]
            assert t.ndim == 2
            t = t.clamp(0., 1. - 1e-4)
        alpha_t = 1 - t + torch.zeros_like(xt)
        alpha_s = 1 - (t - dt) + torch.zeros_like(xt)

        x0 = x0.to(torch.int64)
        log_x_theta_at_x0 = torch.gather(model_output, -1, x0[:, :, None]).squeeze(-1)
        log_x_theta_at_m = model_output[:, :, self.mask_index]
        x_theta_at_m = log_x_theta_at_m.exp()
        
        term_1_coef = dt / t
        term_1_log_nr = torch.log(alpha_t * x_theta_at_m / t + 1)
        term_1_log_dr = log_x_theta_at_x0
        
        term_2_coef = 1 - dt / t
        term_2_log_nr = term_1_log_nr
        term_2_log_dr = torch.log(alpha_s * x_theta_at_m / (t - dt) + 1)

        L_vb_masked = (
            term_1_coef * (term_1_log_nr - term_1_log_dr)
            + term_2_coef * (term_2_log_nr - term_2_log_dr))

        L_vb = L_vb_masked * (xt == self.mask_index)

        return self.T * L_vb

    def forward_diffusion(self, latents):
        """Forward diffusion process, adds noise to the latents."""

        t = self.sample_timestep(latents.shape[0], latents.device)
        if self.T > 0:
            t = (t * self.T).to(torch.int)
            t = t / self.T
            # t \in {1/T, 2/T, ..., 1}
            t += (1 / self.T)

        sigma, dsigma = self.noise(t)
        unet_conditioning = sigma[:, None]
        move_chance = 1 - torch.exp(-sigma[:, None])
        
        noised_latents = self.q_xt(latents, move_chance)
        model_output = self.forward(noised_latents, unet_conditioning)
    
        if self.T > 0:
            diffusion_loss = self.d3pm_loss(model_output=model_output, xt=noised_latents, x0=latents, t=t)
            return diffusion_loss
        # SUBS parameterization, continuous time.
        else:
            log_p_theta = torch.gather(input=model_output, dim=-1, index=latents[:, :, None]).squeeze(-1)
            return - log_p_theta * (dsigma / torch.expm1(sigma))[:, None]
    

    ######### LOSS CALCULATIONS #########
    def maybe_sub_sample(self, x0, attention_mask):
        # seqlen = x0.shape[1]
        # print(seqlen)
        # if seqlen > self.config.model.length:
        #     assert seqlen == 2 * self.config.model.length
        #     # cropping is needed for text8-crop dataset
        #     # try the same starting point for now
        #     start = np.random.choice(self.config.model.length)
        #     end = start + self.config.model.length
        #     input_tokens = x0[:, start: end]
        #     output_tokens = x0[:, start + 1: end + 1]
        #     new_attention_mask = attention_mask[:, start: end]

        #     # Helps with validation PPL, since the val
        #     # examples will all start and end with BOS/EOS
        #     input_tokens[:, 0] = self.tokenizer.bos_token_id
        #     output_tokens[:, -1] = self.tokenizer.eos_token_id
    
        # elif self.parameterization == 'ar':
        #     input_tokens = x0[:, :-1]
        #     output_tokens = x0[:, 1:]
        #     new_attention_mask = attention_mask[:, 1:]
        # else:
        input_tokens = x0
        output_tokens = None
        new_attention_mask = attention_mask
    
        return input_tokens, output_tokens, new_attention_mask

    def compute_loss(self, latents, attention_mask):
        """"Average of MLM losses to stabilize training"""
        (input_tokens, output_tokens, attention_mask) = self.maybe_sub_sample(latents, attention_mask)
        loss = self.forward_diffusion(input_tokens)

        nlls = loss * attention_mask
        count = attention_mask.sum()
        batch_nll = nlls.sum()
        token_nll = batch_nll / count

        return Loss(loss=token_nll, nlls=nlls, token_mask=attention_mask)


    ######### TRAINING #########
    def training_step(self, batch, batch_idx):
        latents, attention_mask = batch
        loss = self.compute_loss(latents, attention_mask)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        return optimizer

    def validation_step(self, batch):
        latents, attention_mask = batch
        loss = self.compute_loss(latents, attention_mask)
        return loss
    

    ######### GENERATION #########
    def sample_prior(self, *batch_dims):
        return self.mask_index * torch.ones(* batch_dims, dtype=torch.int64)

    def sample_categorical(categorical_probs):
        gumbel_norm = (1e-10 - (torch.rand_like(categorical_probs) + 1e-10).log())
        return (categorical_probs / gumbel_norm).argmax(dim=-1)

    def ddpm_caching_update(self, x, t, dt, p_x0=None):
        assert self.config.noise.type == 'loglinear'
        sigma_t, _ = self.noise(t)
        if t.ndim > 1:
            t = t.squeeze(-1)
        assert t.ndim == 1
        move_chance_t = t[:, None, None]
        move_chance_s = (t - dt)[:, None, None]
        assert move_chance_t.ndim == 3, move_chance_t.shape
        if p_x0 is None:
            p_x0 = self.forward(x, sigma_t).exp()
    
        assert move_chance_t.ndim == p_x0.ndim
        q_xs = p_x0 * (move_chance_t - move_chance_s)
        q_xs[:, :, self.mask_index] = move_chance_s[:, :, 0]
        _x = self.sample_categorical(q_xs)
        
        copy_flag = (x != self.mask_index).to(x.dtype)
        return p_x0, copy_flag * x + (1 - copy_flag) * _x


    @torch.no_grad()
    def sample_subs_guidance(self, n_samples, stride_length, num_strides, dt=0.001):
        ones = torch.ones(n_samples, dtype=self.dtype,device=self.device)
        num_steps = int(1 / dt)
        sampling_steps = 0
        intermediate_tokens = []
        target = None

        for _ in range(num_strides + 1):
            p_x0_cache = None
            x = self._sample_prior(n_samples,self.config.model.length).to(self.device)
            
            if target is not None:
                x[:, : -stride_length] = target
            
            for i in range(num_steps + 1):
                p_x0_cache, x_next = self.ddpm_caching_update(x=x, t=(1 - i * dt) * ones, dt=dt, p_x0=p_x0_cache)
                if (not torch.allclose(x_next, x) or self.time_conditioning):
                    p_x0_cache = None
                    sampling_steps += 1
                x = x_next
            x = self.forward(x, 0 * ones).argmax(dim=-1)
            intermediate_tokens.append(x[:, :stride_length].cpu().numpy())
            target = x[:, stride_length:]
    
        intermediate_tokens.append(target.cpu().numpy())
        intermediate_text_samples = []
        sequence_lengths = ((np.concatenate(intermediate_tokens, axis=1)[:, 1:]
                                 == self.tokenizer.eos_token_id).cumsum(-1) == 0).sum(-1)
    
        for i in range(2, len(intermediate_tokens) + 1):
            intermediate_text_samples.append(self.tokenizer.decode(np.concatenate(intermediate_tokens[:i], axis=1)))
        
        return (sampling_steps, intermediate_text_samples,
            sequence_lengths)

    def restore_model_and_semi_ar_sample(self, stride_length, num_strides, dt=0.001):
        """Generate samples from the model."""
        # Lightning auto-casting is not working in this method for some reason
        self.backbone.eval()
        self.noise.eval()
    
        (sampling_steps, samples, sequence_lengths) = self.sample_subs_guidance(n_samples=self.config.Loader.BATCH_SIZE,stride_length=stride_length,num_strides=num_strides,dt=dt)

        self.backbone.train()
        self.noise.train()
        return sampling_steps, samples, sequence_lengths