"""r04_figures.py: the five figures of the budget (plan section 4), drawn from the data and the values files.

Every figure is an EPS file at its final size (widths 84 mm or 174 mm, heights from the registered judgments
fig1_height_mm to fig5_height_mm), lettering in Liberation Sans (metrically identical to Arial) at fig_font_pt,
embedded as Type 42 fonts, lines of at least fig_line_pt, patterns or markers in addition to colour, and no
title inside the artwork. Plan Fig. 2 and Fig. 3 swap numbers because the Methods cite the spectra first. The spectra examples are chosen by the rule fixed in the plan; the choice is
registered in values4.json. Numbers drawn in the figures come from values4.json or from the data directly.
"""
import sys, json
sys.path.insert(0, "/home/claude/mill/phase4/code")
sys.path.insert(0, "/home/claude/mill/phase3/code")
sys.path.insert(0, "/home/claude/mill/code")
import r_freeze
r_freeze.assert_freeze()

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
from vals import Values
import q_common as QC

P3, P4 = "/home/claude/mill/phase3", "/home/claude/mill/phase4"
V2 = json.load(open("/home/claude/mill/outputs/values.json"))
V3 = json.load(open(f"{P3}/outputs/values3.json"))
V4 = json.load(open(f"{P4}/outputs/values4.json"))
V = Values(path=f"{P4}/outputs/values4.json", script="r04_figures.py")
g = lambda n: (V4[n] if n in V4 else V3[n] if n in V3 else V2[n])["value"]
ci = lambda n: (V4[n]["lo"], V4[n]["hi"])
MM = 1 / 25.4
POP = "labelled cutting paths (label.csv rows)"
FONT, LINE = float(g("fig_font_pt")), float(g("fig_line_pt"))
OUT = f"{P4}/figures"
plt.rcParams.update({
    "font.family": "Liberation Sans", "font.size": FONT, "axes.labelsize": FONT, "axes.titlesize": FONT,
    "xtick.labelsize": FONT, "ytick.labelsize": FONT, "legend.fontsize": FONT, "legend.title_fontsize": FONT,
    "ps.fonttype": 42, "pdf.fonttype": 42, "axes.linewidth": LINE, "xtick.major.width": LINE, "ytick.major.width": LINE,
    "xtick.minor.width": LINE, "ytick.minor.width": LINE, "lines.linewidth": 1.0, "patch.linewidth": LINE,
    "hatch.linewidth": LINE, "hatch.color": "black", "legend.frameon": False, "axes.spines.top": False,
    "axes.spines.right": False, "savefig.dpi": 600, "text.color": "black", "axes.labelcolor": "black",
    "xtick.color": "black", "ytick.color": "black", "axes.edgecolor": "black"})
# Okabe and Ito palette (distinguishable with colour-vision deficiency); every series also has a pattern or marker
BLUE, VERM, GREEN, ORANGE, SKY, GREY = "#0072B2", "#D55E00", "#009E73", "#E69F00", "#56B4E9", "#BBBBBB"
T = pd.read_parquet("/home/claude/mill/outputs/segments.parquet")
T["family"] = T.Tool.str.replace("*", "", regex=False)
MACH = ("imi_vm20i", "imi_vmx30ui")


def panel_label(ax, s, x=-0.02, y=1.0):
    ax.text(x, y, s, transform=ax.transAxes, fontweight="bold", va="bottom", ha="right")


def save(fig, name):
    fig.savefig(f"{OUT}/{name}.eps", format="eps")
    fig.savefig(f"{OUT}/{name}_preview.png", format="png", dpi=600)
    plt.close(fig)


