import sys
import pandas as pd
from transformers import AutoTokenizer

MODEL_PATH = "./Qwen3-0.6B"
CSV_FILE = "gravity_scan_0.6B.csv"

def query(word):
    print(f"[*] 正在从 {MODEL_PATH} 探测词汇: '{word}' ...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    
    if not pd.io.common.file_exists(CSV_FILE):
        print(f"[-] 错误: 未找到基准表 {CSV_FILE}！")
        return
        
    df = pd.read_csv(CSV_FILE)
    
    # 1. 真实分词切分分析
    token_ids = tokenizer.encode(word, add_special_tokens=False)
    print("\n" + "=" * 75)
    print(f"🎯 目标词: 【{word}】 ➔ 分词器切分 ID: {token_ids}")
    print("-" * 75)
    print(f"{'Token ID':<10} | {'当前字符':<16} | {'强吸引至 (Top-1)':<18} | {'置信度':<8} | {'引力熵'}")
    print("-" * 75)
    
    for tid in token_ids:
        row = df[df["token_id"] == tid]
        if not row.empty:
            r = row.iloc[0]
            print(f"{tid:<10} | {str(r['char']):<16} | {str(r['top1_target']):<18} | {r['top1_prob']*100:>6.1f}% | {r['entropy']:.4f}")
        else:
            print(f"{tid:<10} | [未找到记录]")
            
    print("=" * 75)
    
    # 2. 全词表模糊检索相关 Token
    print(f"\n[*] 全词表中所有包含 '{word}' 的独立 Token 吸引子分布:")
    fuzzy = df[df["char"].str.contains(word, case=False, na=False)]
    if not fuzzy.empty:
        for _, r in fuzzy.iterrows():
            print(f"  ➜ ID: {int(r['token_id']):<6} | Token: {str(r['char']):<14} ➔ 强吸引至: {str(r['top1_target']):<14} (置信度: {r['top1_prob']*100:>5.1f}%, 熵: {r['entropy']:.3f})")
    else:
        print("  [未找到其他衍生 Token]")
    print("=" * 75 + "\n")

if __name__ == "__main__":
    target_word = sys.argv[1] if len(sys.argv) > 1 else "bikini"
    query(target_word)