"""Log-log Chamfer and p95 against atoms for the two versions in horse_ladder.csv; marks every step up."""
import csv
from collections import defaultdict
from pathlib import Path

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

R = Path(__file__).resolve().parent
rows = defaultdict(list)
for r in csv.DictReader((R / "horse_ladder.csv").open()):
    rows[r["version"]].append(r)
STYLE = {0: dict(color="#2a78d6", marker="o", ls="-"), 1: dict(color="#eb6834", marker="s", ls="--")}
plt.rcParams.update({"font.size": 11, "axes.edgecolor": "#888", "axes.labelcolor": "#222", "xtick.color": "#444", "ytick.color": "#444"})
fig, axes = plt.subplots(1, 2, figsize=(16, 6.4))
ups = []
for ax, key, lab in ((axes[0], "chamfer_pct", "Chamfer, % of the diagonal"), (axes[1], "p95_pct", "p95 distance, % of the diagonal")):
    for i, (ver, rs) in enumerate(rows.items()):
        pts = [(int(r["atoms"]), float(r[key])) for r in rs if r[key]]
        x, y = zip(*pts)
        st = STYLE[i]
        ax.plot(x, y, color=st["color"], ls=st["ls"], lw=2, marker=st["marker"], ms=6, label=f"{ver}, {rs[0]['parts']} parts")
        for (x0, y0), (x1, y1) in zip(pts, pts[1:]):
            if y1 > y0 * 1.0005:  # a step up (more than 0.05 %)
                ax.plot([x1], [y1], "o", ms=15, mfc="none", mec="#222", mew=1.5)
                ax.annotate(f"{x0}→{x1}: +{100 * (y1 / y0 - 1):.1f} %", (x1, y1), xytext=(12, 12) if i == 0 else (-6, -30),
                            textcoords="offset points", fontsize=9.5, color="#222")
                if key == "chamfer_pct":
                    ups.append((ver, x0, x1, y0, y1))
        missing = [int(r["atoms"]) for r in rs if not r[key]]
        if missing and key == "chamfer_pct":
            ax.annotate(f"{ver.split(' ')[0]}: no surface at {', '.join(map(str, missing))} atoms", (0.02, 0.04),
                        xycoords="axes fraction", fontsize=9.5, color="#444")
    ax.set_xscale("log", base=2)
    ax.set_xticks([2 ** k for k in range(3, 17, 2)], [f"{2 ** k:,}" for k in range(3, 17, 2)])
    ax.set_yscale("log")
    lo, hi = ax.get_ylim()
    ax.set_ylim(lo, hi * 1.8)  # headroom for the step labels
    ax.set_xlabel("atoms kept (prefix of one encode)")
    ax.set_ylabel(lab)
    ax.grid(True, which="major", color="#e3e3e3", lw=0.8)
    ax.grid(True, which="minor", color="#f1f1f1", lw=0.5)
    ax.spines[["top", "right"]].set_visible(False)
axes[0].legend(frameon=False, loc="upper right")
fig.suptitle("Horse, 256 layers, N 64, M 24, tracked-fit, decoded on 128³: error at each prefix (ringed: a step where the error rises)",
             fontsize=12)
fig.tight_layout()
fig.savefig(R / "horse_monotonic.png", dpi=110, facecolor="white")
for u in ups:
    print("step up:", u)