# ------------------------------------------------------------------ Fig. 1: label 1 share by tool family and material, a panel per machine
MATS = ["AA 6061", "AISI 1018", "AISI 4140"]
HATCH = {"AA 6061": "", "AISI 1018": "////", "AISI 4140": "...."}
FILL = {"AA 6061": BLUE, "AISI 1018": "white", "AISI 4140": ORANGE}
fig, axes = plt.subplots(2, 1, figsize=(84 * MM, g("fig1_height_mm") * MM), sharey=True)
for ax, mach, lab in zip(axes, MACH, "ab"):
    S = T[(T.machine == mach) & T.label.isin([0, 1])]
    fams = sorted(S.family.unique())
    w = 0.26
    for j, mat in enumerate(MATS):
        for i, f in enumerate(fams):
            c = S[(S.family == f) & (S.Workpiece == mat)]
            x = i + (j - 1) * w
            if len(c) == 0:
                continue
            share = (c.label == 1).mean()
            if share > 0:
                ax.bar(x, share, w, color=FILL[mat], hatch=HATCH[mat], edgecolor="black", linewidth=LINE)
            else:
                ax.plot(x, 0.045, marker="o", markersize=2.6, markerfacecolor="white", markeredgecolor="black", markeredgewidth=0.6, linestyle="none")
    ax.set_xticks(range(len(fams)))
    ax.set_xticklabels(fams)
    ax.set_xlim(-0.6, len(fams) - 0.4)
    ax.set_ylim(0, 1.05)
    ax.set_yticks([0, 0.5, 1.0])
    ax.set_ylabel("Label 1 share")
    ax.text(0.99, 0.97, mach, transform=ax.transAxes, ha="right", va="top")
    panel_label(ax, lab, x=-0.14)
axes[1].set_xlabel("Tool family")
handles = [Patch(facecolor=FILL[m], hatch=HATCH[m], edgecolor="black", linewidth=LINE, label=m) for m in MATS]
fig.legend(handles=handles, loc="upper center", bbox_to_anchor=(0.56, 1.0), ncol=3, handlelength=1.6, columnspacing=1.2, borderaxespad=0.2)
fig.subplots_adjust(left=0.15, right=0.98, top=0.86, bottom=0.17, hspace=0.45)
save(fig, "Fig1")

# ------------------------------------------------------------------ Fig. 3 (plan Fig. 2): Sound0 full-scale share by label, empirical distributions, a panel per machine
fig, axes = plt.subplots(1, 2, figsize=(84 * MM, g("fig2_height_mm") * MM), sharey=True)
THR = float(g("d9_threshold"))
for ax, mach, lab in zip(axes, MACH, "ab"):
    S = T[(T.machine == mach) & T.label.isin([0, 1])]
    for L, style, col in ((0, (0, (3, 1.5)), BLUE), (1, "-", VERM)):
        x = np.sort(S[S.label == L].s0_fullscale.values * 100)
        yv = np.arange(1, len(x) + 1) / len(x)
        ax.step(np.concatenate([[0], x]), np.concatenate([[0], yv]), where="post", linestyle=style, color=col, linewidth=1.2)
    ax.axvline(THR * 100, color="black", linestyle=":", linewidth=0.8)
    ax.set_xlim(0, 50)
    ax.set_ylim(0, 1.0)
    ax.set_xticks([0, 10, 20, 30, 40, 50])
    ax.set_xlabel("Samples at full scale (%)")
    ax.text(0.97, 0.03, mach, transform=ax.transAxes, ha="right", va="bottom")
    panel_label(ax, lab, x=0.0, y=1.05)
axes[0].set_ylabel("Cumulative share of paths")
axes[1].legend(handles=[Line2D([], [], color=BLUE, linestyle=(0, (3, 1.5)), linewidth=1.2, label="Label 0"),
                        Line2D([], [], color=VERM, linestyle="-", linewidth=1.2, label="Label 1")],
               loc="center right", bbox_to_anchor=(1.0, 0.45), handlelength=2.2)
fig.subplots_adjust(left=0.14, right=0.98, top=0.88, bottom=0.2, wspace=0.12)
save(fig, "Fig3")

