#!/usr/bin/env python3
"""Generate all paper figures from real experiment artifacts."""
import json
import sys

import numpy as np

import os as _os
ROOT = _os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))
ART = f"{ROOT}/artifacts"
sys.path.insert(0, ROOT)

import matplotlib
matplotlib.use("Agg")
import matplotlib.font_manager as fm
fm.fontManager.addfont('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf')
import matplotlib.pyplot as plt
plt.rcParams.update({
    "font.family": "DejaVu Sans",
    "font.size": 9,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.grid": True,
    "grid.alpha": 0.25,
    "figure.dpi": 200,
})

OUT = f"{ART}/figures"
import os
os.makedirs(OUT, exist_ok=True)

metrics = json.load(open(f"{ART}/metrics.json"))
ACCENT = "#2e7d5b"
ACCENT2 = "#b8860b"
GRAY = "#666666"

# ---------------------------------------------------------------- fig: EPT
from spira import SPIRA
model = SPIRA.load(f"{ART}/spira_model.npz")
sample = (b"It was the best of times, it was the worst of times, it was the age "
          b"of wisdom, it was the age of foolishness, it was the epoch of belief, "
          b"it was the epoch of incredulity, it was the season of Light, "
          b"it was the season of Darkness")
sur = model.ept.surprise(sample)
patches = model.ept.segment(sample)
fig, axes = plt.subplots(2, 1, figsize=(7.0, 3.2), height_ratios=[1.2, 1],
                         constrained_layout=True, sharex=True)
axes[0].fill_between(range(len(sur)), sur, color=ACCENT, alpha=0.55, step="mid")
axes[0].plot(range(len(sur)), sur, color=ACCENT, lw=0.8)
axes[0].set_ylabel("surprisal (bits)")
axes[0].set_title("EPT: per-byte surprisal and constant-information patch boundaries", fontsize=9)
b = model.ept.boundaries(sample)
for bx in b[:-1]:
    axes[0].axvline(bx, color=ACCENT2, lw=0.7, alpha=0.85)
# patch length bars
axes[1].bar(b[:-1], [b[i+1]-b[i] for i in range(len(b)-1)], width=1.0,
            color=ACCENT, alpha=0.7, align="edge")
axes[1].axhline(float(np.mean([b[i+1]-b[i] for i in range(len(b)-1)])),
                color=ACCENT2, lw=1, ls="--", label=f"mean {np.mean([b[i+1]-b[i] for i in range(len(b)-1)]):.1f} B")
axes[1].set_xlabel("byte position")
axes[1].set_ylabel("patch length (B)")
axes[1].legend(frameon=False, fontsize=8)
fig.savefig(f"{OUT}/fig_ept.pdf")
plt.close(fig)
print("fig_ept done:", len(patches), "patches, avg", np.mean([len(p) for p in patches]))

# ------------------------------------------------------- fig: compression
comp = metrics["compression"]
methods = ["gzip -9", "BPE-4k\n(unigram)", "BPE-4k\n(bigram)", "BPE-16k\n(bigram)", "EPT\n(bigram)", "EPT\n(unigram)"]
vals = [comp["gzip"]["bits_per_byte"],
        comp["bpe_4096"]["bits_per_byte_unigram"],
        comp["bpe_4096"]["bits_per_byte_bigram"],
        comp["bpe_16384"]["bits_per_byte_bigram"],
        comp["ept"]["bits_per_byte_bigram"],
        comp["ept"]["bits_per_byte_unigram"]]
colors = [GRAY, "#9aa5b1", "#9aa5b1", "#9aa5b1", ACCENT, ACCENT]
fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True)
bars = ax.bar(range(len(vals)), vals, color=colors, width=0.65)
for i, v in enumerate(vals):
    ax.text(i, v + 0.08, f"{v:.2f}", ha="center", fontsize=8)
