fig = plt.figure(layout="tight", figsize=(12, 6.5))
gs = GridSpec(3, 4, figure=fig, height_ratios=[1, 1, 0.0])

for i, (name, title) in enumerate(
    {
        "nan.json": "Missing NaN Values",
        "missing.json": "Missing Non-NaN Values",
        "specific-humidity-log10.json": r"Pointwise $\log_{10}$",
        "kinetic-energy.json": "Kinetic Energy",
        "derivative-radial.json": r"Laplacian ($\Delta x$=const)",
        "derivative-log-exp.json": "Logarithm of Laplacian",
        "vorticity.json": "Relative Vorticity",
        "dssim.json": "dSSIM (3x3)",
    }.items()
):
    with Path("observations").joinpath(name).open() as f:
        observations = json.load(f)

    scatter_colors = []
    is_iteratives = []
    is_losslesss = []
    codec_encode_timings = []
    codec_decode_timings = []
    corrections_compute_timings = []
    corrections_encode_timings = []

    raw_codec_encode_timings = {c: [] for c in colors.keys()}

    zstd_encode_timing = None
    optzconfig_encode_timing = {c: None for c in colors.keys()}
    qpet_sperr_encode_timing = None

    for result in observations:
        if result["codec"]["id"] == "zstd.rs":
            (zstd_encode_timing,) = result["encode_timing"][
                observe.json_hash(result["codec"])
            ]

        if result["codec"]["id"] == "pressio.rs":
            codec_id = result["codec"]["early_config"]["numcodecs.rs:id"].removeprefix(
                "e-"
            )
            if codec_id == "sz3.3.rs":  # skip when no corrections are needed
                continue
            if codec_id == "sz3.2.rs":
                codec_id = "sz3.rs"
            assert codec_id in ["zfp.rs", "sz3.rs", "sperr.rs"]
            optzconfig_encode_timing[codec_id] = result["encode_timing"][
                observe.json_hash(result["codec"])
            ]

        if result["codec"]["id"] == "qpet-sperr.rs":
            (qpet_sperr_encode_timing,) = result["encode_timing"][
                observe.json_hash(result["codec"])
            ]

        if result["codec"]["id"] != "safeguards":
            continue

        codec_id = result["codec"]["codec"]["id"]

        if codec_id == "combinators.stack":
            a, b = result["codec"]["codec"]["codecs"]
            assert a["id"] == "replace.filter", a["id"]
            codec_id = b["id"]
        elif codec_id == "mask.meta":
            codec_id = result["codec"]["codec"]["codec"]["id"]
        elif codec_id == "pw_ratio":
            codec_id = result["codec"]["codec"]["log_codec"]["id"]
        if codec_id in ["sz3.2.rs", "sz3.3.rs"]:
            codec_id = "sz3.rs"

        assert codec_id in ["zero", "zfp.rs", "sz3.rs", "sperr.rs"], codec_id

        if codec_id == "zero":
            continue

        is_iterative = result["codec"]["compute"]["unstable_iterative"]
        is_lossless = result["codec"]["compute"]["unstable_lossless_corrections"]

        assert result["codec"]["lossless"]["for_codec"] is None

        (safeguarded_encode_timing,) = result["encode_timing"][
            observe.json_hash(result["codec"])
        ]
        (codec_encode_encode_timing,) = result["encode_timing"][
            observe.json_hash(result["codec"]["codec"])
        ]
        codec_encode_decode_timing, codec_decode_timing = result["decode_timing"][
            observe.json_hash(result["codec"]["codec"])
        ]
        # corrections lossless codec may be skipped if no corrections are needed
        (lossless_for_corrections_encode_timing,) = result["encode_timing"].get(
            observe.json_hash(result["codec"]["lossless"]["for_corrections"]), [None]
        )

        # skip when no corrections were needed
        # - missing with mask meta-codec
        # - log10 with pwe meta-codec
        # - NaN with SZ3[v3.3]
        # - radial Laplacian with the 3rd error bound
        if lossless_for_corrections_encode_timing is None:
            continue

        raw_codec_encode_timings[codec_id].append(codec_encode_encode_timing)

        scatter_colors.append(colors[codec_id])
        is_iteratives.append(is_iterative and not is_lossless)
        is_losslesss.append(is_iterative and is_lossless)
        codec_encode_timings.append(
            codec_encode_encode_timing / codec_encode_encode_timing
        )
        codec_decode_timings.append(
            codec_encode_decode_timing / codec_encode_encode_timing
        )
        corrections_compute_timings.append(
            lossless_for_corrections_encode_timing / codec_encode_encode_timing
        )
        corrections_encode_timings.append(
            (
                safeguarded_encode_timing
                - codec_encode_encode_timing
                - codec_encode_decode_timing
                - lossless_for_corrections_encode_timing
            )
            / codec_encode_encode_timing
        )

    is_losslesss = np.array(is_losslesss)
    is_iteratives = np.array(is_iteratives)

    ax = fig.add_subplot(gs[i % 2, i // 2])

    ax.set_title(title)

    rng = np.random.default_rng(seed=42)

    b1 = ax.boxplot(codec_encode_timings, notch=True, showfliers=False, positions=[0])
    b2 = ax.boxplot(
        np.array(codec_decode_timings) + np.array(codec_encode_timings),
        notch=False,
        showfliers=False,
        positions=[0.2],
    )
    b3 = ax.boxplot(
        np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings),
        notch=False,
        showfliers=False,
        positions=[0.4],
    )
    b4 = ax.boxplot(
        np.array(corrections_compute_timings)
        + np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings),
        notch=False,
        showfliers=False,
        positions=[0.6],
    )

    total_timings = (
        np.array(corrections_compute_timings)
        + np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings)
    )

    extra = (ax.get_ylim()[1] - ax.get_ylim()[0]) / 10

    ax.set_xlim(-0.1, 0.7)
    ax.set_xticklabels(["CE", "+CD", "+Sg", "+CC"])
    ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1] + extra)
    ax.set_yticks(
        [
            1 if y == 0 else y
            for y in ax.get_yticks()
            if (y >= 1 and y <= ax.get_ylim()[1]) or y == 0
        ]
    )
    ax.get_yticklabels()[-1].set_fontweight("bold")
    ax.text(
        0.68,
        b4["caps"][1].get_data()[1][0] + extra / 2,
        rf"max: $\times${np.round(np.amax(total_timings), 1)}",
        ha="right",
        va="bottom",
    )
    ax.text(
        0.6,
        b4["medians"][0].get_data()[1][0],
        f"{np.round(np.median(total_timings), 1)}",
        ha="center",
        va="bottom",
        color=b4["medians"][0].get_color(),
        path_effects=[mpe.withStroke(linewidth=2, foreground="white", alpha=0.5)],
    )

    cs = np.array(scatter_colors)

    xs = rng.normal(loc=0.0, scale=0.1 / 5, size=len(codec_encode_timings))
    ys = np.array(codec_encode_timings)
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    xs = rng.normal(loc=0.2, scale=0.1 / 5, size=len(codec_decode_timings))
    ys = np.array(codec_decode_timings) + np.array(codec_encode_timings)
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    xs = rng.normal(loc=0.4, scale=0.1 / 5, size=len(corrections_encode_timings))
    ys = (
        np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings)
    )
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    xs = rng.normal(loc=0.6, scale=0.1 / 5, size=len(corrections_compute_timings))
    ys = (
        np.array(corrections_compute_timings)
        + np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings)
    )
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    extra_texts = []

    extra_texts.append(
        mof.TextArea(
            "ZSTD(22):",
            textprops=dict(ha="left", va="center", color="gray", fontsize="small"),
        )
    )
    for c in colors.keys():
        zstd_rel_encode_timings = zstd_encode_timing / np.array(
            raw_codec_encode_timings[c]
        )
        extra_texts.append(
            mof.TextArea(
                rf"$\times${np.round(np.mean(zstd_rel_encode_timings), 1)}$\pm${np.round(np.std(zstd_rel_encode_timings), 1)}",
                textprops=dict(
                    ha="left", va="center", color=colors[c], fontsize="small"
                ),
            )
        )

    extra_texts.append(
        mof.TextArea("", textprops=dict(ha="left", va="center", fontsize=4))
    )
    extra_texts.append(
        mof.TextArea(
            "OptZConfig:",
            textprops=dict(
                ha="left", va="center", color=mpl.cm.tab20b.colors[10], fontsize="small"
            ),
        )
    )
    for c in colors.keys():
        optzconfig_rel_encode_timings = optzconfig_encode_timing[c] / np.array(
            raw_codec_encode_timings[c]
        )
        extra_texts.append(
            mof.TextArea(
                rf"$\times${np.round(np.mean(optzconfig_rel_encode_timings), 1)}$\pm${np.round(np.std(optzconfig_rel_encode_timings), 1)}",
                textprops=dict(
                    ha="left", va="center", color=colors[c], fontsize="small"
                ),
            )
        )

    if qpet_sperr_encode_timing is not None:
        qpet_sperr_rel_encode_timings = qpet_sperr_encode_timing / np.array(
            raw_codec_encode_timings["sperr.rs"]
        )
        extra_texts.append(
            mof.TextArea("", textprops=dict(ha="left", va="center", fontsize=4))
        )
        extra_texts.append(
            mof.TextArea(
                "QPET-SPERR:",
                textprops=dict(
                    ha="left",
                    va="center",
                    color=mpl.cm.tab20b.colors[4],
                    fontsize="small",
                ),
            )
        )
        extra_texts.append(
            mof.TextArea(
                rf"$\times${np.round(np.mean(qpet_sperr_rel_encode_timings), 1)}$\pm${np.round(np.std(qpet_sperr_rel_encode_timings), 1)}",
                textprops=dict(
                    ha="left", va="center", color=colors["sperr.rs"], fontsize="small"
                ),
            )
        )

    extra_text = mof.VPacker(
        children=extra_texts,
        sep=2,
    )
    extra_annot = mof.AnnotationBbox(
        extra_text,
        xy=(
            -0.08,
            b4["caps"][1].get_data()[1][0] + extra * 1.35,
        ),
        box_alignment=(0, 1),
        pad=0,
        frameon=False,
    )
    ax.add_artist(extra_annot)