# ------------------------------------------------------------------ Fig. 2 (plan Fig. 3): spectra of a label 1 and a label 0 path from one run (selection rule of the plan)
P = QC.load_paths()
agree = P[(P.N >= g("d5_min_rpm")) & P.acc0_nh_peak_ratio.notna() & P.s2_nh_peak_ratio.notna()]
cand = agree.groupby("run").filter(lambda d: set(d.label) >= {0, 1})
counts = cand[cand.label == 1].groupby("run").size()
run = int(sorted(counts[counts == counts.max()].index)[0])
sel = {}
for L in (1, 0):
    d = cand[(cand.run == run) & (cand.label == L)]
    med = d.acc0_rms.median()
    sel[L] = d.loc[(d.acc0_rms - med).abs().sort_values(kind="stable").index[0]]
V.set("fig3_run", run, "run", {"population": POP, "filter": "imi_vm20i runs holding both labels among paths with Acc0 and Sound2 spectra and spindle speed at least d5_min_rpm",
                                "statistic": "run with the most label 1 paths (ties: lower run number), the Fig. 3 rule"})
V.set("fig3_run_l1", int(counts.max()), "paths", {"population": POP, "filter": "that run", "statistic": "label 1 paths in the agreement population"})
for L in (1, 0):
    r = sel[L]
    V.set(f"fig3_l{L}_start_s", float(r.start), "s", {"population": POP, "filter": f"run fig3_run, label {L}", "statistic": "start time of the path whose Acc0 RMS is closest to the median Acc0 RMS of its label in the run (Fig. 3 rule)"})
    V.set(f"fig3_l{L}_rpm", float(r.N), "r/min", {"population": POP, "filter": f"run fig3_run, label {L}, the selected path", "statistic": "commanded spindle speed"})
    V.set(f"fig3_l{L}_tool", str(r.Tool), "text", {"population": POP, "filter": f"run fig3_run, label {L}, the selected path", "statistic": "tool"}, kind="text")
    V.set(f"fig3_l{L}_material", str(r.Workpiece), "text", {"population": POP, "filter": f"run fig3_run, label {L}, the selected path", "statistic": "workpiece material"}, kind="text")
z = np.load(f"{QC.SEG}/imi_vm20i_dataset{run}.npz", allow_pickle=True)
meta = json.loads(str(z["meta"]))
lo_hz, hi_hz, comb = float(g("d5_band_lo_hz")), float(g("d5_band_hi_hz")), int(g("d5_comb_bins"))
fig, axes = plt.subplots(2, 2, figsize=(174 * MM, g("fig3_height_mm") * MM), sharex=True)
for row, (ch, chname) in enumerate((("acc0", "Acc0"), ("s2", "Sound2"))):
    for col, L in enumerate((1, 0)):
        ax = axes[row, col]
        r = sel[L]
        j = int(np.argmin(np.abs(z["start"] - r.start) + np.abs(z["end"] - r.end)))
        Pw = z[f"{ch}_psd"][j].astype(float)
        dfh = meta["channels"][ch]["df"]
        f = np.arange(len(Pw)) * dfh
        band = (f >= lo_hz) & (f < hi_hz)
        fsp = r.N / 60.0
        idx = np.round(np.arange(1, int(10000.0 / fsp) + 2) * fsp / dfh).astype(int)
        harm = np.zeros(len(f), bool)
        for dd in range(-comb, comb + 1):
            ii = idx + dd; harm[ii[(ii >= 0) & (ii < len(f))]] = True
        ref = Pw[band].max()
        db = 10 * np.log10(np.maximum(Pw, 1e-30) / ref)
        for i0 in idx:
            a, b = (i0 - comb - 0.5) * dfh, (i0 + comb + 0.5) * dfh
            if b >= lo_hz and a <= hi_hz:
                ax.axvspan(a / 1000, b / 1000, color=GREY, linewidth=0, zorder=0)
        ax.plot(f[band] / 1000, db[band], color="black", linewidth=0.6, zorder=2)
        nh, hh = band & ~harm, band & harm
        on = Pw[nh].max() > Pw[hh].max()
        k = np.where(nh if on else hh)[0][np.argmax(Pw[nh if on else hh])]
        ax.plot(f[k] / 1000, db[k] + 3, marker="v" if on else "o", markersize=4.5, markerfacecolor=VERM if on else "white",
                markeredgecolor="black", markeredgewidth=0.6, linestyle="none", zorder=3)
        ax.set_xlim(lo_hz / 1000, hi_hz / 1000)
        ymin = {"acc0": -70, "s2": -110}[ch]
        assert db[band].min() > ymin, (ch, L, db[band].min())
        ax.set_ylim(ymin, 8)
        ax.set_yticks(list(range(ymin + 10, 1, 20)) if ch == "acc0" else [-100, -80, -60, -40, -20, 0])
        if col == 0:
            ax.set_ylabel(f"{chname} power (dB)")
        if row == 1:
            ax.set_xlabel("Frequency (kHz)")
        # white box behind the label so that the comb bands do not run through the lettering (final-gate inspection)
        ax.text(0.99, 0.97, f"Label {L}, {int(r.N):,} r/min", transform=ax.transAxes, ha="right", va="top", zorder=4,
                bbox=dict(facecolor="white", edgecolor="none", boxstyle="square,pad=0.15"))
        panel_label(ax, "abcd"[row * 2 + col], x=-0.03)
