"""Write task files from corpus.jsonl, the photo manifest and downloaded images (run on the server).

  splits/mt_{train,val,test}.jsonl        transliteration -> translation
  splits/ocr_real_{train,val,test}.jsonl  photo -> transliteration (test/val: first photo per record)
  splits/ocr_syn_{train,val}.jsonl        synthetic photo -> transliteration
"""
import json
import os
import re
import sys
from collections import Counter

sys.path.insert(0, "/home/jovyan/humonen/runic_access/code")
from runic_norm import clean_translation  # noqa: E402

OCR_ONLY = "--ocr_only" in sys.argv

R = "/home/jovyan/humonen/runic_access"
OUT = f"{R}/splits"
os.makedirs(OUT, exist_ok=True)

corpus = [json.loads(l) for l in open(f"{R}/data/corpus.jsonl")]
by_id = {r["id"]: r for r in corpus}

counts = Counter()
if not OCR_ONLY:
  with open(f"{OUT}/mt_train.jsonl", "w") as ftr, open(f"{OUT}/mt_val.jsonl", "w") as fva, \
          open(f"{OUT}/mt_test.jsonl", "w") as fte:
      files = {"train": ftr, "val": fva, "test": fte}
      for r in corpus:
          tgt = clean_translation(r["translation"])
          if len(re.sub(r"[^\w]", "", tgt)) < 3 or (r["split"] == "train" and len(tgt) > 1000):
              counts[("mt_dropped", r["split"])] += 1
              continue
          rec = dict(id=r["id"], src=r["norm"], tgt=tgt, lang=r["lang"], cluster=r["cluster"],
                     source=r["src"], signum=r.get("signum"), max_ratio_to_train=r.get("max_ratio_to_train"))
          files[r["split"]].write(json.dumps(rec, ensure_ascii=False) + "\n")
          counts[("mt", r["split"], r["lang"])] += 1

manifest = [json.loads(l) for l in open(f"{R}/data/image_manifest.jsonl")]
with open(f"{OUT}/ocr_real_train.jsonl", "w") as ftr, open(f"{OUT}/ocr_real_val.jsonl", "w") as fva, \
        open(f"{OUT}/ocr_real_test.jsonl", "w") as fte:
    files = {"train": ftr, "val": fva, "test": fte}
    for m in manifest:
        if not os.path.exists(f"{R}/images/{m['file']}"):
            counts[("ocr_missing", m["split"])] += 1
            continue
        if m["split"] != "train" and m["k"] != 0:
            continue
        r = by_id[m["row_id"]]
        if m["split"] == "train" and len(r["norm"]) > 600:
            continue
        rec = dict(id=m["file"].rsplit(".", 1)[0], image=f"images/{m['file']}", target=r["norm"],
                   row_id=r["id"], lang=r["lang"], tgt=clean_translation(r["translation"]), source=r["src"], cluster=r["cluster"],
                   max_ratio_to_train=r.get("max_ratio_to_train"))
        files[m["split"]].write(json.dumps(rec, ensure_ascii=False) + "\n")
        counts[("ocr", m["split"], r["src"])] += 1

for name in ("train", "val"):
    with open(f"{R}/synth/syn_{name}.jsonl") as f, open(f"{OUT}/ocr_syn_{name}.jsonl", "w") as g:
        for l in f:
            r = json.loads(l)
            r["image"] = "synth/" + r["image"]
            g.write(json.dumps(r, ensure_ascii=False) + "\n")
            counts[("syn", name)] += 1

for k, v in sorted(counts.items(), key=str):
    print(k, v)