fig.supylabel("relative compression time")

ax = fig.add_subplot(gs[2, :])
ax.axis("off")
l1 = ax.legend(
    handles=[
        mpl.lines.Line2D(
            [0],
            [0],
            marker="o",
            color="none",
            markerfacecolor=col,
            markeredgecolor=col,
            label={
                "zfp.rs": "ZFP",
                "sz3.rs": "SZ3",
                "sperr.rs": "SPERR",
            }[c],
        )
        for c, col in colors.items()
    ],
    ncol=3,
    loc="center left",
    title="Compressor",
)
ax.legend(
    handles=[
        mpl.lines.Line2D(
            [0],
            [0],
            marker="D",
            color="none",
            markerfacecolor=colors["zfp.rs"],
            markeredgecolor=colors["zfp.rs"],
            label="lossless",
        ),
        mpl.lines.Line2D(
            [0],
            [0],
            marker="o",
            color="none",
            markerfacecolor=colors["zfp.rs"],
            markeredgecolor=colors["zfp.rs"],
            label="one-shot",
        ),
        mpl.lines.Line2D(
            [0],
            [0],
            marker="*",
            color="none",
            markerfacecolor=colors["zfp.rs"],
            markeredgecolor=colors["zfp.rs"],
            label="iterative",
        ),
    ],
    ncol=3,
    loc="center right",
    title="Safeguards Corrections",
)
ax.add_artist(l1)