fig.subplots_adjust(left=0.085, right=0.99, top=0.95, bottom=0.11, hspace=0.18, wspace=0.17)
save(fig, "Fig2")

# ------------------------------------------------------------------ Fig. 4: AUC by evaluation scheme, scale of misordered pairs
# the row names are registered as text values and printed in Table 3 through the same tokens (review item R12)
rows = [("highrate", "High-rate RMS", {"random": "e4_highrate_auc_random", "run": "h1_auc_run", "family": "h1_auc_family", "material": "e7_highrate_auc_material"}),
        ("mtc", "MTConnect", {"random": "e4_mtc_auc_random", "run": "h2_auc_run", "family": "h2_auc_family", "material": "e7_mtc_auc_material"}),
        ("nos0", "High-rate, no Sound0", {"random": "e4_nos0_auc_random", "run": "e2_nos0_auc_run", "family": "e2_nos0_auc_family", "material": "e7_nos0_auc_material"}),
        ("clip", "Clipping only", {"random": "e4_clip_auc_random", "run": "e2_clip_auc_run", "family": "e2_clip_auc_family", "material": "e7_clip_auc_material"}),
        ("metanf", "Design, no family", {"random": "e4_metanf_auc_random", "run": "e3_metanf_auc_run", "family": "e4_metanf_auc_family"}),
        ("meta", "Design, with family", {"random": "e4_meta_auc_random", "run": "e3_meta_auc_run"})]
for tag, name, _ in rows:
    V.set(f"auc_row_{tag}", name, "text", {"population": "presentation", "filter": "Fig. 4 and Table 3", "statistic": "row name of the feature set, shared by the figure and the table"}, kind="text")
rows = [(name, cells) for _, name, cells in rows]
SCH = [("random", "Random paths", "s", SKY, "white"), ("run", "Runs held out", "o", BLUE, BLUE),
       ("family", "Families held out", "^", VERM, VERM), ("material", "Materials held out", "D", GREEN, "white")]
fig, ax = plt.subplots(figsize=(84 * MM, g("fig4_height_mm") * MM))
tr = lambda a: np.log10(np.clip(1 - np.asarray(a, float), 1e-4, None))
for i, (name, cells) in enumerate(rows):
    for k, (key, lab, mk, col, face) in enumerate(SCH):
        if key not in cells:
            continue
        n = cells[key]
        yv = i + (k - 1.5) * 0.17
        est_ = V4[n]["value"]; lo, hi = ci(n + "_ci")
        ax.plot([tr(lo), tr(hi)], [yv, yv], color="black", linewidth=0.8)
        ax.plot(tr(est_), yv, marker=mk, markersize=4, markerfacecolor=face, markeredgecolor=col if face == "white" else "black",
                markeredgewidth=0.8, linestyle="none", color=col)
