"""
Figures for PIN v5, drawn from the measured values recorded in each run's
output. Figure 6 carries the three-base numbers. No data is generated here and nothing is smoothed: every point is a
number that appears in the text, so a reader comparing a figure against a
table finds them identical.
"""

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

plt.rcParams.update({
    "font.family": "DejaVu Serif", "font.size": 9,
    "axes.grid": True, "grid.alpha": 0.25, "grid.linewidth": 0.5,
    "axes.spines.top": False, "axes.spines.right": False,
    "figure.dpi": 160, "savefig.bbox": "tight",
})
INK, ALT, MUT = "#1a1a1a", "#c1440e", "#6b7f8f"


# ---- 1. pooling, and the depth curve inverting -------------------------
depths = [1, 2, 3, 4, 5]
pool = {7:  ([.8657, .8617, .8632, .8627, .8588],
             [.8200, .8541, .8670, .8664, .8676]),
        14: ([.8798, .8766, .8766, .8790, .8818],
             [.7946, .8582, .8745, .8798, .8778]),
        28: ([.8685, .8699, .8752, .8748, .8784],
             [.6662, .7645, .8077, .8347, .8477])}
# each panel scales to its own data: at 7x7 the whole flat curve spans
# 0.007, which a shared axis wide enough for 28x28 renders as a flat line
fig, ax = plt.subplots(1, 3, figsize=(7.8, 2.6))
for a, g in zip(ax, (7, 14, 28)):
    flat, pooled = pool[g]
    a.plot(depths, flat, "o-", color=MUT, lw=1.4, ms=4, label="flat head")
    a.plot(depths, pooled, "s-", color=ALT, lw=1.6, ms=4, label="pooled")
    lo = min(min(flat), min(pooled)); hi = max(max(flat), max(pooled))
    pad = max(0.1*(hi-lo), 0.004)
    a.set_ylim(lo-pad, hi+pad)
    a.set_title(f"{g}x{g}", fontsize=9)
    a.set_xlabel("layers"); a.set_xticks(depths)
    a.tick_params(labelsize=8)
ax[0].set_ylabel("accuracy")
ax[1].legend(frameon=False, fontsize=8, loc="lower right")
ax[0].annotate("flat prefers shallow;\npooled rises to a\nplateau at three",
               (3, .8670), (1.12, .8300), fontsize=7.5, color=ALT,
               arrowprops=dict(arrowstyle="->", color=ALT, lw=0.8))
ax[2].text(1.2, .70, "4 channels here, so the\npooled head is 50 values:\n"
           "this panel is confounded", fontsize=7, color=MUT)
fig.suptitle("Pooling before the head, and the depth curve inverting",
             fontsize=10, y=1.04)
fig.savefig("fig1_pooling.png"); plt.close(fig)


# ---- 2. criticality --------------------------------------------------
fig, ax = plt.subplots(1, 2, figsize=(7.0, 2.6))
li = [2, 3, 4, 5]
ax[0].axhline(1, color=MUT, ls=":", lw=1)
ax[0].plot(li, [.9698, 1.2398, .9999, 1.5406], "o-", color=MUT, lw=1.4,
           ms=4, label="at initialisation")
ax[0].plot(li, [1.4709, 1.5563, 1.5360, 2.8993], "s-", color=ALT, lw=1.6,
           ms=4, label="after 40 epochs")
ax[0].set_xlabel("layer"); ax[0].set_ylabel("chi"); ax[0].set_xticks(li)
ax[0].set_title("a layer's expansion", fontsize=9)
ax[0].legend(frameon=False, fontsize=8)
ax[0].text(2.1, 2.55, "product\n1.85 -> 10.2", fontsize=7.5, color=ALT)

d = [3, 5, 7]
ax2 = ax[1]
ax2.plot(d, [2.608, 1.585, 1.309], "o-", color=MUT, lw=1.4, ms=4,
         label="per layer")
ax2.plot(d, [5.96, 5.78, 4.62], "s-", color=ALT, lw=1.6, ms=4,
         label="whole stack")
ax2.set_xlabel("depth"); ax2.set_ylabel("expansion"); ax2.set_xticks(d)
ax2.set_title("the stack regulates itself", fontsize=9)
ax2.legend(frameon=False, fontsize=8)
ax2.text(4.6, 3.6, "a fixed budget,\nspread thinner", fontsize=7.5,
         color=ALT)
fig.suptitle("Sensitivity through a folded stack", fontsize=10, y=1.04)
fig.savefig("fig2_criticality.png"); plt.close(fig)