ax.set_xticks(range(len(methods)))
ax.set_xticklabels(methods, fontsize=7.5)
ax.set_ylabel("bits per byte (with dictionary)")
ax.set_title("Coded size on the held-out test stream", fontsize=9)
fig.savefig(f"{OUT}/fig_compression.pdf")
plt.close(fig)
print("fig_compression done")

# -------------------------------------------------- fig: one-shot+corrupt
one = metrics["one_shot"]
corr = metrics["correction"]
corr_sp = metrics.get("correction_sparse_am", {})
ks = [0, 1, 2, 3]
am_curve = [corr[f"corrupt_{k}_of_4"]["am_acc"] for k in ks]
fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True)
ax.plot(ks, [v * 100 for v in am_curve], "o-", color=ACCENT, lw=1.6,
        label="full hippocampus (46% load)")
if corr_sp:
    am_sp = [corr_sp[f"corrupt_{k}_of_4"]["am_acc"] for k in ks]
    ax.plot(ks, [v * 100 for v in am_sp], "s--", color=ACCENT2, lw=1.4,
            label="sparse hippocampus (book only)")
    for k, v in zip(ks, am_sp):
        ax.annotate(f"{v*100:.1f}%", (k, v * 100), textcoords="offset points",
                    xytext=(0, 7), ha="center", fontsize=8, color=ACCENT2)
for k, v in zip(ks, am_curve):
    ax.annotate(f"{v*100:.1f}%", (k, v * 100), textcoords="offset points",
                xytext=(0, -13), ha="center", fontsize=8)
ax.set_xlabel("corrupted context slots (of 4)")
ax.set_ylabel("exact recall (%)")
ax.set_ylim(-14, 110)
ax.set_xticks(ks)
ax.set_title("Pattern completion under context corruption\n(completion trades against hippocampal load)", fontsize=9)
ax.legend(frameon=False, fontsize=8, loc="center right")
fig.savefig(f"{OUT}/fig_correction.pdf")
plt.close(fig)
print("fig_correction done:", am_curve)

# ------------------------------------------------- fig: vocab coverage
vc = metrics["vocab_coverage_curve"]
vs = [c["vocab"] for c in vc]
ct = [c["cov_test"] * 100 for c in vc]
co = [c["cov_oneshot"] * 100 for c in vc]
fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True)
ax.plot(vs, ct, "o-", color=ACCENT, lw=1.6, ms=3.5, label="held-out test")
ax.plot(vs, co, "s-", color=ACCENT2, lw=1.4, ms=3.5, label="unseen book (one-shot)")
ax.axvline(4096, color=GRAY, lw=0.9, ls=":")
ax.text(4096 * 1.06, 88, "v1: 4096", fontsize=7.5, color=GRAY)
ax.axvline(16384, color=GRAY, lw=0.9, ls=":")
ax.text(16384 * 1.06, 88, "v2: 16384", fontsize=7.5, color=GRAY)
ax.annotate(f"{ct[-1]:.1f}%", (vs[-1], ct[-1]), textcoords="offset points",
            xytext=(2, 6), fontsize=8)
ax.set_xscale("log", base=2)
ax.set_xticks(vs)
ax.set_xticklabels([str(v) for v in vs], fontsize=7.5)
ax.set_xlabel("readout vocabulary size (top-V by train frequency)")
ax.set_ylabel("held-out target coverage (%)")
ax.set_title("Zipf knee of the readout vocabulary\n(constant-information patches flatten the tail)", fontsize=9)
ax.legend(frameon=False, fontsize=8, loc="lower right")
fig.savefig(f"{OUT}/fig_vocab.pdf")
plt.close(fig)
print("fig_vocab done")

