"""Write LaTeX tables for the results section from results/results.json (no hand-copied numbers)."""
import json
import os

ROOT = os.path.expanduser("~/runic_access")
R = json.load(open(f"{ROOT}/results/" + os.environ.get("RESULTS", "results.json")))
OUT = f"{ROOT}/paper/tables"
os.makedirs(OUT, exist_ok=True)


def ci(v, c, d=1):
    return f"{v:.{d}f} {{\\scriptsize[{c[0]:.{d}f}, {c[1]:.{d}f}]}}"


# ---------------- E1 ----------------
e1 = R.get("E1", {})
rows = [("Constant string", "--", "ocr_test_constant"), ("Zero-shot", "0.8B", "ocr_test_zs08"), ("Zero-shot", "4B", "ocr_test_zs4b"),
        ("Procedural", "4B", "ocr_test_ocr4b_syn"), ("Diffusion", "4B", "ocr_test_ocr4b_diff"),
        ("Real", "4B", "ocr_test_ocr4b_real"), ("Real + procedural", "4B", "ocr_test_ocr4b_mix"),
        ("Real + diffusion", "4B", "ocr_test_ocr4b_diffmix"), ("Real + proc., 1\\,MP", "4B", "ocr_test_ocr4b_mix_hr"),
        ("Real crops + proc.", "4B", "ocr_test_ocr4b_crop"), ("Real + proc., on crops", "4B", "ocr_test_ocr4b_mix_oncrop"),
        ("Real + procedural", "0.8B", "ocr_test_ocr08_mix")]


def cell(v, key):
    return f"{v[key]['cer']:.1f}" if key in v else "--"


n = {k: e1.get("ocr_test_ocr4b_mix", {}).get(k, {}).get("n", 0) for k in ("all", "gr", "rdb_photo", "card")}
lines = [r"\begin{table*}[t]", r"\centering", r"\color{blue}",
         r"\caption{Recognition across image domains: CER (\%%, lower is better) on held-out synthetic images of each style and on the real test images by domain (GR: Gamla runor field photographs; RuneS: photographs of objects; Cards: scanned archive cards). \emph{All}: the %d real test images with a 95\%% bootstrap interval. Rare: recall of words outside the 150 most frequent training words.}" % n["all"],
         r"\label{tab:ocr}", r"\small", r"\begin{tabular}{llrrrrrrr}", r"\toprule",
         r" & & \multicolumn{2}{c}{Synthetic} & \multicolumn{3}{c}{Real test} & & \\",
         r"\cmidrule(lr){3-4}\cmidrule(lr){5-7}",
         r"Training data & Model & Proc. & Diff. & GR (%d) & RuneS (%d) & Cards (%d) & All & Rare \\" % (n["gr"], n["rdb_photo"], n["card"]),
         r"\midrule"]
for name, size, tag in rows:
    if tag not in e1:
        continue
    v = e1[tag]
    a = v["all"]
    rare = f"{v['rare_recall']:.1f}" if "rare_recall" in v else "--"
    lines.append(f"{name} & {size} & {cell(v, 'proc_syn')} & {cell(v, 'diff_syn')} & {cell(v, 'gr')} & {cell(v, 'rdb_photo')} & {cell(v, 'card')} & {ci(a['cer'], a['cer_ci'])} & {rare} \\\\")
lines += [r"\bottomrule", r"\end{tabular}", r"\end{table*}"]
open(f"{OUT}/ocr.tex", "w").write("\n".join(lines) + "\n")

# ---------------- E2 ----------------
e2 = R.get("E2", {})
rows = [("Zero-shot", "0.8B", "mt_test_zs08"), ("Zero-shot", "4B", "mt_test_zs4b"), ("Retrieval", "--", "mt_test_retrieval"),
        ("LoRA", "0.8B", "mt_test_mt08"), ("LoRA", "4B", "mt_test_mt4b"), ("LoRA + noise", "4B", "mt_test_mt4b_noise")]
lines = [r"\begin{table*}[t]", r"\centering", r"\color{blue}",
         r"\caption{Translation of reference transliterations on the held-out test set (chrF++ and BLEU, higher is better; 95\% bootstrap intervals for chrF++). \emph{Strict}: test pairs whose closest training transliteration is below 80\% similarity.}",
         r"\label{tab:mt}", r"\begin{tabular}{llrrrrrr}", r"\toprule",
         r" & & \multicolumn{2}{c}{English ($n{=}%d$)} & \multicolumn{2}{c}{Swedish ($n{=}%d$)} & \multicolumn{2}{c}{Strict ($n{=}%d$)} \\" % (
             e2.get("mt_test_retrieval", {}).get("en", {}).get("n", 0), e2.get("mt_test_retrieval", {}).get("sv", {}).get("n", 0),
             e2.get("mt_test_retrieval", {}).get("strict", {}).get("n", 0)),
         r"\cmidrule(lr){3-4}\cmidrule(lr){5-6}\cmidrule(lr){7-8}",
         r"System & Model & chrF++ & BLEU & chrF++ & BLEU & chrF++ & BLEU \\", r"\midrule"]
