"""Findings for the write-up: memorisation vs generalisation, and context from
production tokenizers. Writes findings.json."""
import json, pathlib, random
from fractions import Fraction
from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
import tiktoken
from eval_common import units, visible
from train_tokenizer import BYTES

HERE = pathlib.Path(__file__).resolve().parent
LANGS = ["en", "hi", "te", "mai"]          # the evaluation languages (India pages)
HELD = ["en", "hi", "te"]                   # held-out articles exist for these (none for Maithili)
TRAIN_LANGS = ["en", "hi", "te", "mr"]      # held-out training text; Marathi stands in for Devanagari
INDIA = {l: (HERE / "corpus" / f"{l}.faithful.txt").read_text(encoding="utf-8") for l in LANGS}
idx = json.loads((HERE / "heldout" / "index.json").read_text())
topics = sorted(idx)
random.Random(0).shuffle(topics)
TRAIN, TEST = topics[:18], topics[18:]
def held(topic_list, l):
    return "\n".join((HERE / "heldout" / f"{l}__{t.replace(' ', '_')}.md").read_text(encoding="utf-8") for t in topic_list)
TEST_TEXT = {l: held(TEST, l) for l in HELD}

def hf_count(tok, text):
    return len(tok.encode(text).ids)            # whole text in one call, like the course's evaluator

def summary(counter):  # counter(lang, text) -> tokens
    out = {}
    for name, texts, langs in (("india_pages", INDIA, LANGS), ("heldout_test", TEST_TEXT, HELD)):
        f = {l: counter(l, texts[l]) / units(texts[l]) for l in langs}
        sp = max(f.values()) - min(f.values())
        out[name] = {"fertility": {l: round(v, 4) for l, v in f.items()}, "spread": round(sp, 4), "score": round(1000 / sp, 1)}
    return out

res = {"heldout_topics_train": TRAIN, "heldout_topics_test": TEST}

final = Tokenizer.from_file(str(HERE / "tokenizer.json"))
res["submission"] = summary(lambda l, t: hf_count(final, t))

base = Tokenizer.from_file(str(HERE / "tokenizer.base.json"))
res["before_equalising"] = summary(lambda l, t: hf_count(base, t))

# same design, trained only on held-out topics (never sees the India pages)
def train_generalist(weights):
    tok = Tokenizer(models.BPE(byte_fallback=True))
    tok.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True)
    tok.decoder = decoders.Sequence([decoders.Metaspace(replacement="▁", prepend_scheme="always", split=True), decoders.ByteFallback(), decoders.Fuse()])
    tr = trainers.BpeTrainer(vocab_size=10000 - 256, min_frequency=2, show_progress=False, limit_alphabet=1000)
    data = {l: held(TRAIN, l) for l in TRAIN_LANGS}
    tok.train_from_iterator((data[l] for l in TRAIN_LANGS for _ in range(weights[l])), trainer=tr)
    spec = json.loads(tok.to_str()); v = spec["model"]["vocab"]
    for b in BYTES:
        v.setdefault(b, len(v))
    return Tokenizer.from_str(json.dumps(spec, ensure_ascii=False))

best = None
for w in [(1, 1, 1, 1), (1, 2, 4, 3), (1, 2, 6, 4), (1, 3, 8, 5), (1, 3, 10, 6), (1, 4, 12, 8)]:
    g = train_generalist(dict(zip(TRAIN_LANGS, w)))
    s = summary(lambda l, t: hf_count(g, t))
    s["weights"] = dict(zip(TRAIN_LANGS, w))
    print("generalist", w, s["india_pages"])
    if best is None or s["india_pages"]["spread"] < best["india_pages"]["spread"]:
        best = s
res["generalist_same_design"] = best

for name in ("cl100k_base", "o200k_base"):
    enc = tiktoken.get_encoding(name)
    res[name] = summary(lambda l, t: len(enc.encode(t, disallowed_special=())))
    res[name]["vocab"] = enc.n_vocab
    print(name, res[name]["india_pages"])

(HERE / "findings.json").write_text(json.dumps(res, ensure_ascii=False, indent=1), encoding="utf-8")
print(json.dumps({k: v for k, v in res.items() if k.startswith(("submission", "before"))}, indent=1))