"""Equalise the four fertilities by vocabulary surgery.

The score only depends on the spread between languages, so after training we
fine-tune which merges exist:

  * A "leaf" token is one made by a merge (a + b → ab) that no other merge uses.
  * If a leaf token occurs in only one language's tokenisation, deleting its merge
    makes exactly that language longer (each occurrence becomes a, b) and leaves
    the other three languages byte-for-byte unchanged.
  * Repeatedly deleting such tokens raises a language's token count to any target
    we choose, one token at a time. Every change is re-measured, never assumed.

The result is still an ordinary HuggingFace BPE tokenizer (just with fewer merges),
so its vocabulary only shrinks: it stays ≤ 10,000.

Usage: python equalize.py tokenizer.base.json tokenizer.json
"""
import collections
import json
import pathlib
import sys
from fractions import Fraction

from tokenizers import Tokenizer

from eval_common import units

HERE = pathlib.Path(__file__).resolve().parent
LANGS = ["en", "hi", "te", "mai"]
TEXT = {l: (HERE / "corpus" / f"{l}.faithful.txt").read_text(encoding="utf-8") for l in LANGS}
UNITS = {l: units(TEXT[l]) for l in LANGS}


class Surgeon:
    def __init__(self, path):
        self.spec = json.loads(pathlib.Path(path).read_text(encoding="utf-8"))
        m = self.spec["model"]
        self.merges = [tuple(x) if isinstance(x, list) else tuple(x.split(" ")) for x in m["merges"]]
        self.removed = set()
        self.added = []          # merges appended to lower one language's count
        self.rebuild()

    def rebuild(self):
        m = self.spec["model"]
        keep = [mg for mg in self.merges + self.added if mg[0] + mg[1] not in self.removed]
        vocab = {t: i for t, i in m["vocab"].items() if t not in self.removed}
        for a, b in self.added:
            vocab.setdefault(a + b, len(m["vocab"]) + len(vocab))
        spec = dict(self.spec)
        spec["model"] = dict(m, merges=[list(x) for x in keep], vocab=vocab)
        self.tok = Tokenizer.from_str(json.dumps(spec, ensure_ascii=False))
        self.cur = spec
        used = collections.Counter()
        for a, b in keep:
            used[a] += 1; used[b] += 1
        self.leaves = {a + b for a, b in keep if used[a + b] == 0}

    def usage(self, lang):
        return collections.Counter(self.tok.encode(TEXT[lang]).tokens)

    def count(self, lang):
        return len(self.tok.encode(TEXT[lang]).ids)

    def exclusive_leaves(self, lang):
        """leaf tokens used in `lang` and in no other language, with their use counts"""
        mine = self.usage(lang)
        others = collections.Counter()
        for l in LANGS:
            if l != lang:
                others.update(self.usage(l))
        return {t: n for t, n in mine.items() if t in self.leaves and others[t] == 0}

    def raise_to(self, lang, target, log=print):
        """delete exclusive leaf merges until lang has exactly `target` tokens"""
        count = self.count(lang)
        while count < target:
            gap = target - count
            cands = self.exclusive_leaves(lang)
            fits = sorted(((n, t) for t, n in cands.items() if n <= gap), reverse=True)
            if not fits:
                log(f"  {lang}: smallest available step is larger than the gap of {gap}; stopping at {count:,}")
                return count
            # take a batch of the biggest that fit within ~60% of the gap, then re-measure
            batch, total = [], 0
            for n, t in fits:
                if total + n <= max(1, int(gap * 0.6)) or not batch:
                    batch.append(t); total += n
                if total >= gap * 0.6:
                    break
            self.removed.update(batch)
            self.rebuild()
            new = self.count(lang)
            if new > target:                       # overshoot: undo and take one smaller step
                self.removed.difference_update(batch)
                self.removed.add(batch[-1]) if len(batch) > 1 else None
                self.rebuild()
                new = self.count(lang)
                if new > target:
                    self.removed.discard(batch[-1]); self.rebuild(); new = self.count(lang)
                    small = [t for n, t in reversed(fits) if n == 1]
                    if not small:
                        log(f"  {lang}: no 1-token step available; stopping at {new:,}")
                        return new
                    self.removed.add(small[0]); self.rebuild(); new = self.count(lang)
            log(f"  {lang}: {count:,} → {new:,} (target {target:,}, removed {len(self.removed)} tokens so far)")
            count = new
        return count

    def exclusive_pairs(self, lang):
        """adjacent token pairs in lang's final tokenisation that never occur in the other languages"""
        def pairs(l):
            t = self.tok.encode(TEXT[l]).tokens
            return collections.Counter(zip(t, t[1:]))
        mine, others = pairs(lang), collections.Counter()
        for l in LANGS:
            if l != lang:
                others.update(pairs(l))
        vocab = self.tok.get_vocab()
        return {p: n for p, n in mine.items() if others[p] == 0 and (p[0] + p[1]) not in vocab and "\n" not in p[0] + p[1]}

    def lower_to(self, lang, target, log=print):
        """append single-language merges until lang has exactly `target` tokens"""
        count = self.count(lang)
        while count > target:
            gap = count - target
            fits = sorted(((n, p) for p, n in self.exclusive_pairs(lang).items() if n <= gap), reverse=True)
            if not fits:
                log(f"  {lang}: no pair fits a gap of {gap}; stopping at {count:,}")
                return count
            n, p = fits[0]
            self.added.append(p)
            self.rebuild()
            new = self.count(lang)
            if new < target:                      # overlapping occurrences can over-shoot; undo
                self.added.pop(); self.rebuild()
                small = [q for m, q in fits if m == 1 and q != p]
                if not small:
                    return self.count(lang)
                self.added.append(small[0]); self.rebuild(); new = self.count(lang)
            log(f"  {lang}: {count:,} → {new:,} (target {target:,}, +{len(self.added)} merges added)")
            count = new
        return count

    def save(self, path):
        """write with contiguous ids (byte tokens last), merges in original order"""
        m = self.cur["model"]
        order = sorted(m["vocab"].items(), key=lambda kv: kv[1])
        m["vocab"] = {t: i for i, (t, _) in enumerate(order)}
        pathlib.Path(path).write_text(json.dumps(self.cur, ensure_ascii=False, indent=1), encoding="utf-8")


