"""Shared helpers: tokenizer, bucket choice, scoring API over any backend that maps ids[B,T] -> lp[B,T-1]."""
import json, os
import numpy as np
from tokenizers import Tokenizer
from qdec import HF

TOK = Tokenizer.from_file(f"{HF}/tokenizer.json")
BOS = TOK.token_to_id("<bos>")
PAD = TOK.token_to_id("<pad>")
BUCKETS = [16, 32, 64, 96]
CACHE = os.environ.get("MCBPMF_LM_WORK", ".") + "/results/walk2/dec_cache.v2.f3c75593dbd9.7ab4ee7.fdde1635a116.jsonl"


def ids(text):
    return [BOS] + TOK.encode(text).ids


def bucket(n):
    for b in BUCKETS:
        if n <= b:
            return b
    raise ValueError(f"sequence of {n} tokens > {BUCKETS[-1]}")


def pack(seqs, B, T):
    x = np.full((B, T), PAD, dtype=np.int32)
    for i, s in enumerate(seqs):
        x[i, :len(s)] = s
    return x


def sums(lp, seqs):
    return [float(np.asarray(lp[i, :len(s) - 1], dtype=np.float64).sum()) for i, s in enumerate(seqs)]


def load_cache():
    return [json.loads(l) for l in open(CACHE, encoding="utf-8")]