from __future__ import annotations
import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask
from transformers.utils import ModelOutput
from dataclasses import dataclass
from transformers.cache_utils import Cache

@dataclass
class SimpleOutputWithPast(ModelOutput):
    loss: torch.FloatTensor | None = None
    logits: torch.FloatTensor | None = None
    causal_logits: torch.FloatTensor | None = None
    past_key_values: Cache | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None
from .utils import maybe_truncate_last_dim, pad_along_last_dim
from .configuration_nlduni import get_ministral_shift_sbd_logits
from .sbd_attention import build_image_attention_doc_ids

def _resolve_shift_sbd_logits(model, shift_logits):
    model_module = model.module if hasattr(model, 'module') else model
    config = getattr(model_module, 'config', None)
    configured = get_ministral_shift_sbd_logits(config) if config is not None else None
    has_alignment_contract = config is not None
    if shift_logits is None:
        return bool(configured) if configured is not None else False
    requested = bool(shift_logits)
    if has_alignment_contract and requested != configured:
        raise ValueError('Inference shift_logits must match config.shift_sbd_logits to avoid a one-token alignment error. Update the model config explicitly if this override is intentional.')
    return requested

def wte(model, x, t2i_inference=False, gen_shape=None, x_gen=None, inputs_embeds_curr=None, new_token_mask=None, x_gen_child_embeddings=None, gen_parent_mask=None, gen_source_null_mask=None, gen_role=None):
    if t2i_inference:
        if (x_gen is None) == (x_gen_child_embeddings is None):
            raise ValueError('provide exactly one of x_gen or x_gen_child_embeddings')
        if new_token_mask is None:
            new_token_mask = x >= INT_MAX
        if gen_parent_mask is not None or gen_source_null_mask is not None or gen_role is not None:
            raise ValueError('structured generation masks and roles require Infinity BSQ')
        gen_latents_comp_embeds = model.call_gen_embedding(x_gen, gen_shape, child_embeddings=x_gen_child_embeddings)
        if inputs_embeds_curr is None:
            x_txt_only = x.clone()
            x_txt_only[new_token_mask] = 0
            inputs_embeds_curr = model.embed_tokens(x_txt_only)
        inputs_embeds_curr[new_token_mask] = pad_along_last_dim(gen_latents_comp_embeds, inputs_embeds_curr.shape[-1]).view(-1, inputs_embeds_curr.shape[-1])
    else:
        inputs_embeds_curr = model.embed_tokens(x)
        new_token_mask = None
    return (inputs_embeds_curr, new_token_mask)
INT_MAX = 1000000
SBD_FLEX_ATTENTION_PAD_MULTIPLE = 512

def _round_up_to_multiple(length, multiple):
    if isinstance(multiple, bool) or not isinstance(multiple, int) or multiple < 1:
        raise ValueError('padding multiple must be a positive integer')
    return (length + multiple - 1) // multiple * multiple

def single_stream_sbd_sampling_mask(block_size, past_length, b, h, q_idx, kv_idx, query_image_doc_id, key_value_image_doc_id):
    del block_size, h
    absolute_q = q_idx + past_length
    token_causal = absolute_q >= kv_idx
    image_query = query_image_doc_id[b, q_idx] > 0
    same_image_doc = image_query & (query_image_doc_id[b, q_idx] == key_value_image_doc_id[b, kv_idx])
    return token_causal | same_image_doc

