#!/usr/bin/env python3
"""Plot results/trie-seed{1,2,3}.txt -> ../trie-proof-size.svg (median over seeds)."""
import pathlib
import statistics

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = pathlib.Path(__file__).resolve().parent
plt.rcParams["svg.hashsalt"] = "merkle-auth-ds"

BG, FG, HI, MUTED = "#2d333b", "#adbac7", "#cdd9e5", "#768390"
COLORS = {2: "#3fb950", 4: "#388bfd", 16: "#f0883e", 256: "#a371f7"}

rows = {}
for s in (1, 2, 3):
    for line in (HERE / "results" / f"trie-seed{s}.txt").read_text().splitlines():
        if line.startswith("#") or not line.strip():
            continue
        f = line.split()
        k, logn, mean, est = int(f[0]), int(f[1]), float(f[6]), float(f[10])
        rows.setdefault((k, logn), {"mean": [], "est": est})["mean"].append(mean)

fig, ax = plt.subplots(figsize=(8.4, 4.8), dpi=100)
fig.patch.set_facecolor(BG)
ax.set_facecolor(BG)
for k, c in COLORS.items():
    xs = sorted(logn for (kk, logn) in rows if kk == k)
    meas = [32 * statistics.median(rows[(k, x)]["mean"]) for x in xs]
    est = [32 * rows[(k, x)]["est"] for x in xs]
    ax.plot(xs, meas, "o-", color=c, lw=2, ms=5, label=f"k={k} measured")
    ax.plot(xs, est, "--", color=c, lw=1.2, alpha=0.8, label=f"k={k} 32(k-1)log_k N")
ax.set_yscale("log")
ax.set_xlabel("log2 N (number of random keys)", color=FG)
ax.set_ylabel("sibling hash bytes per proof (32 B each)", color=FG)
ax.set_title("Non-empty sibling hashes on a membership path, path-compressed k-ary trie",
             color=HI, fontsize=11)
ax.tick_params(colors=FG)
for sp in ax.spines.values():
    sp.set_color(MUTED)
ax.grid(True, which="both", color="#444c56", lw=0.5)
leg = ax.legend(ncol=2, fontsize=8, facecolor=BG, edgecolor=MUTED, loc="upper left")
for t in leg.get_texts():
    t.set_color(FG)
fig.tight_layout()
fig.savefig(HERE.parent / "trie-proof-size.svg", facecolor=BG, metadata={"Date": None})
print("wrote trie-proof-size.svg")