plt.tight_layout()

plt.savefig(Path("plots") / "summary-overhead-compression.pdf", dpi=300)

plt.show()

fig, (ax, ax2, ax3) = plt.subplots(
    1, 3, figsize=(7, 4), width_ratios=[19.5, 1.5, 1.5], gridspec_kw=dict(wspace=0.2)
)

rng = np.random.default_rng(seed=42)

scatter_colors = []
codec_decode_timings = []
corrections_decode_timings = []
corrections_apply_timings = []
zstd_rel_decode_timings = {c: [] for c in colors.keys()}
optzconfig_rel_decode_timings = {c: [] for c in colors.keys()}
qpet_sperr_rel_decode_timings = []

for name in [
    "derivative-log-exp.json",
    "derivative-radial.json",
    "dssim.json",
    "kinetic-energy.json",
    "missing.json",
    "nan.json",
    "specific-humidity-log10.json",
    "vorticity.json",
]:
    with Path("observations").joinpath(name).open() as f:
        observations = json.load(f)

    raw_codec_decode_timings = {c: [] for c in colors.keys()}

    zstd_decode_timing = None
    optzconfig_decode_timing = {c: None for c in colors.keys()}
    qpet_sperr_decode_timing = None

    for result in observations:
        if result["codec"]["id"] == "zstd.rs":
            (zstd_decode_timing,) = result["decode_timing"][
                observe.json_hash(result["codec"])
            ]

        if result["codec"]["id"] == "pressio.rs":
            codec_id = result["codec"]["early_config"]["numcodecs.rs:id"].removeprefix(
                "e-"
            )
            if codec_id == "sz3.3.rs":  # skip when no corrections are needed
                continue
            if codec_id == "sz3.2.rs":
                codec_id = "sz3.rs"
            assert codec_id in ["zfp.rs", "sz3.rs", "sperr.rs"]
            optzconfig_decode_timing[codec_id] = result["decode_timing"][
                observe.json_hash(result["codec"])
            ]

        if result["codec"]["id"] == "qpet-sperr.rs":
            (qpet_sperr_decode_timing,) = result["decode_timing"][
                observe.json_hash(result["codec"])
            ]

        if result["codec"]["id"] != "safeguards":
            continue

        codec_id = result["codec"]["codec"]["id"]

        if codec_id == "combinators.stack":
            a, b = result["codec"]["codec"]["codecs"]
            assert a["id"] == "replace.filter", a["id"]
            codec_id = b["id"]
        elif codec_id == "mask.meta":
            codec_id = result["codec"]["codec"]["codec"]["id"]
        elif codec_id == "pw_ratio":
            codec_id = result["codec"]["codec"]["log_codec"]["id"]
        if codec_id in ["sz3.2.rs", "sz3.3.rs"]:
            codec_id = "sz3.rs"

        assert codec_id in ["zero", "zfp.rs", "sz3.rs", "sperr.rs"], codec_id

        if codec_id == "zero":
            continue

        assert result["codec"]["lossless"]["for_codec"] is None

        codec_encode_decode_timing, codec_decode_timing = result["decode_timing"][
            observe.json_hash(result["codec"]["codec"])
        ]
        (safeguarded_decode_timing,) = result["decode_timing"][
            observe.json_hash(result["codec"])
        ]
        # corrections lossless codec may be skipped if no corrections are needed
        (lossless_for_corrections_decode_timing,) = result["decode_timing"].get(
            observe.json_hash(result["codec"]["lossless"]["for_corrections"]), [None]
        )

        # skip when no corrections were needed
        # - missing with mask meta-codec
        # - log10 with pwe meta-codec
        # - NaN with SZ3[v3.3]
        # - radial Laplacian with the 3rd error bound
        if lossless_for_corrections_decode_timing is None:
            continue

        raw_codec_decode_timings[codec_id].append(codec_decode_timing)

        scatter_colors.append(colors[codec_id])
        codec_decode_timings.append(codec_decode_timing / codec_decode_timing)
        corrections_decode_timings.append(
            lossless_for_corrections_decode_timing / codec_decode_timing
        )
        corrections_apply_timings.append(
            (
                safeguarded_decode_timing
                - codec_decode_timing
                - lossless_for_corrections_decode_timing
            )
            / codec_decode_timing
        )

    for c in colors.keys():
        zstd_rel_decode_timings[c].extend(
            zstd_decode_timing / np.array(raw_codec_decode_timings[c])
        )
        optzconfig_rel_decode_timings[c].extend(
            optzconfig_decode_timing[c] / np.array(raw_codec_decode_timings[c])
        )
    if qpet_sperr_decode_timing is not None:
        qpet_sperr_rel_decode_timings.extend(
            qpet_sperr_decode_timing / np.array(raw_codec_decode_timings["sperr.rs"])
        )