def _create_single_stream_sbd_sampling_block_mask(input_embeddings, *, past_key_values, block_size, input_modality_indices, generation_token_mask=None, pad_to_multiple=None):
    if isinstance(block_size, bool) or not isinstance(block_size, int) or block_size < 1:
        raise ValueError('SBD diffusion sampling requires a positive integer block_size')
    query_length = input_embeddings.shape[1]
    past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
    key_value_length = past_length + query_length
    padded_query_length = _round_up_to_multiple(query_length, pad_to_multiple) if pad_to_multiple is not None else query_length
    padded_key_value_length = _round_up_to_multiple(key_value_length, pad_to_multiple) if pad_to_multiple is not None else key_value_length
    if input_modality_indices is None:
        raise ValueError('SBD diffusion sampling requires input_modality_indices to route same-block generation attention')
    if tuple(input_modality_indices.shape) != tuple(input_embeddings.shape[:2]):
        raise ValueError('SBD diffusion sampling requires input_modality_indices to align with the current input embeddings')
    query_transformer_routing = input_modality_indices.to(dtype=torch.bool)
    if generation_token_mask is None:
        generation_token_mask = query_transformer_routing
    elif tuple(generation_token_mask.shape) != tuple(input_embeddings.shape[:2]):
        raise ValueError('SBD diffusion sampling requires generation_token_mask to align with the current input embeddings')
    query_target_generation_routing = generation_token_mask.to(dtype=torch.bool)
    cached_generation_routing = torch.zeros(query_transformer_routing.shape[0], past_length, dtype=torch.bool, device=query_transformer_routing.device)
    padded_generation_routing = torch.zeros(query_transformer_routing.shape[0], padded_key_value_length - key_value_length, dtype=torch.bool, device=query_transformer_routing.device)
    key_value_target_generation_routing = torch.cat((cached_generation_routing, query_target_generation_routing, padded_generation_routing), dim=1)
    key_value_valid = (torch.arange(padded_key_value_length, device=query_transformer_routing.device).unsqueeze(0) < key_value_length).expand(query_transformer_routing.shape[0], -1).contiguous()
    key_value_image_doc_id = build_image_attention_doc_ids(key_value_target_generation_routing, valid_tokens=key_value_valid)
    query_image_doc_id = torch.cat((key_value_image_doc_id[:, past_length:past_length + query_length], torch.zeros(query_transformer_routing.shape[0], padded_query_length - query_length, dtype=key_value_image_doc_id.dtype, device=query_transformer_routing.device)), dim=1)
    past_length_tensor = torch.tensor(past_length, dtype=torch.long, device=query_transformer_routing.device)
    query_valid = (torch.arange(padded_query_length, device=query_transformer_routing.device).unsqueeze(0) < query_length).expand(query_transformer_routing.shape[0], -1).contiguous()

    def mask_mod(b, h, q_idx, kv_idx):
        logical_attention = single_stream_sbd_sampling_mask(block_size, past_length_tensor, b, h, q_idx, kv_idx, query_image_doc_id, key_value_image_doc_id) & query_valid[b, q_idx] & key_value_valid[b, kv_idx]
        padded_query_attention = ~query_valid[b, q_idx] & (kv_idx == 0)
        return logical_attention | padded_query_attention
    return create_block_mask(mask_mod, B=input_embeddings.shape[0], H=None, Q_LEN=padded_query_length, KV_LEN=padded_key_value_length, device=input_embeddings.device)

def get_logits(model, input_emnbeddings, modality_indices=None, t2i_inference=False, past_key_values=None, gen_shape=None, timesteps=None, input_modality_indices=None, generation_attention_mask=None, sbd_diffusion_sampling=False, full_attention_approx=False, loop_k=None):
    if t2i_inference:
        resolved_loop_k = None
        if input_modality_indices is None:
            input_modality_indices = modality_indices
        attention_kwargs = {}
        if full_attention_approx:
            attention_kwargs['overwrite_block_mask'] = 'full'
            attention_kwargs['overwrite_attn_impl'] = 'flash_attn'
        elif sbd_diffusion_sampling:
            if generation_attention_mask is None:
                generation_attention_mask = modality_indices
            attention_kwargs['overwrite_block_mask'] = _create_single_stream_sbd_sampling_block_mask(input_emnbeddings, past_key_values=past_key_values, block_size=getattr(model.config, 'block_size', None), input_modality_indices=input_modality_indices, generation_token_mask=generation_attention_mask, pad_to_multiple=SBD_FLEX_ATTENTION_PAD_MULTIPLE)
        else:
            attention_kwargs['overwrite_attn_impl'] = 'flash_attn'
        output = model(None, input_embeddings=input_emnbeddings, modality_indices=input_modality_indices, output_hidden_states=True, past_key_values=past_key_values, use_cache=None, single_stream_image_diffusion=True, loop_k=None, **attention_kwargs)
        hidden_states = getattr(output, 'last_hidden_state', None)
        if hidden_states is None:
            hidden_states = output.hidden_states[-1]
        gen_hidden_states = hidden_states[modality_indices]
        gen_hidden_states = maybe_truncate_last_dim(gen_hidden_states, model.config.d_model_gen)
        gen_logits = model.call_gen_predictor(gen_hidden_states, gen_shape, timesteps=timesteps)
        seq_len_per_img = np.prod(gen_shape)
        if len(gen_logits.shape) == 2:
            gen_logits = gen_logits.view(-1, seq_len_per_img, gen_logits.shape[-1])
        else:
            gen_logits = gen_logits.view(-1, seq_len_per_img, *gen_logits.shape[-2:])
        return gen_logits
    else:
        raise ValueError('get_logits is the image sampler helper')