ax.set_yticks(range(len(rows)))
ax.set_yticklabels([r[0] for r in rows])
ax.invert_yaxis()
ticks = [0.5, 0.9, 0.99, 0.999]
ax.set_xticks(tr(ticks))
ax.set_xticklabels([f"{t:g}" for t in ticks])
XMIN = 0.2   # left end of the axis; every interval must lie inside it with room to spare (review item R12)
lows = [ci(c[k] + "_ci")[0] for _, c in rows for k in c]
assert min(lows) > XMIN + 0.1, min(lows)
ax.set_xlim(tr(XMIN), tr(0.9995))
ax.set_xlabel("AUC (log scale of 1 minus AUC)")
ax.legend(handles=[Line2D([], [], marker=mk, markersize=4, markerfacecolor=face, markeredgecolor=col if face == "white" else "black",
                          linestyle="none", color=col, label=lab) for key, lab, mk, col, face in SCH],
          loc="lower center", bbox_to_anchor=(0.35, 1.0), ncol=2, handletextpad=0.3, columnspacing=0.8, borderaxespad=0.2)
fig.subplots_adjust(left=0.365, right=0.95, top=0.83, bottom=0.16)
save(fig, "Fig4")

# ------------------------------------------------------------------ Fig. 5: indicator combinations among label 1 paths by speed class (a); Acc0 line share by family (b)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(174 * MM, g("fig5_height_mm") * MM), gridspec_kw={"width_ratios": [1.4, 1]})
COMBO = [("both", "Both lines", VERM, ""), ("acc0_only", "Acc0 only", ORANGE, "////"), ("s2_only", "Sound2 only", SKY, "\\\\\\\\"), ("neither", "Neither", "white", "....")]
edge = int(g("speed_class_edge_rpm"))
for i, cls in enumerate(("lo", "hi")):
    n = sum(V4[f"h3_tab_{cls}_{c}"]["value"] for c, *_ in COMBO)
    bottom = 0.0
    for c, lab, col, hat in COMBO:
        v = V4[f"h3_tab_{cls}_{c}"]["value"] / n
        ax1.bar(i, v, 0.55, bottom=bottom, color=col, hatch=hat, edgecolor="black", linewidth=LINE)
        bottom += v
    ax1.text(i, 1.02, f"n = {n}", ha="center", va="bottom")
ax1.set_xticks([0, 1])
ax1.set_xticklabels([f"Below\n{edge:,} r/min", f"{edge:,} r/min\nand above"])
ax1.set_xlim(-0.6, 2.5)
ax1.set_ylim(0, 1.12)
ax1.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
ax1.set_ylabel("Share of label 1 paths")
ax1.legend(handles=[Patch(facecolor=col, hatch=hat, edgecolor="black", linewidth=LINE, label=lab) for c, lab, col, hat in COMBO][::-1],
           loc="center right", handlelength=1.6)
panel_label(ax1, "a", x=-0.06, y=1.02)
for i, (fk, lab) in enumerate((("a", "A"), ("e", "E"))):
    v = V4[f"e11_share_{fk}"]["value"]; lo, hi = ci(f"e11_share_{fk}_ci")
    ax2.plot([i, i], [lo, hi], color="black", linewidth=0.8)
    ax2.plot(i, v, marker="o" if fk == "a" else "s", markersize=5, markerfacecolor="black" if fk == "a" else "white",
             markeredgecolor="black", markeredgewidth=0.8, linestyle="none")
    ax2.text(i + 0.08, v, f"{V4[f'e11_n_{fk}']['value']} paths, {V4[f'e11_runs_{fk}']['value']} runs", va="center", ha="left")
ax2.set_xticks([0, 1])
ax2.set_xticklabels(["Family A", "Family E"])
ax2.set_xlim(-0.5, 2.1)
ax2.set_ylim(0, 1.0)
ax2.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
ax2.set_ylabel("Label 1 paths with an off-harmonic Acc0 line")
panel_label(ax2, "b", x=-0.1, y=1.04)
fig.subplots_adjust(left=0.07, right=0.99, top=0.9, bottom=0.17, wspace=0.28)
save(fig, "Fig5")
V.save()
print("figures written")