"""Figure 1: the certified sum profiles psi (step functions on lattice cells) for m = 2, 3, 4, 5."""
import os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
plt.rcParams.update({"font.family": "serif", "font.serif": ["cmr10"], "mathtext.fontset": "cm",
                     "axes.formatter.use_mathtext": True, "axes.unicode_minus": False,
                     "font.size": 9, "axes.linewidth": 0.6, "xtick.major.width": 0.6, "ytick.major.width": 0.6,
                     "pdf.fonttype": 42})
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(HERE, "..", "certify")))
from trial_io import load_trial
TRIALS = os.path.join(HERE, "..", "certify", "trials")
OUTPUT = os.path.abspath(os.path.join(HERE, "..", "..", "fig"))
TR = [("A_m2_k2355", 2, 2355), ("A_m3_k111300", 3, 111300), ("A_m4_k5129500", 4, 5129500), ("A_m5_k234360000", 5, 234360000)]
COL = ["#86b6ef", "#3987e5", "#1c5cab", "#0d366b"]        # ordinal blue ramp (validated: dataviz validator, --ordinal)
LS = ["-", (0, (5, 1.6)), (0, (6, 1.4, 1.4, 1.4)), (0, (1.4, 1.2))]   # secondary encoding for greyscale print
fig, ax = plt.subplots(figsize=(5.6, 2.9))
for (f, m, k), col, ls in zip(TR, COL, LS):
    d = load_trial(os.path.join(TRIALS, f + ".json"))
    s = d["E"] * float(d["delta"]); psi = d["h"]
    xs = np.repeat(s, 2)[1:-1]; ys = np.repeat(psi, 2)
    xs = xs[2:]; ys = ys[2:]                                   # drop the cell [0, s_1), where psi = 0
    ax.plot(xs, ys, color=col, lw=1.4, ls=ls, label=r"$m=%d$, $k=%s$" % (m, f"{k:,}".replace(",", "{,}")))
    ax.text(s[1] - 0.008, psi[1], r"$m=%d$" % m, ha="right", va="center", fontsize=8, color="#333333")
ax.set_xlim(0.38, 1.005); ax.set_ylim(0, 1.05)
ax.set_xlabel(r"$s=t_1+\!\cdot\!\cdot\!\cdot\!+t_k$"); ax.set_ylabel(r"$\psi(s)$")
ax.grid(True, color="#e6e6e3", lw=0.5); ax.set_axisbelow(True)
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
ax.legend(loc="upper left", frameon=False, fontsize=8, handlelength=3.2)
fig.tight_layout()
fig.savefig(os.path.join(OUTPUT, "h_profiles.pdf"))
fig.savefig(os.path.join(OUTPUT, "h_profiles.png"), dpi=200)
print("ok")