# import json

# def read_jsonl(path: str):
#     data = []
#     with open(path, "r", encoding="utf-8") as f:
#         for line in f:
#             line = line.strip()
#             if line:
#                 data.append(json.loads(line))
#     return data

# data_1 = read_jsonl("/mnt/bn/life-mllm/users/cxr/hallucination/OPERA/log/Qwen2.5-VL-7B-Instruct/ours_20250826_143459.jsonl")
# data_2 = read_jsonl("/mnt/bn/life-mllm/users/cxr/hallucination/OPERA/log/Qwen2.5-VL-7B-Instruct/ours_20250826_143737.jsonl")

# for i, (a, b) in enumerate(zip(data_1, data_2)):
#     if a != b:
#         print('tcd:', a)
#         print('qwen2.5vl:', b)
#         print()
import torch 
import torch.nn.functional as F

def get_candidate_tokens_logits(logits, beta):
    """
    直接基于logits筛选满足条件的token索引（不做归一化），返回格式同torch.topk的topk_idx
    
    参数:
        logits: 模型输出的未归一化分数，形状为 (batch_size, vocab_size)
        beta: 阈值系数，范围通常为 [0, 1]
    
    返回:
        list[torch.Tensor]: 每个元素为对应批次样本的候选token索引张量，
                           形状为 (num_candidates,)（num_candidates随样本变化）
    """
    # 计算每个样本的最大logits值（保留维度用于广播）
    max_logits, _ = logits.max(dim=-1, keepdim=True)  # 形状: (batch_size, 1)
    
    # 计算每个样本的阈值（beta × 最大logits）
    thresholds = beta * max_logits  # 形状: (batch_size, 1)
    
    # 生成布尔掩码：每个位置标记是否满足 logits ≥ 阈值
    masks = logits >= thresholds  # 形状: (batch_size, vocab_size)
    
    # 按批次提取满足条件的索引
    candidate_indices = []
    for i in range(logits.size(0)):
        indices = torch.where(masks[i])[0]  # 第i个样本的候选索引
        candidate_indices.append(indices)
    
    return torch.stack(candidate_indices, dim=0)

# 模拟批次输入（batch_size=2，vocab_size=5）
logits = torch.tensor([
    [1.2, 3.5, 2.1, 0.8, 4.0],  # 样本1的logits
    [1.2, 3.5, 2.1, 0.8, 4.0],
])
beta = 0.6

# 获取候选索引
candidates = get_candidate_tokens_logits(logits, beta)

print("样本1候选索引:", candidates)  # tensor([1, 4])（满足概率≥0.6×最大概率）
# topk_vals, topk_idx = torch.topk(logits, 3, dim=-1)
topk_vals = torch.gather(logits, dim=-1, index=candidates)
print(topk_vals)

def logits_entropy(logits: torch.Tensor, dim: int = -1, k: int = 8) -> torch.Tensor:
    """
    计算 logits 对应的熵 (Shannon entropy)。
    
    参数:
        logits: torch.Tensor, shape [..., vocab_size]
        dim: 计算概率分布的维度，一般是最后一维 (vocab_size)
    返回:
        熵值张量, shape [...]
    """
    # softmax 得到概率分布
    probs = F.softmax(logits, dim=dim)
    # 避免 log(0)
    # topk_probs, _ = torch.topk(probs, k)
    topk_probs = probs
    log_probs = torch.log(topk_probs + 1e-12)
    # 计算熵
    entropy = -(topk_probs * log_probs).sum(dim=dim).item()
    return entropy

print(logits_entropy(torch.tensor([[25.]])))