"""Figure for report.html, from data/ only.

    python reproduce/figures.py      # writes report/speedup.png
"""

import csv
import statistics as st
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

HERE = Path(__file__).resolve().parent.parent
rows = list(csv.DictReader(open(HERE / "data" / "timings.csv", newline="")))
natoms = {r["idx"]: int(r["n_atoms"]) for r in csv.DictReader(open(HERE / "data" / "reactions.csv", newline=""))}
calls = sorted({(r["idx"], r["nimages"]) for r in rows})
t = {(c, v): st.median(float(r["seconds"]) for r in rows if (r["idx"], r["nimages"]) == c and r["version"] == v)
     for c in calls for v in ("before", "after")}

fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.0), dpi=150)
colors = {"12": "#2a5db0", "30": "#d9822b"}
for n, col in colors.items():
    cs = [c for c in calls if c[1] == n]
    axes[0].scatter([t[c, "before"] for c in cs], [t[c, "after"] for c in cs], s=14, color=col, label=f"{n} images")
    axes[1].scatter([natoms[c[0]] for c in cs], [t[c, "before"] / t[c, "after"] for c in cs], s=14, color=col)
lim = [0.03, 2]
axes[0].plot(lim, lim, color="0.6", lw=0.8, ls="--")
axes[0].text(0.75, 0.95, "same speed", color="0.45", fontsize=7, rotation=38)
axes[0].set(xscale="log", yscale="log", xlim=lim, ylim=lim, xlabel="before (s per call)", ylabel="with numba (s per call)")
axes[0].legend(frameon=False, fontsize=8)
axes[1].axhline(1, color="0.6", lw=0.8, ls="--")
axes[1].set(xlabel="atoms", ylabel="times faster", ylim=(0, 3.6))
for ax in axes:
    ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout()
(HERE / "report").mkdir(exist_ok=True)
fig.savefig(HERE / "report" / "speedup.png", metadata={"Software": None})