import torch
import numpy as np
import torch.nn.functional as F

from transformers import AutoTokenizer, AutoModel


def add_gumbel_noise(logits, temperature):
    '''
    The Gumbel max is a method for sampling categorical distributions.
    According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality.
    Thus, we use float64.
    '''
    if temperature == 0:
        return logits
    logits = logits.to(torch.float64)
    noise = torch.rand_like(logits, dtype=torch.float64)
    gumbel_noise = (- torch.log(noise)) ** temperature
    return logits.exp() / gumbel_noise


def get_num_transfer_tokens(mask_index, steps):
    '''
    In the reverse process, the interval [0, 1] is uniformly discretized into steps intervals.
    Furthermore, because LLaDA employs a linear noise schedule (as defined in Eq. (8)),
    the expected number of tokens transitioned at each step should be consistent.

    This function is designed to precompute the number of tokens that need to be transitioned at each step.
    '''
    mask_num = mask_index.sum(dim=1, keepdim=True)

    base = mask_num // steps
    remainder = mask_num % steps

    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base

    for i in range(mask_num.size(0)):
        num_transfer_tokens[i, :remainder[i]] += 1

    return num_transfer_tokens


@ torch.no_grad()
def generate(model, prompt, steps=128, gen_length=128, block_length=128, temperature=0.,
             cfg_scale=0., remasking='low_confidence', mask_id=126336):
    '''
    Args:
        model: Mask predictor.
        prompt: A tensor of shape (1, L).
        steps: Sampling steps, less than or equal to gen_length.
        gen_length: Generated answer length.
        block_length: Block length, less than or equal to gen_length. If less than gen_length, it means using semi_autoregressive remasking.
        temperature: Categorical distribution sampling temperature.
        cfg_scale: Unsupervised classifier-free guidance scale.
        remasking: Remasking strategy. 'low_confidence' or 'random'.
        mask_id: The toke id of [MASK] is 126336.
    '''
    # 1. Initialization: Create the full sequence tensor 'x'.
    # It starts with the prompt, followed by `gen_length` [MASK] tokens.
    # This is the "canvas" that will be filled in.
    x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
    x[:, :prompt.shape[1]] = prompt.clone()

    # Keep track of where the original prompt is, so we don't modify it.
    prompt_index = (x != mask_id)

    # 2. Semi-Autoregressive Setup (Blocking)
    # The generation is split into 'num_blocks' chunks. This handles long generation
    # by generating `block_length` tokens at a time before moving to the next chunk.
    assert gen_length % block_length == 0
    num_blocks = gen_length // block_length
    
    # The total number of refinement steps is distributed among the blocks.
    assert steps % num_blocks == 0
    steps_per_block = steps // num_blocks

    # 3. Outer Loop: Process each block sequentially.
    for num_block in range(num_blocks):
        # Define the current working area (the block to be filled).
        # Note: The original code has a small typo `...:]`, corrected here for clarity.
        start_pos = prompt.shape[1] + num_block * block_length
        end_pos = prompt.shape[1] + (num_block + 1) * block_length
        block_mask_index = (x[:, start_pos:end_pos] == mask_id)

        # Calculate how many tokens to "unmask" or "confirm" in each refinement step for this block.
        # This ensures a steady, linear progression of unmasking.
        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps_per_block)

        # 4. Inner Loop: Iteratively refine the current block.
        for i in range(steps_per_block):
            # Get the indices of all currently masked tokens in the entire sequence.
            mask_index = (x == mask_id)

            # --- 4a. Prediction with optional Classifier-Free Guidance (CFG) ---
            if cfg_scale > 0.:
                # Create an unconditional version of the input by masking the prompt.
                un_x = x.clone()
                un_x[prompt_index] = mask_id
                
                # Run the model on both the conditional (x) and unconditional (un_x) inputs.
                x_ = torch.cat([x, un_x], dim=0)
                logits_cat = model(x_).logits
                logits, un_logits = torch.chunk(logits_cat, 2, dim=0)
                
                # Combine logits to steer the generation towards the prompt.
                # The formula is: unconditional + scale * (conditional - unconditional)
                # An algebraic simplification gives the line below.
                logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
            else:
                # If no CFG, just do a single forward pass.
                logits = model(x).logits

            # --- 4b. Candidate Generation ---
            # Add Gumbel noise for stochastic sampling. If temperature is 0, this is a simple argmax.
            logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
            # Get the most likely token prediction for every position in the sequence.
            x0 = torch.argmax(logits_with_noise, dim=-1)

            # --- 4c. Confidence Scoring for Remasking ---
            # Determine which of the new predictions to keep for the next step.
            if remasking == 'low_confidence':
                # Calculate the softmax probabilities of the predicted tokens.
                p = F.softmax(logits, dim=-1)
                # Get the probability of the chosen token `x0` at each position. This is the "confidence".
                x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1)
            elif remasking == 'random':
                # Use random scores as confidence for random unmasking.
                x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
            else:
                raise NotImplementedError(remasking)
            
            # --- 4d. Selecting Tokens to Update ---
            # We are only interested in updating tokens inside the current block.
            # Set confidence outside the current active generation area to -infinity to ignore them.
            x0_p[:, end_pos:] = -np.inf

            # Only consider predictions for positions that are currently masked.
            # Original tokens (prompt and previously confirmed tokens) should have -infinity confidence.
            confidence = torch.where(mask_index, x0_p, -np.inf)
            
            # Replace the content of `x` at masked positions with the new predictions (`x0`).
            x0 = torch.where(mask_index, x0, x)
            
            # This will hold the indices of the tokens we decide to "confirm" in this step.
            transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
            
            # For each item in the batch (here, just 1)...
            for j in range(confidence.shape[0]):
                # ...select the `k` tokens with the HIGHEST confidence scores among the masked positions.
                # `k` is determined by `num_transfer_tokens` for the current step `i`.
                _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j, i])
                # Mark these high-confidence positions to be updated.
                transfer_index[j, select_index] = True

            # --- 4e. State Update ---
            # Update the main tensor 'x' by replacing [MASK] tokens with the selected high-confidence predictions.
            # The other [MASK]s remain for the next refinement iteration.
            x[transfer_index] = x0[transfer_index]

    # 5. Return Final Generation
    # Once all blocks and steps are complete, return the generated part of the sequence.
    return x