b1 = ax.boxplot(codec_decode_timings, notch=False, showfliers=False, positions=[0])
b2 = ax.boxplot(
    np.array(corrections_decode_timings) + np.array(codec_decode_timings),
    notch=False,
    showfliers=False,
    positions=[0.2],
)
b3 = ax.boxplot(
    np.array(corrections_apply_timings)
    + np.array(corrections_decode_timings)
    + np.array(codec_decode_timings),
    notch=False,
    showfliers=False,
    positions=[0.4],
)

total_timings = (
    np.array(corrections_apply_timings)
    + np.array(corrections_decode_timings)
    + np.array(codec_decode_timings)
)

ax.text(
    0.4,
    b3["caps"][1].get_data()[1][0] + 0.005,
    rf"max: $\times${np.amax(total_timings):0.3}",
    ha="center",
    va="bottom",
    path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
)
ax.text(
    0.4,
    b3["medians"][0].get_data()[1][0],
    rf"$\times${np.median(total_timings):0.3}",
    ha="center",
    va="bottom",
    color=b3["medians"][0].get_color(),
    path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
)

ax.set_xlim(-0.1, 0.5)
ax.set_xticks([0, 0.2, 0.4])
ax.set_xticklabels(["decompress codec", "+ corrections", "+ apply corrections"])
ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1] + 0.01)
ax.set_ylabel("relative decompression time")

b4 = ax2.boxplot(
    sum((list(t) for t in zstd_rel_decode_timings.values()), start=[]),
    notch=False,
    showfliers=False,
    positions=[0.0],
)
ax2.set_xlim(-0.15, 0.15)
ax2.set_xticks([0])
ax2.set_xticklabels(["ZSTD"])

b5 = ax3.boxplot(
    sum((list(t) for t in optzconfig_rel_decode_timings.values()), start=[]),
    notch=False,
    showfliers=False,
    positions=[0.0],
)
ax3.set_xlim(-0.15, 0.15)
ax3.set_xticks([0])
ax3.set_xticklabels(["OptZConfig"])
ax3.yaxis.tick_right()

ax.scatter(
    rng.normal(loc=0.0, scale=0.1 / 5, size=len(codec_decode_timings)),
    codec_decode_timings,
    marker=".",
    c=scatter_colors,
)
ax.scatter(
    rng.normal(loc=0.2, scale=0.1 / 5, size=len(corrections_decode_timings)),
    np.array(corrections_decode_timings) + np.array(codec_decode_timings),
    marker=".",
    c=scatter_colors,
)
ax.scatter(
    rng.normal(loc=0.4, scale=0.1 / 5, size=len(corrections_apply_timings)),
    np.array(corrections_apply_timings)
    + np.array(corrections_decode_timings)
    + np.array(codec_decode_timings),
    marker=".",
    c=scatter_colors,
)
ax2.scatter(
    rng.normal(
        loc=0.0,
        scale=0.1 / 5,
        size=sum(len(z) for z in zstd_rel_decode_timings.values()),
    ),
    sum((list(z) for z in zstd_rel_decode_timings.values()), start=[]),
    marker=".",
    c=[colors[c] for c in colors.keys() for _ in zstd_rel_decode_timings[c]],
    s=4,
)
ax3.scatter(
    rng.normal(
        loc=0.0,
        scale=0.1 / 5,
        size=sum(len(o) for o in optzconfig_rel_decode_timings.values()),
    ),
    sum((list(o) for o in optzconfig_rel_decode_timings.values()), start=[]),
    marker=".",
    c=[colors[c] for c in colors.keys() for _ in optzconfig_rel_decode_timings[c]],
    s=4,
)
ax3.set_ylim(*ax2.get_ylim())

zstd_text = mof.HPacker(
    children=[
        mof.TextArea(
            "ZSTD(22):",
            textprops=dict(
                ha="left",
                va="center",
                color="gray",
                fontsize="small",
                path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
            ),
        ),
    ]
    + [
        t
        for c in zstd_rel_decode_timings.keys()
        for t in [
            mof.TextArea(
                rf"$\times${np.round(np.mean(zstd_rel_decode_timings[c]), 1)}$\pm${np.round(np.std(zstd_rel_decode_timings[c]), 1)}",
                textprops=dict(
                    ha="left",
                    va="center",
                    color=colors[c],
                    fontsize="small",
                    path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
                ),
            ),
            mof.TextArea(
                "/",
                textprops=dict(
                    ha="left",
                    va="center",
                    color="gray",
                    fontsize="small",
                    path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
                ),
            ),
        ]
    ][:-1]
)
optzconfig_text = mof.HPacker(
    children=[
        mof.TextArea(
            "OptZConfig:",
            textprops=dict(
                ha="left",
                va="center",
                color=mpl.cm.tab20b.colors[10],
                fontsize="small",
                path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
            ),
        ),
    ]
    + [
        t
        for c in optzconfig_rel_decode_timings.keys()
        for t in [
            mof.TextArea(
                rf"$\times${np.round(np.mean(optzconfig_rel_decode_timings[c]), 1)}$\pm${np.round(np.std(optzconfig_rel_decode_timings[c]), 1)}",
                textprops=dict(
                    ha="left",
                    va="center",
                    color=colors[c],
                    fontsize="small",
                    path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
                ),
            ),
            mof.TextArea(
                "/",
                textprops=dict(
                    ha="left",
                    va="center",
                    color=mpl.cm.tab20b.colors[10],
                    fontsize="small",
                    path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
                ),
            ),
        ]
    ][:-1]
)
qpet_sperr_text = mof.HPacker(
    children=[
        mof.TextArea(
            "QPET-SPERR:",
            textprops=dict(
                ha="left",
                va="center",
                color=mpl.cm.tab20b.colors[4],
                fontsize="small",
                path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
            ),
        ),
        mof.TextArea(
            rf"$\times${np.round(np.mean(qpet_sperr_rel_decode_timings), 1)}$\pm${np.round(np.std(qpet_sperr_rel_decode_timings), 1)}",
            textprops=dict(
                ha="left",
                va="center",
                color=colors["sperr.rs"],
                fontsize="small",
                path_effects=[mpe.withStroke(linewidth=3, foreground="white")],
            ),
        ),
    ]
)