def add_gumbel_noise(logits, temperature):
    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_transfer_index(logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None, neg_entropy=False):
    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
    x0 = torch.argmax(logits_with_noise, dim=-1)
    if remasking == 'low_confidence':
        p = F.softmax(logits, dim=-1)
        x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1)
    elif remasking == 'top_p_margin':
        p = F.softmax(logits, dim=-1)
        top2 = torch.topk(p, k=2, dim=-1).values
        margin = top2[..., 0] - top2[..., 1]
        plus_inf = torch.full_like(margin, 1e309)
        minus_inf = torch.full_like(margin, -1e309)
        masked_for_min = torch.where(mask_index, margin, plus_inf)
        masked_for_max = torch.where(mask_index, margin, minus_inf)
        row_min = masked_for_min.amin(dim=1, keepdim=True)
        row_max = masked_for_max.amax(dim=1, keepdim=True)
        denom = row_max - row_min
        normalized = torch.zeros_like(margin)
        nonzero = denom > 0
        normalized = torch.where(mask_index & nonzero, (margin - row_min) / (denom + 1e-12), normalized)
        normalized = torch.where(mask_index & ~nonzero, torch.ones_like(normalized), normalized)
        x0_p = normalized
    elif remasking == 'random':
        x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
    else:
        raise NotImplementedError(remasking)
    if neg_entropy:
        p = F.softmax(logits, dim=-1)
        epsilon = 1e-10
        log_probs = torch.log(p + epsilon)
        confidence_scores = torch.sum(p * log_probs, dim=-1)
    else:
        confidence_scores = x0_p
    x0 = torch.where(mask_index, x0, x)
    confidence = torch.where(mask_index, confidence_scores, -np.inf)
    transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
    if threshold is not None:
        num_transfer_tokens = mask_index.sum(dim=1, keepdim=True)
    for j in range(confidence.shape[0]):
        (_, select_index) = torch.topk(confidence[j], k=num_transfer_tokens[j])
        transfer_index[j, select_index] = True
        if threshold is not None:
            for k in range(1, num_transfer_tokens[j]):
                if confidence[j, select_index[k]] < threshold:
                    transfer_index[j, select_index[k]] = False
    return (x0, transfer_index)

def get_num_transfer_tokens(mask_index, steps: int):
    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, :int(remainder[i])] += 1
    return num_transfer_tokens

def simple_fwd(model, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, past_key_values=None, **kwargs):
    enc_out = model.encoder(past_key_values=past_key_values, input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, overwrite_attn_impl='flash_attn', **kwargs)
    logits = model.diffusion_head(enc_out.last_hidden_state)
    return SimpleOutputWithPast(loss=logits, logits=logits, causal_logits=None, past_key_values=enc_out.past_key_values, hidden_states=None, attentions=None)

