"""E4 figure. (a) chrF++ of both 4B translators against the actual CER of synthetically corrupted
transliterations, with the end-to-end operating point of the recognizer on all test photographs.
(b) The same curves relative to uncorrupted input, and the real recognizer output grouped by the CER of each
photograph, relative to the translation of the reference transliterations of the same photographs.
Recognizer output is drawn for the clean translator; the noise-trained one differs by less than one point
(printed below for the caption). The last CER group (mean CER far above 100%) is drawn at the axis edge."""
import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

ROOT = os.path.expanduser("~/runic_access")
R = json.load(open(f"{ROOT}/results/" + os.environ.get("RESULTS", "results.json")))
e4, e3 = R["E4"], R["E3"]
tag = sys.argv[1] if len(sys.argv) > 1 else "ocr4b_mixcut"
act = {float(k): v for k, v in e4["actual_cer"].items()}
TR = (("mt4b", "LoRA 4B", "#2a78d6", "o"), ("mt4b_noise", "LoRA 4B + noise", "#eb6834", "s"))
REC = "#0b0b0b"                      # recognizer output: its own identity, not a translator colour
INK, MUTED, GRID = "#0b0b0b", "#52514e", "#e6e5e0"
XMAX = 100

plt.rcParams.update({"font.size": 7.5, "axes.edgecolor": MUTED, "axes.labelcolor": INK, "xtick.color": MUTED,
                     "ytick.color": MUTED, "axes.spines.top": False, "axes.spines.right": False,
                     "pdf.fonttype": 42})
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(3.45, 4.1), sharex=True)
for tr, label, color, marker in TR:
    lv = sorted(e4[tr], key=float)
    xs = [act[float(l)] for l in lv]
    ys = [e4[tr][l]["chrfpp"] for l in lv]
    ax1.plot(xs, ys, color=color, marker=marker, ms=4, lw=1.8, zorder=3)
    ax2.plot(xs, [100 * y / ys[0] for y in ys], color=color, marker=marker, ms=4, lw=1.8, zorder=3)

# recognizer output grouped by the CER of each photograph: (a) chrF++, (b) relative to the reference input
sub = e3["subsets"]
bins = [b for b in ("cer<25", "cer25-50", "cer50-75", "cer>=75") if b in sub]
bx = [sub[b]["ocr"]["cer"] for b in bins]
by_abs = [sub[b]["mt4b_fromocr"]["chrfpp"] for b in bins]
by = [100 * sub[b]["mt4b_fromocr"]["chrfpp"] / sub[b]["mt4b_fromgold"]["chrfpp"] for b in bins]


def groups(ax, ys, label_off):
    inside = [(x, y) for x, y in zip(bx, ys) if x <= XMAX]
    off = [(x, y) for x, y in zip(bx, ys) if x > XMAX]
    ax.plot(*zip(*inside), color=REC, ls=(0, (3, 2)), lw=1.2, marker="o", ms=6, mfc="white", mec=REC, mew=1.5,
            zorder=4)
    for x, y in off:
        xe = XMAX - 2
        ax.plot([inside[-1][0], xe], [inside[-1][1], y], color=REC, ls=(0, (1, 2)), lw=1.0, zorder=3)
        ax.plot([xe], [y], ls="none", marker=">", ms=6, mfc="white", mec=REC, mew=1.5, zorder=4, clip_on=False)
        if label_off:
            ax.annotate(f"group with mean\nCER {x:.0f}%", (xe, y), xytext=(2, -8), textcoords="offset points",
                        ha="right", va="top", fontsize=6.5, color=MUTED)


groups(ax1, by_abs, False)
groups(ax2, by, True)
# (a) the whole pipeline over all test photographs
k = f"e2e_{tag}_mt4b_fromocr"
x0, y0 = e3["ocr_cer"]["all"]["cer"], e3[k]["all"]["chrfpp"]
ax1.plot([x0], [y0], ls="none", marker="D", ms=5, color=REC, zorder=5)
ax1.annotate("all test\nphotographs", (x0, y0), xytext=(5, 3), textcoords="offset points", ha="left",
             va="bottom", fontsize=6.5, color=MUTED)

for b, x, y in zip(bins, bx, by):
    print(b, round(x, 1), "clean", round(y, 1), "noise",
          round(100 * sub[b]["mt4b_noise_fromocr"]["chrfpp"] / sub[b]["mt4b_noise_fromgold"]["chrfpp"], 1))

for ax in (ax1, ax2):
    ax.grid(color=GRID, lw=0.5)
    ax.set_axisbelow(True)
    ax.set_xlim(-3, XMAX)
ax1.set_ylabel("chrF++")
ax2.set_ylabel("chrF++ relative to\nreference input (%)")
ax2.set_xlabel("CER of translator input (%)")
ax1.set_ylim(0, 80)
ax2.set_ylim(0, 105)
ax1.text(0.98, 0.94, "(a)", transform=ax1.transAxes, ha="right", va="top", color=INK)
ax2.text(0.98, 0.94, "(b)", transform=ax2.transAxes, ha="right", va="top", color=INK)
handles = [Line2D([], [], color=c, marker=m, ms=4, lw=1.8) for _, _, c, m in TR] + [
    Line2D([], [], color=REC, ls=(0, (3, 2)), lw=1.2, marker="o", ms=5, mfc="white", mew=1.5),
    Line2D([], [], color=REC, ls="none", marker="D", ms=4.5)]
labels = ["synthetic noise, LoRA 4B", "synthetic noise, LoRA 4B + noise",
          "recognizer output by CER group, LoRA 4B", "recognizer output, all photographs"]
fig.legend(handles, labels, loc="upper center", ncol=1, fontsize=6.5, frameon=False, labelcolor=INK,
           bbox_to_anchor=(0.56, 1.0), handlelength=2.6)
plt.tight_layout(h_pad=0.4, rect=(0, 0, 1, 0.83))
plt.savefig(f"{ROOT}/paper/figures/noise_curve.pdf")
plt.savefig(f"{ROOT}/results/noise_curve_preview.png", dpi=200)
print("saved")