for name, size, tag in rows:
    if tag not in e2:
        continue
    v = e2[tag]
    cells = []
    for k in ("en", "sv", "strict"):
        cells += [ci(v[k]["chrfpp"], v[k]["chrfpp_ci"]), f"{v[k]['bleu']:.1f}"]
    lines.append(f"{name} & {size} & " + " & ".join(cells) + r" \\")
lines += [r"\bottomrule", r"\end{tabular}", r"\end{table*}"]
open(f"{OUT}/mt.tex", "w").write("\n".join(lines) + "\n")

# ---------------- E3 ----------------
e3 = R.get("E3", {})
if e3.get("subsets"):
    tag = next(k[4:].rsplit("_mt4b", 1)[0] for k in e3 if k.startswith("e2e_") and "_mt4b" in k)
    sub = e3["subsets"]
    lines = [r"\begin{table}[t]", r"\centering", r"\color{blue}",
             r"\caption{End-to-end translation of the test photographs (chrF++). \emph{Reference}: the 4B translator applied to the reference transliterations; the other columns translate the recognizer output. Rows group the photographs by the recognizer's CER on each image and by image domain.}",
             r"\label{tab:e2e}", r"\small", r"\setlength{\tabcolsep}{3pt}", r"\begin{tabular}{lrrrrrrr}", r"\toprule",
             r" & & & & \multicolumn{4}{c}{From photograph} \\", r"\cmidrule(lr){5-8}",
             r"Photographs & $n$ & CER & Ref. & Const. & Retr. & LoRA & +noise \\", r"\midrule"]
    a = e3["ocr_cer"]["all"]
    lines.append(f"All & {a['n']} & {a['cer']:.1f} & {e3[f'e2e_{tag}_mt4b_fromgold']['all']['chrfpp']:.1f} & "
                 f"{e3[f'e2e_{tag}_constant']['all']['chrfpp']:.1f} & {e3[f'e2e_{tag}_retrieval_fromocr']['all']['chrfpp']:.1f} & "
                 f"{e3[f'e2e_{tag}_mt4b_fromocr']['all']['chrfpp']:.1f} & {e3[f'e2e_{tag}_mt4b_noise_fromocr']['all']['chrfpp']:.1f} \\\\")
    names = (("cer<25", r"CER $<$25\%"), ("cer25-50", r"CER 25--50\%"), ("cer50-75", r"CER 50--75\%"), ("cer>=75", r"CER $\geq$75\%"),
             ("gr", "Gamla runor"), ("rdb_photo", "RuneS photos"), ("card", "RuneS cards"))
    for i, (k, name) in enumerate(names):
        if i in (0, 4):
            lines.append(r"\midrule")
        if k not in sub:
            continue
        v = sub[k]
        lines.append(f"{name} & {v['n']} & {v['ocr']['cer']:.1f} & {v['mt4b_fromgold']['chrfpp']:.1f} & {v['constant']['chrfpp']:.1f} & "
                     f"{v['retrieval_fromocr']['chrfpp']:.1f} & {v['mt4b_fromocr']['chrfpp']:.1f} & {v['mt4b_noise_fromocr']['chrfpp']:.1f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\end{table}"]
    open(f"{OUT}/e2e.tex", "w").write("\n".join(lines) + "\n")
# ---------------- appendix: raw recognition and output length ----------------
lines = [r"\begin{table}[h]", r"\centering", r"\color{blue}",
         r"\caption{Recognition on the real test images: corpus-level CER (\%) without and with repetition truncation, median CER of individual images, ratio of output to reference length, and share of outputs more than twice as long as the reference (after truncation).}",
         r"\label{tab:ocr-raw}", r"\footnotesize", r"\setlength{\tabcolsep}{2.5pt}",
         r"\begin{tabular}{llrrrrr}", r"\toprule", r"Training data & Model & Raw & Cut & Median & Length & $>$2$\times$ \\", r"\midrule"]
for name, size, tag in [("Zero-shot", "0.8B", "ocr_test_zs08"), ("Zero-shot", "4B", "ocr_test_zs4b"),
        ("Procedural", "4B", "ocr_test_ocr4b_syn"), ("Diffusion", "4B", "ocr_test_ocr4b_diff"),
        ("Real", "4B", "ocr_test_ocr4b_real"), ("Real + procedural", "4B", "ocr_test_ocr4b_mix"),
        ("Real + diffusion", "4B", "ocr_test_ocr4b_diffmix"), ("Real + proc., 1\\,MP", "4B", "ocr_test_ocr4b_mix_hr"),
        ("Real crops + proc.", "4B", "ocr_test_ocr4b_crop"), ("Real + proc., on crops", "4B", "ocr_test_ocr4b_mix_oncrop"),
        ("Real + procedural", "0.8B", "ocr_test_ocr08_mix")]:
    if tag in e1 and "raw" in e1[tag]:
        v = e1[tag]
        lines.append(f"{name} & {size} & {v['raw']['cer']:.1f} & {v['all']['cer']:.1f} & {v['median_cer']:.1f} & "
                     f"{v['len_ratio']:.2f} & {v['over2x']:.1f} \\\\")