@torch.no_grad()
def generate_with_prefix_cache_block_diff(model, prompt=None, prompt_embeds=None, steps=128, gen_length=128, block_length=128, temperature=0.0, remasking='low_confidence', mask_id=126336, threshold=None, factor=None, shift_logits=None, neg_entropy=False, causal_context=False, eos_token_id=None, max_thinking_tokens=None, end_think_token_id=None):
    shift_logits = _resolve_shift_sbd_logits(model, shift_logits)
    dream_style = shift_logits
    if (prompt is None) == (prompt_embeds is None):
        raise ValueError('Exactly one of `prompt` or `prompt_embeds` must be provided.')
    if prompt is not None:
        prompt_ids = prompt
        prompt_len = prompt_ids.shape[1]
        x_accum = prompt_ids.clone()
        B = prompt_ids.shape[0]
        token_device = prompt_ids.device
        token_dtype = prompt_ids.dtype
    else:
        prompt_ids = None
        prompt_len = prompt_embeds.shape[1]
        B = prompt_embeds.shape[0]
        token_device = prompt_embeds.device
        token_dtype = torch.long
        x_accum = torch.full((B, prompt_len), mask_id, dtype=token_dtype, device=token_device)
    if gen_length % block_length != 0:
        gen_length = (gen_length // block_length + 1) * block_length
    num_blocks = gen_length // block_length
    assert steps % num_blocks == 0
    steps_per_block = steps // num_blocks
    nfe = 0
    model_module = model.module if hasattr(model, 'module') else model
    for layer in model_module.encoder.layers:
        layer.self_attn.mode = 'bidirectional'
    if causal_context:
        for layer in model_module.encoder.layers:
            if hasattr(layer.self_attn, 'diffusion_lm'):
                layer.self_attn.diffusion_lm = False
    output = simple_fwd(model, input_ids=prompt_ids, inputs_embeds=prompt_embeds, use_cache=True, use_causal_mask=causal_context)
    past_key_values = output.past_key_values
    if causal_context:
        for layer in model_module.encoder.layers:
            if hasattr(layer.self_attn, 'diffusion_lm'):
                layer.self_attn.diffusion_lm = True
    next_token = None
    if causal_context:
        last_logit = output.logits[:, -1, :]
        if temperature > 0:
            probs = torch.softmax(last_logit / temperature, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
        else:
            next_token = torch.argmax(last_logit, dim=-1, keepdim=True)
    next_logits_context = None
    if dream_style:
        next_logits_context = output.logits[:, -1:, :]
    for num_block in range(num_blocks):
        mask_block = torch.ones((B, block_length), dtype=token_dtype, device=token_device) * mask_id
        if causal_context:
            mask_block[:, 0] = next_token[:, 0]
        x_accum = torch.cat([x_accum, mask_block], dim=1)
        current_block_start = prompt_len + num_block * block_length
        block_slice = slice(current_block_start, current_block_start + block_length)
        if end_think_token_id is not None and max_thinking_tokens is not None:
            tokens_before_block = num_block * block_length
            tokens_after_block = tokens_before_block + block_length
            if tokens_after_block > max_thinking_tokens:
                gen_so_far = x_accum[:, prompt_len:current_block_start]
                has_end_think = (gen_so_far == end_think_token_id).any(dim=1) if gen_so_far.size(1) > 0 else torch.zeros(B, dtype=torch.bool, device=token_device)
                if not has_end_think.all():
                    if tokens_before_block < max_thinking_tokens:
                        offset = max_thinking_tokens - tokens_before_block
                    else:
                        offset = 0
                    inject_pos = current_block_start + offset
                    for b in range(B):
                        if not has_end_think[b]:
                            x_accum[b, inject_pos] = end_think_token_id
        mask_block_idx0 = x_accum[:, block_slice] == mask_id
        if dream_style:
            schedule_mask = mask_block_idx0
        else:
            schedule_mask = mask_block_idx0
        num_transfer_tokens = get_num_transfer_tokens(schedule_mask, steps_per_block)
        for i in range(steps_per_block):
            mask_block_idx = x_accum[:, block_slice] == mask_id
            if mask_block_idx.sum() == 0:
                break
            nfe += 1
            logits_block = simple_fwd(model, x_accum[:, block_slice], past_key_values=past_key_values, use_cache=False).logits
            if dream_style:
                if block_length == 1:
                    logits_use = next_logits_context
                else:
                    logits_use = torch.cat([next_logits_context, logits_block[:, :-1, :]], dim=1)
                mask_use = mask_block_idx
                x_use = x_accum[:, block_slice]
                (x0, transfer_idx) = get_transfer_index(logits_use, temperature, remasking, mask_use, x_use, num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, neg_entropy=neg_entropy)
                cur = x_accum[:, block_slice].clone()
                cur[transfer_idx] = x0[transfer_idx]
                x_accum[:, block_slice] = cur
            else:
                (x0, transfer_idx) = get_transfer_index(logits_block, temperature, remasking, mask_block_idx, x_accum[:, block_slice], num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, neg_entropy=neg_entropy)
                cur = x_accum[:, block_slice].clone()
                cur[transfer_idx] = x0[transfer_idx]
                x_accum[:, block_slice] = cur
            if eos_token_id is not None:
                block_tokens = x_accum[:, block_slice]
                eos_mask = block_tokens == eos_token_id
                any_eos = eos_mask.any(dim=1)
                if any_eos.any():
                    after_eos = eos_mask.cumsum(dim=1).bool()
                    mask_before = (block_tokens == mask_id) & ~after_eos
                    if (any_eos & ~mask_before.any(dim=1)).any():
                        break
        if causal_context:
            for layer in model_module.encoder.layers:
                if hasattr(layer.self_attn, 'diffusion_lm'):
                    layer.self_attn.diffusion_lm = False
        output = simple_fwd(model, x_accum[:, block_slice], past_key_values=past_key_values, use_cache=True, use_causal_mask=causal_context)
        past_key_values = output.past_key_values
        nfe += 1
        if causal_context:
            for layer in model_module.encoder.layers:
                if hasattr(layer.self_attn, 'diffusion_lm'):
                    layer.self_attn.diffusion_lm = True
            last_logit = output.logits[:, -1, :]
            if temperature > 0:
                probs = torch.softmax(last_logit / temperature, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(last_logit, dim=-1, keepdim=True)
        if dream_style and num_block < num_blocks - 1:
            next_logits_context = output.logits[:, -1:, :]
        if eos_token_id is not None:
            gen_so_far = x_accum[:, prompt_len:]
            is_eos = gen_so_far == eos_token_id
            has_eos = is_eos.any(dim=1)
            if has_eos.all():
                first_eos_pos = is_eos.to(torch.int64).argmax(dim=1)
                max_eos = first_eos_pos.max().item()
                if prompt_ids is None:
                    return (x_accum[:, prompt_len:prompt_len + max_eos + 1], nfe)
                return (x_accum[:, :prompt_len + max_eos + 1], nfe)
    if prompt_ids is None:
        return (x_accum[:, prompt_len:], nfe)
    return (x_accum, nfe)