"""Image-blind baselines: the single training transliteration (translation) that minimizes the
average error (maximizes chrF++) on a sample of training items, applied to every test item."""
import json
import random
import sys
from collections import Counter

import sacrebleu
from rapidfuzz.distance import Levenshtein

ROOT = "/Users/humonen/runic_access"
rng = random.Random(0)
train = [json.loads(l) for l in open(f"{ROOT}/splits_local/mt_train.jsonl")]
out = {}

# recognition: medoid transliteration under CER
cands = [s for s, _ in Counter(r["src"] for r in train if len(r["src"]) >= 5).most_common(300)]
sample = rng.sample([r["src"] for r in train], 1500)
best = min(cands, key=lambda c: sum(Levenshtein.distance(c, s) for s in sample) / sum(len(s) for s in sample))
out["ocr_constant"] = best

# translation: per-language medoid translation under chrF++
chrf = sacrebleu.CHRF(word_order=2)
for lang in ("en", "sv"):
    tr = [r["tgt"] for r in train if r["lang"] == lang]
    cands = [t for t, _ in Counter(tr).most_common(200)]
    cands += rng.sample(tr, 100)
    sample = rng.sample(tr, 400)
    best = max(cands, key=lambda c: chrf.corpus_score([c] * len(sample), [sample]).score)
    out[f"mt_constant_{lang}"] = best
json.dump(out, open(f"{ROOT}/results/prior_baselines.json", "w"), ensure_ascii=False, indent=1)
print(json.dumps(out, ensure_ascii=False, indent=1))