# ------------------------------------------------- fig: replay consolidation
rep = metrics["replay"]
xs = [0] + [c["replayed_tokens"] for c in rep["curve"]]
ys = [rep["before"]["hier_acc"] * 100] + [c["hier_acc"] * 100 for c in rep["curve"]]
ysi = [rep["before"]["hier_acc_invocab"] * 100] + [c["hier_acc_invocab"] * 100 for c in rep["curve"]]
fig, ax = plt.subplots(figsize=(4.9, 3.0), constrained_layout=True)
ax.plot(xs, ys, "o-", color=ACCENT, lw=1.6, ms=3.5, label="cortex-only recall")
ax.plot(xs, ysi, "s--", color=ACCENT2, lw=1.4, ms=3.2, label="cortex-only, in-vocab only")
ax.axhline(rep["am_recall_after_replay"]["am_acc"] * 100, color=GRAY, lw=1.1, ls=":")
ax.text(xs[-1] * 0.06, rep["am_recall_after_replay"]["am_acc"] * 100 - 4.5,
        f"hippocampus (unchanged): {rep['am_recall_after_replay']['am_acc']*100:.1f}%",
        fontsize=7.5, color=GRAY)
ax.axhline(rep["after"]["hier_acc"] * 100, color=ACCENT, lw=0.8, ls="--", alpha=0.5)
ax.annotate(f"full eval after replay: {rep['after']['hier_acc']*100:.1f}%\n"
            f"(in-vocab {rep['after']['hier_acc_invocab']*100:.1f}%)",
            (xs[-1], rep["after"]["hier_acc"] * 100), textcoords="offset points",
            xytext=(-96, 6), fontsize=7.5, color=ACCENT)
ax.set_xlabel("hippocampal tokens replayed through the cortex")
ax.set_ylabel("exact next-patch recall (%)")
ax.set_title("Systems consolidation: the cortex learns the book\nonly through hippocampal replay (never seen awake)", fontsize=9)
ax.legend(frameon=False, fontsize=8, loc="center right")
fig.savefig(f"{OUT}/fig_replay.pdf")
plt.close(fig)
print("fig_replay done")

# ------------------------------------------------------------ fig: raster
z = np.load(f"{ART}/spira_model.npz.raster.npz")
ticks, neurons = z["ticks"], z["neurons"]
fig, ax = plt.subplots(figsize=(6.8, 2.7), constrained_layout=True)
ax.scatter(ticks, neurons, s=0.55, c=ACCENT, linewidths=0)
ax.axhline(512 * 0.8, color=ACCENT2, lw=0.8, ls="--")
ax.text(5, 512 * 0.8 + 6, "inhibitory population (Dale's law)", fontsize=7, color=ACCENT2)
ax.set_xlabel("tick (token)")
ax.set_ylabel("reservoir neuron")
ax.set_title(f"Spike raster — {len(ticks)} spikes over 1500 ticks "
             f"({len(ticks)/1500:.0f} spikes/tick)", fontsize=9)
fig.savefig(f"{OUT}/fig_raster.pdf")
plt.close(fig)
print("fig_raster done")

# ------------------------------------------------- fig: learning curve
import glob
logs = []
for line in open(f"{ART}/chunks.jsonl"):
    try:
        logs.append(json.loads(line))
    except Exception:
        pass
# training acc EMAs were logged inside msg lines; rebuild a coarse curve from
# chunk results recorded in the metrics/history + chunks.jsonl messages
accs = []
for e in logs:
    if "acc=" in e.get("msg", ""):
        try:
            accs.append(float(e["msg"].split("acc=")[1].split()[0]))
        except Exception:
            pass
fig, ax = plt.subplots(figsize=(4.6, 2.9), constrained_layout=True)
if accs:
    xs = np.arange(len(accs)) * 25000
    ax.plot(xs, [a * 100 for a in accs], "-", color=ACCENT, lw=1.4)
    ax.set_xlabel("tick (token)")
    ax.set_ylabel("online next-patch accuracy (%)")
    ax.set_title("Online learning during training (vocab-masked, incl. UNK targets)", fontsize=9)
fig.savefig(f"{OUT}/fig_learning.pdf")
plt.close(fig)
print("fig_learning done:", len(accs), "points")