extra_texts = mof.VPacker(
    children=[
        zstd_text,
        mof.TextArea("", textprops=dict(ha="left", va="center", fontsize=4)),
        optzconfig_text,
        mof.TextArea("", textprops=dict(ha="left", va="center", fontsize=4)),
        qpet_sperr_text,
    ],
    sep=2,
)

extra_annot = mof.AnnotationBbox(
    extra_texts,
    xy=(
        -0.08,
        b3["caps"][1].get_data()[1][0] + 0.011,
    ),
    box_alignment=(0, 1),
    pad=0,
    frameon=False,
)
ax.add_artist(extra_annot)

ax.legend(
    handles=[
        mpl.lines.Line2D(
            [0],
            [0],
            marker="o",
            color="none",
            markerfacecolor=col,
            markeredgecolor=col,
            label={
                "zfp.rs": "ZFP",
                "sz3.rs": "SZ3",
                "sperr.rs": "SPERR",
            }[c],
        )
        for c, col in colors.items()
    ],
    loc="center left",
    title="Compressor",
)

# plt.tight_layout()

plt.savefig(
    Path("plots") / "summary-overhead-decompression.pdf", dpi=300, bbox_inches="tight"
)

plt.show()

fig, (ax1, axm, ax2) = plt.subplots(
    1, 3, width_ratios=[4, 0.5, 1], gridspec_kw=dict(wspace=0)
)

rng = np.random.default_rng(seed=42)

xticklabels = []

for i, (name, title) in enumerate(
    {
        "nan.tex": "Missing NaN Values",
        "missing.tex": "Missing Non-NaN Values",
        "specific-humidity-log10.tex": r"Pointwise $\log_{10}$",
        "kinetic-energy.tex": "Kinetic Energy",
        "derivative-radial.tex": r"Laplacian ($\Delta x$=const)",
        "derivative-log-exp.tex": "Logarithm of Laplacian",
        "vorticity.tex": "Relative Vorticity",
        "dssim.tex": "dSSIM (3x3)",
    }.items()
):
    crc = "CR (any NaN)" if name == "nan.tex" else "CR"

    cr = parse_latex_table(Path("tables") / name)[[crc, "C"]]

    # exclude results with no corrections
    cr = cr[cr["C"] != "0"][crc]

    cr = pd.to_numeric(cr.str.removeprefix(r"$\times$ "), errors="coerce")

    lossless = cr[cr.index.get_level_values("Corrections") == "lossless"].droplevel(
        "Corrections"
    )
    oneshot = cr[cr.index.get_level_values("Corrections") == "one-shot"].droplevel(
        "Corrections"
    )
    iterative = cr[cr.index.get_level_values("Corrections") == "iterative"].droplevel(
        "Corrections"
    )

    zstd = (
        cr[cr.index.get_level_values("Compressor") == "ZSTD(22)"].values
        / oneshot.values
    )
    optzconfig = {}
    for c in colors.keys():
        cname = {
            "zfp.rs": "ZFP",
            "sz3.rs": "SZ3",
            "sperr.rs": "SPERR",
        }[c]
        if (c == "sz3.rs") and (
            r"SZ3[v3.2]($\epsilon_{abs}$)" in cr.index.get_level_values("Compressor")
        ):
            cname = "SZ3[v3.2]"
        optzconfig[c] = (
            cr[cr.index.get_level_values("Compressor") == f"OptZConfig({cname})"].values
            / oneshot[
                oneshot.index.get_level_values("Compressor").isin(
                    [cname, rf"{cname}($\epsilon_{{abs}}$)"]
                )
            ].values
        )
    if np.any(cr.index.get_level_values("Compressor") == "QPET-SPERR"):
        qpet_sperr = (
            cr[cr.index.get_level_values("Compressor") == "QPET-SPERR"].values
            / oneshot[
                oneshot.index.get_level_values("Compressor").isin(
                    ["SPERR", r"SPERR($\epsilon_{abs}$)"]
                )
            ].values
        )
    else:
        qpet_sperr = np.zeros(0)

    ax1.boxplot(
        lossless / oneshot,
        notch=False,
        showfliers=False,
        positions=[-i * 0.2],
        orientation="horizontal",
    )

    ax2.boxplot(
        iterative / oneshot,
        notch=False,
        showfliers=False,
        positions=[-i * 0.2],
        orientation="horizontal",
    )

    scatter_colors = [
        colors_with_zero[c] for c in lossless.index.get_level_values("Compressor")
    ]

    ax1.scatter(
        lossless / oneshot,
        rng.normal(loc=-i * 0.2, scale=0.1 / 5, size=len(oneshot)),
        marker="D",
        s=8,
        color=scatter_colors,
    )
    ax2.scatter(
        iterative / oneshot,
        rng.normal(loc=-i * 0.2, scale=0.1 / 5, size=len(oneshot)),
        marker="*",
        s=16,
        color=scatter_colors,
    )

    qpet_scatter = rng.normal(loc=-i * 0.2, scale=0.1 / 5, size=len(qpet_sperr))
    ax1.scatter(
        [x for x in qpet_sperr if x < 1],
        [y for x, y in zip(qpet_sperr, qpet_scatter) if x < 1],
        marker="X",
        s=24,
        color=mpl.cm.tab20b.colors[4],
    )
    ax2.scatter(
        [x for x in qpet_sperr if x >= 1],
        [y for x, y in zip(qpet_sperr, qpet_scatter) if x >= 1],
        marker="X",
        s=24,
        color=mpl.cm.tab20b.colors[4],
    )

    zstd_scatter = -i * 0.2 + 0.125 * (1 - np.linspace(0.0, 1.0, len(zstd)))
    axm.scatter(
        [1 / 3 for x in zstd if x < 1],
        [y for x, y in zip(zstd, zstd_scatter) if x < 1],
        marker="<",
        s=16,
        color=[c for x, c in zip(zstd, scatter_colors) if x < 1],
        edgecolors="gray",
        linewidths=1,
    )
    axm.scatter(
        [1 / 3 for x in zstd if x >= 1],
        [y for x, y in zip(zstd, zstd_scatter) if x >= 1],
        marker=">",
        s=16,
        color=[c for x, c in zip(zstd, scatter_colors) if x >= 1],
        edgecolors="gray",
        linewidths=1,
    )
    optzconfig_scatter = -i * 0.2 + 0.125 * (
        1
        - np.linspace(0.0, 1.0, sum(len(o) for o in optzconfig.values()) * 4 // 3)[
            -sum(len(o) for o in optzconfig.values()) :
        ]
    )
    axm.scatter(
        [2 / 3 for c in colors.keys() for x in optzconfig[c] if x < 1],
        optzconfig_scatter[[x < 1 for c in colors.keys() for x in optzconfig[c]]],
        marker="<",
        s=16,
        color=[colors[c] for c in colors.keys() for x in optzconfig[c] if x < 1],
        edgecolors=mpl.cm.tab20b.colors[10],
        linewidths=1,
    )
    axm.scatter(
        [2 / 3 for c in colors.keys() for x in optzconfig[c] if x >= 1],
        optzconfig_scatter[[x >= 1 for c in colors.keys() for x in optzconfig[c]]],
        marker=">",
        s=16,
        color=[colors[c] for c in colors.keys() for x in optzconfig[c] if x >= 1],
        edgecolors=mpl.cm.tab20b.colors[10],
        linewidths=1,
    )

    xticklabels.append(title)