# ---- 3. sample efficiency --------------------------------------------
N = [2, 5, 10, 20, 50, 100, 200, 500, 1000, 3000]
flat = [.9065, .9099, .9216, .9295, .9306, .9309, .9349, .9376, .9402, .9361]
pooled = [.8907, .9064, .9198, .9235, .9266, .9263, .9284, .9263, .9262, .9282]
fig, ax = plt.subplots(figsize=(4.4, 2.7))
ax.axhline(.9071, color=MUT, ls="--", lw=1)
ax.text(2.2, .9040, "the base, unaided", fontsize=7.5, color=MUT)
ax.semilogx(N, flat, "o-", color=ALT, lw=1.6, ms=4, label="3,136 features")
ax.semilogx(N, pooled, "s-", color=MUT, lw=1.4, ms=4, label="16, pooled")
ax.axvline(20, color=INK, ls=":", lw=0.9)
ax.text(21, .8955, "20 examples:\n68% of the gain", fontsize=7.5)
ax.set_xlabel("examples the member was fitted on")
ax.set_ylabel("accuracy on its task")
ax.legend(frameon=False, fontsize=8, loc="lower right")
ax.set_title("A member is cheap to fit", fontsize=10)
fig.savefig("fig3_sample_efficiency.png"); plt.close(fig)


# ---- 4. the weight as a dial -----------------------------------------
w = [0, .25, .5, .75, 1, 1.25, 1.5, 2, 3, 5, 8]
own = [.9071, .9356, .9442, .9434, .9444, .9452, .9429, .9407, .9384, .9328, .9222]
rest = [.8440, .7556, .6641, .5854, .4897, .3755, .2654, .0690, .0002, 0, 0]
conf = [.9480, .9529, .9558, .9561, .9482, .9405, .9359, .9424, .9616, .9712, .9792]
fig, ax = plt.subplots(figsize=(4.8, 2.8))
ax.plot(w, own, "o-", color=ALT, lw=1.6, ms=3.5, label="its own four classes")
ax.plot(w, rest, "s-", color=MUT, lw=1.4, ms=3.5, label="the other six")
ax.plot(w, conf, "^--", color=INK, lw=1.3, ms=3.5, label="confidence")
ax.set_xlabel("member weight w"); ax.set_ylabel("accuracy / confidence")
ax.set_ylim(-0.03, 1.02)
ax.legend(frameon=False, fontsize=8, loc="center right")
ax.annotate("blind to six classes,\n97% confident", (5, .002), (2.1, .18),
            fontsize=7.5, color=MUT,
            arrowprops=dict(arrowstyle="->", color=MUT, lw=0.8))
ax.set_title("Turning a member up: confidence never falls", fontsize=10)
fig.savefig("fig4_dial.png"); plt.close(fig)


# ---- 5. the confidence member and abstention -------------------------
cov = [.2, .4, .6, .8, 1.0]
fig, ax = plt.subplots(1, 2, figsize=(7.0, 2.6), sharey=True)
ax[0].plot(cov, [.5644, .5559, .6394, .6446, .5959], "o-", color=MUT,
           lw=1.4, ms=4, label="base confidence")
ax[0].plot(cov, [.9752, .9426, .8244, .6845, .5959], "s-", color=ALT,
           lw=1.6, ms=4, label="with the member")
ax[0].set_title("mixed input (AUC 0.4073 -> 0.8286)", fontsize=9)
ax[1].plot(cov, [1.0000, .9975, .9760, .9250, .8525], "o-", color=MUT,
           lw=1.4, ms=4, label="base confidence")
ax[1].plot(cov, [.9910, .9902, .9670, .9203, .8525], "s-", color=ALT,
           lw=1.6, ms=4, label="with the member")
ax[1].set_title("familiar input only (0.8681 -> 0.8447)", fontsize=9)
for a in ax:
    a.set_xlabel("coverage"); a.legend(frameon=False, fontsize=8,
                                       loc="lower left")
ax[0].set_ylabel("accuracy on what is answered")
fig.suptitle("A member supplies out-of-depth detection, and nothing else",
             fontsize=10, y=1.05)
fig.savefig("fig5_confidence.png"); plt.close(fig)


# ---- 6. qualification -------------------------------------------------
claims = [.1, .2, .3, .4, .5, .6, .8, 1.0]
fig, ax = plt.subplots(figsize=(4.6, 2.8))
for lab, ys, c, ls in (("w = 3", [.6244, .6417, .6481, .6471, .6416, .6342,
                                  .6215, .6096], MUT, "-"),
                       ("w = 5", [.6768, .7438, .7753, .7759, .7675, .7575,
                                  .7400, .7231], INK, "-"),
                       ("w = 8", [.6818, .7658, .8181, .8219, .8113, .7990,
                                  .7771, .7566], ALT, "-")):
    ax.plot(claims, ys, "o" + ls, color=c, lw=1.5, ms=3.5, label=lab)
ax.axhline(.5936, color=MUT, ls="--", lw=1)
ax.text(.62, .6000, "the base, unaided", fontsize=7.5, color=MUT)
ax.plot([.4], [.8209], "*", color=ALT, ms=13, zorder=5)
ax.annotate("claiming 40% beats\nclaiming everything\nby 0.0636",
            (.4, .8209), (.44, .655), fontsize=7.5, color=ALT,
            arrowprops=dict(arrowstyle="->", color=ALT, lw=0.8))
ax.set_xlabel("share of the task the member claims")
ax.set_ylabel("accuracy over the whole task")
ax.legend(frameon=False, fontsize=8, loc="upper left")
ax.set_title("Qualifying a member: concede, or override", fontsize=10)
fig.savefig("fig6_qualify.png"); plt.close(fig)

print("six figures written")