lines += [r"\bottomrule", r"\end{tabular}", r"\end{table}"]
open(f"{OUT}/ocr_raw.tex", "w").write("\n".join(lines) + "\n")
# ---------------- appendix: reference-free selection of photographs ----------------
sel = R.get("E3", {}).get("selection", {})
names = [(k, n) for k, n in (("nll", "NLL"), ("agree_hr", "Agree-2"), ("agree4", "Agree-5"), ("oracle_cer", "Oracle")) if k in sel]
if names:
    e3a = R["E3"]
    tag = next(k[4:].rsplit("_mt4b", 1)[0] for k in e3a if k.startswith("e2e_") and "_mt4b" in k)
    lines = [r"\begin{table}[!ht]", r"\centering", r"\color{blue}",
             r"\caption{End-to-end translation of the most confidently read test photographs. Photographs are ranked by a confidence signal and the top fraction is kept; each cell gives the recognizer's CER (\%) and chrF++ of the translation from the photograph. NLL: mean per-token negative log-likelihood that the pipeline recognizer assigns to its own output; empty outputs and decodes stopped by the generation limit are ranked last. Agree-2: normalized edit distance between the outputs of the pipeline recognizer and of a recognizer trained on whole images at 1.05\,MP. Agree-5: mean distance to four other adapted recognizers (real only, real + diffusion, 1.05\,MP and the 0.8B model, all on whole images). Oracle: ranking by the true CER.}",
             r"\label{tab:selection}", r"\footnotesize", r"\setlength{\tabcolsep}{2.5pt}",
             r"\begin{tabular}{l" + "rr" * len(names) + "}", r"\toprule",
             " & " + " & ".join(r"\multicolumn{2}{c}{%s}" % n for _, n in names) + r" \\",
             "".join(r"\cmidrule(lr){%d-%d}" % (2 + 2 * j, 3 + 2 * j) for j in range(len(names))),
             "Kept & " + " & ".join("CER & chrF++" for _ in names) + r" \\", r"\midrule"]
    for frac in ("0.1", "0.2", "0.3", "0.5"):
        n = sel[names[0][0]][frac]["n"]
        lines.append(f"{round(100 * float(frac))}\\% ({n}) & " + " & ".join(
            f"{sel[k][frac]['cer']:.1f} & {sel[k][frac]['chrfpp_ocr']:.1f}" for k, _ in names) + r" \\")
    a_ = e3a["ocr_cer"]["all"]["cer"]
    c_ = e3a[f"e2e_{tag}_mt4b_fromocr"]["all"]["chrfpp"]
    lines.append(f"100\\% ({e3a['ocr_cer']['all']['n']}) & " + " & ".join(f"{a_:.1f} & {c_:.1f}" for _ in names) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\end{table}"]
    open(f"{OUT}/selection.tex", "w").write("\n".join(lines) + "\n")
# ---------------- appendix: archive cards by what the crop shows ----------------
cards = R.get("E3", {}).get("cards", {})
if cards:
    lines = [r"\begin{table}[!ht]", r"\centering", r"\color{blue}",
             r"\caption{Archive cards of the test set by what the located crop shows. Whole: the recognizer trained on whole images reads the whole card; Crop: the pipeline recognizer reads the crop, or the whole card when no box was found. CER (\%) of recognition, chrF++ of translation from the recognizer output and from the reference transliterations (Ref.).}",
             r"\label{tab:cards}", r"\small", r"\setlength{\tabcolsep}{3pt}", r"\begin{tabular}{lrrrrrr}", r"\toprule",
             r" & & \multicolumn{2}{c}{CER} & \multicolumn{3}{c}{chrF++} \\", r"\cmidrule(lr){3-4}\cmidrule(lr){5-7}",
             r"Crop shows & $n$ & Whole & Crop & Whole & Crop & Ref. \\", r"\midrule"]
    for k, name in (("photo", "Photograph"), ("drawing", "Drawing"), ("whole", "No box"), (None, None), ("any", "All cards")):
        if k is None:
            lines.append(r"\midrule")
        elif k in cards:
            v = cards[k]
            lines.append(f"{name} & {v['n']} & {v['cer_whole']:.1f} & {v['cer']:.1f} & {v['chrfpp_whole']:.1f} & "
                         f"{v['chrfpp']:.1f} & {v['chrfpp_gold']:.1f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\end{table}"]
    open(f"{OUT}/cards.tex", "w").write("\n".join(lines) + "\n")
print("tables written to", OUT)