ax1.set_title("Lossless")
ax1.set_xscale("log")
ax1.set_xlabel("compression ratio, relative to one-shot corrections")
ax1.set_yticks(np.arange(len(xticklabels)) * -0.2)
ax1.set_yticklabels(xticklabels)
ax1.axvline(1, c="k", zorder=-1, ls=":")
ax1.set_ylim(len(xticklabels) * -0.2 + 0.05, 0.15)
ax1.set_xlim(0.015, 1.25)
ax1.set_xticks([0.02, 0.05, 0.1, 0.2, 0.5, 1.0])
ax1.set_xticklabels(
    [r"$\div$50", r"$\div$20", r"$\div$10", r"$\div$5", r"$\div$2", r"$\div$1"]
)

ax2.set_title("Iterative")
ax2.set_xscale("log")
ax2.set_yticks(np.arange(len(xticklabels)) * -0.2)
ax2.set_yticklabels([""] * len(xticklabels))
ax2.axvline(1, c="k", zorder=-1, ls=":")
ax2.set_ylim(len(xticklabels) * -0.2 + 0.05, 0.15)
ax2.set_xlim(
    1
    / np.power(
        10.0,
        np.log10(2.5) / (1 - 4 * np.log10(1.25) / (np.log10(1.25) - np.log10(0.015)))
        - np.log10(2.5),
    ),
    2.5,
)
ax2.set_xticks([1.0, 2.0])
ax2.get_xaxis().set_major_formatter(
    mpl.ticker.FuncFormatter(lambda x, y: rf"$\times${int(x)}")
)
ax2.get_xaxis().set_minor_formatter("")

axm.axis("off")
axm.set_xlim(0, 1)

l1 = ax1.legend(
    handles=[
        mpl.lines.Line2D(
            [0],
            [0],
            marker="o",
            color="none",
            markerfacecolor=col,
            markeredgecolor=col,
            label=c,
        )
        for c, col in {
            "0": mpl.cm.tab10.colors[8],
            r"SZ3": mpl.cm.tab10.colors[6],
            r"ZFP": mpl.cm.tab10.colors[0],
            r"SPERR": mpl.cm.tab10.colors[2],
        }.items()
    ],
    ncol=2,
    loc="upper left",
    title="Compressor",
)
ax1.legend(
    handles=[
        mpl.lines.Line2D(
            [0],
            [0],
            marker=m,
            color="none",
            markerfacecolor=mpl.cm.tab10.colors[8]
            if c in ["ZSTD(22)", "OptZConfig"]
            else col,
            markeredgecolor=col,
            markeredgewidth=mw,
            label=c,
        )
        for c, (col, m, mw) in {
            "ZSTD(22)": ("gray", "<", 1),
            "OptZConfig": (mpl.cm.tab20b.colors[10], "<", 1),
            "QPET-SPERR": (mpl.cm.tab20b.colors[4], "X", 0),
        }.items()
    ],
    ncol=1,
    loc="center left",
    title="Comparison",
)
ax1.add_artist(l1)

# plt.tight_layout()

plt.savefig(Path("plots") / "summary-corrections.pdf", dpi=300, bbox_inches="tight")

plt.show()

fig = plt.figure(layout="tight", figsize=(9, 3))
gs = GridSpec(1, 3, figure=fig)

