"""Retrieval baseline: output the translation of the most similar training transliteration
(Indel similarity on space-free strings, same target language)."""
import json
import sys

from rapidfuzz import fuzz, process

train = [json.loads(l) for l in open(sys.argv[1])]
test = [json.loads(l) for l in open(sys.argv[2])]
out = open(sys.argv[3], "w")
by_lang = {}
for r in train:
    by_lang.setdefault(r["lang"], []).append(r)
keys = {lang: [r["src"].replace(" ", "") for r in rows] for lang, rows in by_lang.items()}
for r in test:
    q = (r.get(sys.argv[4]) if len(sys.argv) > 4 else r["src"]) or ""
    match = process.extractOne(q.replace(" ", ""), keys[r["lang"]], scorer=fuzz.ratio)
    rec = dict(r)
    rec["pred"] = by_lang[r["lang"]][match[2]]["tgt"] if match else ""
    rec["retrieval_score"] = match[1] if match else 0
    out.write(json.dumps(rec, ensure_ascii=False) + "\n")
out.close()