import json, pathlib
from train_tokenizer import build, LANGS
from evaluate_tokenizer import count_tokens
from eval_common import units
texts = {l: (pathlib.Path("corpus")/f"{l}.faithful.txt").read_text() for l in LANGS}
U = {l: units(t) for l, t in texts.items()}
cache = {}
def spread(w):
    if w in cache: return cache[w]
    tok = build(dict(zip(LANGS, w)))
    f = {l: count_tokens(tok, texts[l])/U[l] for l in LANGS}
    cache[w] = (max(f.values())-min(f.values()), f)
    return cache[w]
w = [10, 11, 20, 14]
best = spread(tuple(w))
improved = True
while improved:
    improved = False
    for i in (1, 2, 3):
        for d in (-3, -1, 1, 3):
            c = list(w); c[i] += d
            if c[i] < 1: continue
            s = spread(tuple(c))
            if s[0] < best[0]:
                best, w, improved = s, c, True
                print(w, f"spread={s[0]:.5f}", ' '.join(f"{l}={s[1][l]:.4f}" for l in LANGS), flush=True)
print("BEST", w, best[0], json.dumps({l: round(v, 5) for l, v in best[1].items()}))