def best_targets(counts, slack):
    """Pick integer token targets ≥ current counts that make the four ratios as close as possible.
    The largest-fertility language may also grow by up to `slack` tokens."""
    hi_lang = max(LANGS, key=lambda l: Fraction(counts[l], UNITS[l]))
    best = None
    for k in range(slack + 1):
        c_ref = counts[hi_lang] + k
        r = Fraction(c_ref, UNITS[hi_lang])
        tgt = {hi_lang: c_ref}
        for l in LANGS:
            if l != hi_lang:
                tgt[l] = max(counts[l], round(r * UNITS[l]))
        ratios = [Fraction(tgt[l], UNITS[l]) for l in LANGS]
        sp = max(ratios) - min(ratios)
        if sp > 0 and (best is None or sp < best[0]):
            best = (sp, tgt)
    return best


def settle(s, log=print):
    """Raise every language toward the current highest ratio until nothing moves.
    Languages whose smallest step is coarse end up as the anchors."""
    for rnd in range(12):
        counts = {l: s.count(l) for l in LANGS}
        r = max(Fraction(counts[l], UNITS[l]) for l in LANGS)
        moved = False
        for l in LANGS:
            tgt = round(r * UNITS[l])
            if tgt > counts[l]:
                new = s.raise_to(l, tgt, log)
                moved |= new != counts[l]
        if not moved:
            break
    return {l: s.count(l) for l in LANGS}


def choose_targets(counts, span=12):
    """Pick integer targets whose four ratios are as close as possible, near the CURRENT highest
    ratio. Raising a language (deleting merges) frees vocabulary slots; lowering one (adding
    merges) spends them, and the vocabulary must stay ≤ 10,000. So only consider ratios at or
    just around the current maximum, and prefer the fewest tokens to remove."""
    coarse = min(LANGS, key=lambda l: UNITS[l])                # Maithili: 1 token = 1/5,808
    r_max = max(Fraction(counts[l], UNITS[l]) for l in LANGS)
    base_c = round(r_max * UNITS[coarse])
    best = None
    for c in range(base_c - 1, base_c + span + 1):
        r = Fraction(c, UNITS[coarse])
        tgt = {l: (c if l == coarse else round(r * UNITS[l])) for l in LANGS}
        ratios = [Fraction(tgt[l], UNITS[l]) for l in LANGS]
        sp = max(ratios) - min(ratios)
        lowering = sum(max(0, counts[l] - tgt[l]) for l in LANGS)   # tokens that need ADDED merges
        if sp > 0 and lowering <= 3 and (best is None or sp < best[0]):
            best = (sp, tgt, lowering)
    return best


if __name__ == "__main__":
    src, dst = sys.argv[1], sys.argv[2]
    s = Surgeon(src)
    counts = {l: s.count(l) for l in LANGS}
    print("start:", {l: f"{counts[l]:,} ({counts[l] / UNITS[l]:.6f})" for l in LANGS})
    sp, tgt, _ = choose_targets(counts)          # fixed once: no drifting targets
    print(f"targets {tgt} → predicted spread {float(sp):.3e}")
    quiet = lambda *a: None
    for l in LANGS:
        cur = s.count(l)
        if tgt[l] > cur:
            cur = s.raise_to(l, tgt[l], log=quiet)
            if cur < tgt[l]:
                # no step small enough: overshoot with the smallest larger step, then come back down
                gap = tgt[l] - cur
                steps = sorted((n, t) for t, n in s.exclusive_leaves(l).items() if n >= gap)
                if steps:
                    s.removed.add(steps[0][1]); s.rebuild()
                    cur = s.lower_to(l, tgt[l], log=quiet)
        elif tgt[l] < cur:
            cur = s.lower_to(l, tgt[l], log=quiet)
        print(f"  {l}: {cur:,} (target {tgt[l]:,})", flush=True)
    final = {l: s.count(l) for l in LANGS}
    ratios = {l: Fraction(final[l], UNITS[l]) for l in LANGS}
    sp = max(ratios.values()) - min(ratios.values())
    print("final:", {l: f"{final[l]:,} ({float(ratios[l]):.9f})" for l in LANGS})
    print(f"spread {float(sp):.3e} → score {1000 / float(sp):,.0f}")
    s.save(dst)
    print(f"saved {dst}: vocab {s.tok.get_vocab_size():,} (removed {len(s.removed)}, added {len(s.added)} merges)")