all_scatter_colors = []
all_codec_decode_timings = []
all_corrections_decode_timings = []
all_corrections_apply_timings = []

i = 0

for name, title in {
    "nan.json": "Missing NaN Values",
    "missing.json": "Missing Non-NaN Values",
    "specific-humidity-log10.json": r"Pointwise $\log_{10}$",
    "kinetic-energy.json": "Kinetic Energy",
    "derivative-radial.json": r"Laplacian ($\Delta x$=const)",
    "derivative-log-exp.json": "Logarithm of Laplacian",
    "vorticity.json": "Relative Vorticity",
    "dssim.json": "dSSIM (3x3)",
}.items():
    with Path("observations").joinpath(name).open() as f:
        observations = json.load(f)

    scatter_colors = []
    is_iteratives = []
    is_losslesss = []
    codec_encode_timings = []
    codec_decode_timings = []
    corrections_compute_timings = []
    corrections_encode_timings = []

    for result in observations:
        if result["codec"]["id"] != "safeguards":
            continue

        codec_id = result["codec"]["codec"]["id"]

        if codec_id == "combinators.stack":
            a, b = result["codec"]["codec"]["codecs"]
            assert a["id"] == "replace.filter", a["id"]
            codec_id = b["id"]
        elif codec_id == "mask.meta":
            codec_id = result["codec"]["codec"]["codec"]["id"]
        elif codec_id == "pw_ratio":
            codec_id = result["codec"]["codec"]["log_codec"]["id"]
        if codec_id in ["sz3.2.rs", "sz3.3.rs"]:
            codec_id = "sz3.rs"

        assert codec_id in ["zero", "zfp.rs", "sz3.rs", "sperr.rs"], codec_id

        if codec_id == "zero":
            continue

        is_iterative = result["codec"]["compute"]["unstable_iterative"]
        is_lossless = result["codec"]["compute"]["unstable_lossless_corrections"]

        assert result["codec"]["lossless"]["for_codec"] is None

        (safeguarded_encode_timing,) = result["encode_timing"][
            observe.json_hash(result["codec"])
        ]
        (codec_encode_encode_timing,) = result["encode_timing"][
            observe.json_hash(result["codec"]["codec"])
        ]
        codec_encode_decode_timing, codec_decode_timing = result["decode_timing"][
            observe.json_hash(result["codec"]["codec"])
        ]
        # corrections lossless codec may be skipped if no corrections are needed
        (lossless_for_corrections_encode_timing,) = result["encode_timing"].get(
            observe.json_hash(result["codec"]["lossless"]["for_corrections"]), [None]
        )

        (safeguarded_decode_timing,) = result["decode_timing"][
            observe.json_hash(result["codec"])
        ]
        # corrections lossless codec may be skipped if no corrections are needed
        (lossless_for_corrections_decode_timing,) = result["decode_timing"].get(
            observe.json_hash(result["codec"]["lossless"]["for_corrections"]), [None]
        )

        # skip when no corrections were needed
        # - missing with mask meta-codec
        # - log10 with pwe meta-codec
        # - NaN with SZ3[v3.3]
        # - radial Laplacian with the 3rd error bound
        if lossless_for_corrections_encode_timing is None:
            continue
        assert lossless_for_corrections_decode_timing is not None

        scatter_colors.append(colors[codec_id])
        is_iteratives.append(is_iterative and not is_lossless)
        is_losslesss.append(is_iterative and is_lossless)
        codec_encode_timings.append(
            codec_encode_encode_timing / codec_encode_encode_timing
        )
        codec_decode_timings.append(
            codec_encode_decode_timing / codec_encode_encode_timing
        )
        corrections_compute_timings.append(
            lossless_for_corrections_encode_timing / codec_encode_encode_timing
        )
        corrections_encode_timings.append(
            (
                safeguarded_encode_timing
                - codec_encode_encode_timing
                - codec_encode_decode_timing
                - lossless_for_corrections_encode_timing
            )
            / codec_encode_encode_timing
        )

        all_scatter_colors.append(colors[codec_id])
        all_codec_decode_timings.append(codec_decode_timing / codec_decode_timing)
        all_corrections_decode_timings.append(
            lossless_for_corrections_decode_timing / codec_decode_timing
        )
        all_corrections_apply_timings.append(
            (
                safeguarded_decode_timing
                - codec_decode_timing
                - lossless_for_corrections_decode_timing
            )
            / codec_decode_timing
        )

    if name not in ["nan.json", "vorticity.json"]:
        continue

    is_losslesss = np.array(is_losslesss)
    is_iteratives = np.array(is_iteratives)

    ax = fig.add_subplot(gs[i])

    rng = np.random.default_rng(seed=42)

    b1 = ax.boxplot(codec_encode_timings, notch=True, showfliers=False, positions=[0])
    b2 = ax.boxplot(
        np.array(codec_decode_timings) + np.array(codec_encode_timings),
        notch=False,
        showfliers=False,
        positions=[0.2],
    )
    b3 = ax.boxplot(
        np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings),
        notch=False,
        showfliers=False,
        positions=[0.4],
    )
    b4 = ax.boxplot(
        np.array(corrections_compute_timings)
        + np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings),
        notch=False,
        showfliers=False,
        positions=[0.6],
    )

    extra = (ax.get_ylim()[1] - ax.get_ylim()[0]) / 10

    ax.set_xlim(-0.1, 0.7)
    ax.set_xticklabels(["CE", "+CD", "+Sg", "+CC"])
    ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1] + extra)
    ax.set_yticks(
        [
            1 if y == 0 else y
            for y in ax.get_yticks()
            if (y >= 1 and y <= ax.get_ylim()[1]) or y == 0
        ]
    )
    ax.get_yticklabels()[-1].set_fontweight("bold")
    ax.text(
        0.68,
        b4["caps"][1].get_data()[1][0] + extra / 2,
        rf"max: $\times${np.round(np.amax(corrections_compute_timings) + np.amax(corrections_encode_timings) + np.amax(codec_decode_timings) + np.amax(codec_encode_timings), 1)}",
        ha="right",
        va="bottom",
    )

    cs = np.array(scatter_colors)

    xs = rng.normal(loc=0.0, scale=0.1 / 5, size=len(codec_encode_timings))
    ys = np.array(codec_encode_timings)
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    xs = rng.normal(loc=0.2, scale=0.1 / 5, size=len(codec_decode_timings))
    ys = np.array(codec_decode_timings) + np.array(codec_encode_timings)
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    xs = rng.normal(loc=0.4, scale=0.1 / 5, size=len(corrections_encode_timings))
    ys = (
        np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings)
    )
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    xs = rng.normal(loc=0.6, scale=0.1 / 5, size=len(corrections_compute_timings))
    ys = (
        np.array(corrections_compute_timings)
        + np.array(corrections_encode_timings)
        + np.array(codec_decode_timings)
        + np.array(codec_encode_timings)
    )
    ax.scatter(xs[is_losslesss], ys[is_losslesss], marker="D", c=cs[is_losslesss], s=10)
    ax.scatter(
        xs[is_iteratives], ys[is_iteratives], marker="*", c=cs[is_iteratives], s=24
    )
    ax.scatter(
        xs[~(is_iteratives | is_losslesss)],
        ys[~(is_iteratives | is_losslesss)],
        marker="o",
        c=cs[~(is_iteratives | is_losslesss)],
        s=16,
    )

    if i == 0:
        ax.legend(
            handles=[
                mpl.lines.Line2D(
                    [0],
                    [0],
                    marker="o",
                    color="none",
                    markerfacecolor=col,
                    markeredgecolor=col,
                    label={
                        "zfp.rs": "ZFP",
                        "sz3.rs": "SZ3",
                        "sperr.rs": "SPERR",
                    }[c],
                )
                for c, col in colors.items()
            ],
            loc="upper left",
            title="Compressor",
        )
    elif i == 1:
        ax.legend(
            handles=[
                mpl.lines.Line2D(
                    [0],
                    [0],
                    marker="D",
                    color="none",
                    markerfacecolor=colors["zfp.rs"],
                    markeredgecolor=colors["zfp.rs"],
                    label="lossless",
                ),
                mpl.lines.Line2D(
                    [0],
                    [0],
                    marker="o",
                    color="none",
                    markerfacecolor=colors["zfp.rs"],
                    markeredgecolor=colors["zfp.rs"],
                    label="one-shot",
                ),
                mpl.lines.Line2D(
                    [0],
                    [0],
                    marker="*",
                    color="none",
                    markerfacecolor=colors["zfp.rs"],
                    markeredgecolor=colors["zfp.rs"],
                    label="iterative",
                ),
            ],
            loc="upper left",
            title="Corrections",
        )

    i += 1

