# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

import torch
import torch.nn.functional as F
import torch.distributions as dists
from typing import Dict, Optional


def get_token_ids_from_config(config) -> Dict[str, int]:
    """Extract all token IDs from the configuration object.
    
    Args:
        config: Configuration object (LocateAnythingConfig or similar)
        
    Returns:
        Dictionary containing all token IDs
    """
    token_ids = {}
    
    # Get from main config
    token_ids['box_start_token_id'] = getattr(config, 'box_start_token_id', 151668)
    token_ids['box_end_token_id'] = getattr(config, 'box_end_token_id', 151669)
    token_ids['coord_start_token_id'] = getattr(config, 'coord_start_token_id', 151677)
    token_ids['coord_end_token_id'] = getattr(config, 'coord_end_token_id', 152677)
    token_ids['ref_start_token_id'] = getattr(config, 'ref_start_token_id', 151672)
    token_ids['ref_end_token_id'] = getattr(config, 'ref_end_token_id', 151673)
    token_ids['none_token_id'] = getattr(config, 'none_token_id', 4064)
    
    # Get from text_config
    text_config = getattr(config, 'text_config', None)
    if text_config is not None:
        token_ids['null_token_id'] = getattr(text_config, 'null_token_id', 152678)
        token_ids['im_end_token_id'] = getattr(text_config, 'eos_token_id', 151645)
        token_ids['switch_token_id'] = getattr(text_config, 'switch_token_id', 152679)
        token_ids['default_mask_token_id'] = getattr(text_config, 'text_mask_token_id', 151676)
    else:
        token_ids['null_token_id'] = 152678
        token_ids['im_end_token_id'] = 151645
        token_ids['switch_token_id'] = 152679
        token_ids['default_mask_token_id'] = 151676
    
    return token_ids


def top_p_logits(
    logits: torch.Tensor,
    top_p: float = None
) -> torch.Tensor:
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_remove = cumulative_probs > top_p
    # Shift the indices to the right to keep the first token above the threshold
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
    return logits


def top_k_logits(
    logits: torch.Tensor,
    top_k: int = None
) -> torch.Tensor:
    top_k = min(top_k, logits.size(-1))  # Safety check
    # Remove all tokens with a probability less than the last token of the top-k
    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
    return logits


def apply_repetition_penalty(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
    repetition_penalty: float = 1.0
) -> torch.Tensor:
    """
    Apply repetition penalty to logits.
    
    Args:
        logits: Shape [batch_size, seq_len, vocab_size] or [batch_size, vocab_size]
        input_ids: Previously generated token ids, shape [batch_size, seq_len]
        repetition_penalty: Penalty factor. > 1.0 penalizes repetition, < 1.0 encourages it.
    
    Returns:
        Modified logits with repetition penalty applied.
    """
    if repetition_penalty == 1.0:
        return logits
    
    # Convert to 3D for vectorized computation
    if logits.dim() == 2:
        logits = logits.unsqueeze(1)  # [B, 1, V]
        squeeze_back = True
    else:
        squeeze_back = False
    
    batch_size, seq_len, vocab_size = logits.shape

    # Construct [B, V] bool mask marking tokens that have appeared in each batch
    device = logits.device
    token_mask = torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device)
    for b in range(batch_size):
        # Apply penalty only based on tokens already generated in this batch
        unique_tokens = input_ids[b].unique()
        # Prevent out-of-bounds: only keep IDs within vocab range
        valid_tokens = unique_tokens[(unique_tokens >= 0) & (unique_tokens < vocab_size)]
        if valid_tokens.numel() > 0:
            token_mask[b, valid_tokens] = True

    # Expand to [B, L, V] to align with logits
    token_mask = token_mask.unsqueeze(1).expand(-1, seq_len, -1)

    # Divide positive values by penalty, multiply negative values by penalty
    positive = logits > 0
    negative = ~positive

    # Apply penalty only at mask positions
    logits = torch.where(token_mask & positive, logits / repetition_penalty, logits)
    logits = torch.where(token_mask & negative, logits * repetition_penalty, logits)
    
    if squeeze_back:
        logits = logits.squeeze(1)
    
    return logits



def sample_tokens_ar(
    logits: torch.Tensor,
    generated: torch.Tensor,
    token_ids: Dict[str, int],
    **generate_kwargs,
):
    """
    Lightweight sampling function for AR single-step sampling only.
    
    Args:
        logits: [batch_size, vocab_size] or [batch_size, 1, vocab_size]
        generated: [batch_size, seq_len]
    """
    # Convert to 3D for reusing repetition penalty and clipping logic
    if logits.dim() == 2:
        logits = logits.unsqueeze(1)  # [B, 1, V]
    batch_size, seq_len, vocab_size = logits.shape
    assert seq_len == 1, "sample_tokens_ar only supports single-step AR sampling (seq_len == 1)"

    repetition_penalty = generate_kwargs.get('repetition_penalty', 1.0)
    temperature = generate_kwargs.get('temperature', 0)
    top_p = generate_kwargs.get('top_p', None)
    top_k = generate_kwargs.get('top_k', None)

    # Apply repetition penalty only based on historically generated tokens
    if repetition_penalty != 1.0:
        logits = apply_repetition_penalty(logits, generated, repetition_penalty)

    if temperature > 0:
        logits = logits / temperature
    if top_p is not None and top_p < 1:
        logits = top_p_logits(logits, top_p)
    if top_k is not None:
        logits = top_k_logits(logits, top_k)

    probs = torch.softmax(logits, dim=-1)

    if temperature > 0:
        try:
            x0 = dists.Categorical(probs=probs).sample()
            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
        except Exception:
            confidence, x0 = probs.max(dim=-1)
    else:
        # For greedy: directly take the token with maximum probability
        confidence, x0 = probs.max(dim=-1)

    # Keep interface consistent with sample_tokens: return [B, 1, V] / [B, 1] shape
    return probs, confidence, x0, None, None