import inspect

import pandas as pd
from transformers import AutoTokenizer
from lnp_ml.dataset import process_dataframe, SMILES_COL
from lnp_ml.modeling.layers.llm_prompt import LLMPromptEncoder

# 从源头读，避免脚本和模型的阈值各说各话
MAXLEN = inspect.signature(LLMPromptEncoder.__init__).parameters["max_length"].default

tok = AutoTokenizer.from_pretrained("models/qwen2.5-7b-instruct", trust_remote_code=True)
df = process_dataframe(pd.read_csv("data/interim/internal.csv"))
smis = sorted(set(df[SMILES_COL].dropna()), key=len)

enc = LLMPromptEncoder.__new__(LLMPromptEncoder)   # 不加载 7B 权重，只借用格式化方法
enc._is_biot5 = False
# _build_rag_prompt 现在还要读分位边界。这里必须给非空值：留空会让 _qbin 返回空串，
# 测出来的 token 数比真实 prompt 少约 8/邻居，等于白测。
# 具体数值不影响长度（Q1/5 和 Q5/5 token 数相同），只决定落在哪个桶。
enc._rag_pool_qedges = {
    "delivery": [-0.80, -0.30, 0.20, 0.70],
    "size": [-0.90, -0.20, 0.40, 1.10],
}

def nb_of(s):
    return {"smiles": s, "sim": 0.812, "delivery": 1.234,
            "extra": {"size": -0.456, "pdi": 1, "ee": 2, "toxic": 0,
                      "biodist": [0.12, 0.34, 0.21, 0.05, 0.18, 0.07, 0.03]}}

worst = smis[-1]
print(f"分子数={len(smis)}  SMILES 长度 {len(smis[0])}~{len(worst)} 字符")
print(f"分子数={len(smis)}  SMILES 长度 {len(smis[0])}~{len(worst)} 字符  max_length={MAXLEN}")
for k in (2, 4, 8):
    # 最坏情况：目标分子和全部邻居都取最长的那条
    p = LLMPromptEncoder._build_rag_prompt(enc, worst, [nb_of(worst)] * k)
    n = len(tok(p)["input_ids"])
    status = "OK" if n <= MAXLEN else f"超出 {n - MAXLEN} tokens，会被截断"
    print(f"最坏 rag_top_k={k}: {n:5d} tokens  余量 {MAXLEN - n:5d}  {status}")