ax = fig.add_subplot(gs[2])

b1 = ax.boxplot(all_codec_decode_timings, notch=False, showfliers=False, positions=[0])
b2 = ax.boxplot(
    np.array(all_corrections_decode_timings) + np.array(all_codec_decode_timings),
    notch=False,
    showfliers=False,
    positions=[0.2],
)
b3 = ax.boxplot(
    np.array(all_corrections_apply_timings)
    + np.array(all_corrections_decode_timings)
    + np.array(all_codec_decode_timings),
    notch=False,
    showfliers=False,
    positions=[0.4],
)

extra = (ax.get_ylim()[1] - ax.get_ylim()[0]) / 10

ax.yaxis.tick_right()

ax.set_xlim(-0.1, 0.5)
ax.set_xticks([0, 0.2, 0.4])
ax.set_xticklabels(["CD", "+ DC", "+ AC"])
ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1] + extra)
ax.set_yticks(
    [
        1 if y == 0 else y
        for y in ax.get_yticks()
        if (y >= 1 and y <= ax.get_ylim()[1]) or y == 0
    ]
)
ax.get_yticklabels()[-1].set_fontweight("bold")

ax.text(
    0.48,
    b3["caps"][1].get_data()[1][0] + extra / 2,
    rf"max: $\times${np.amax(all_corrections_apply_timings) + np.amax(all_corrections_decode_timings) + np.amax(all_codec_decode_timings):0.3}",
    ha="right",
    va="bottom",
)

ax.scatter(
    rng.normal(loc=0.0, scale=0.1 / 5, size=len(all_codec_decode_timings)),
    all_codec_decode_timings,
    marker=".",
    c=all_scatter_colors,
)
ax.scatter(
    rng.normal(loc=0.2, scale=0.1 / 5, size=len(all_corrections_decode_timings)),
    np.array(all_corrections_decode_timings) + np.array(all_codec_decode_timings),
    marker=".",
    c=all_scatter_colors,
)
ax.scatter(
    rng.normal(loc=0.4, scale=0.1 / 5, size=len(all_corrections_apply_timings)),
    np.array(all_corrections_apply_timings)
    + np.array(all_corrections_decode_timings)
    + np.array(all_codec_decode_timings),
    marker=".",
    c=all_scatter_colors,
)

plt.tight_layout()

ax.set_rasterized(True)

plt.tight_layout()

plt.savefig(Path("plots") / "summary-overhead-egu26.pdf", dpi=300)

plt.show()