"""`kernel` figure — the learning kernel: neutral drift fails both neural models, oppositely.

The Layer-1.5 architecture-generality result gets a mechanistic explanation. Neutral Wright-Fisher
(the histogram / Riis baseline) is the null; a real learner adds an estimator bias that can point
either way. Two regimes, each falsifying neutral drift in the OPPOSITE direction, each repaired by
one knob of the learning kernel:

* **VAE regime (n=6000, K=30):** drift is nearly inert — neutral holds ~all modes — yet the real
  VAE collapsed to a single mode. Sharpening (temperature<1) reproduces it. The estimator ADDS
  collapse pressure.
* **RNN regime (n=200, K=256):** neutral drives H to 0, but the real RNN only partially collapses
  (H floors, forward-KL plateaus). Mutation-toward-prior (reset>0) reproduces the floor. The
  estimator REMOVES collapse pressure.

Analytic arms are read from results/kernel_{sharpen,smooth}; the neural reference endpoints
(dashed) are read from the committed results/mnist_collapse and results/grounding parquets — so the
figure is a pure function of committed artifacts.

Usage:  python figures/plot_kernel.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from _figlib import load_bundle, savefig, letter_axes  # noqa: E402

sys.path.insert(0, str(Path(__file__).parents[1] / "src"))
from knowledge.metrics import heterozygosity  # noqa: E402
from knowledge.truth import make_true_distribution  # noqa: E402


def _mean_traj(df, knob, val, col):
    sub = df[df[knob] == val]
    s = sub.groupby("generation")[col].mean()
    return s.index.to_numpy(), s.to_numpy()


def _neural_dry(results_dir, col, stationary_frac=0.0):
    """Mean of ``col`` on the dry (g=0) neural arm — endpoint, or stationary tail if frac>0."""
    df, _ = load_bundle(results_dir)
    dry = df[df["g"] == 0.0]
    if stationary_frac > 0:
        dry = dry[dry["generation"] >= int(dry["generation"].max() * (1 - stationary_frac))]
        return float(dry[col].mean())
    return float(dry[dry["generation"] == dry["generation"].max()][col].mean())


def main() -> None:
    sh, sh_cfg = load_bundle("results/kernel_sharpen")
    sm, sm_cfg = load_bundle("results/kernel_smooth")
    Hstar_sh = heterozygosity(make_true_distribution(
        sh_cfg["truth"]["K"], 1, "zipf", 0.5, sh_cfg["truth"]["zipf_s"], 0,
        tail_threshold=sh_cfg["truth"]["tail_threshold"]).p_star)
    Hstar_sm = heterozygosity(make_true_distribution(
        sm_cfg["truth"]["K"], 1, "zipf", 0.5, sm_cfg["truth"]["zipf_s"], 0,
        tail_threshold=sm_cfg["truth"]["tail_threshold"]).p_star)

    # Neural reference endpoints (dashed) from the committed neural runs.
    vae_H = _neural_dry("results/mnist_collapse", "heterozygosity")
    vae_sup = _neural_dry("results/mnist_collapse", "support_size")
    rnn_H = _neural_dry("results/grounding", "heterozygosity", stationary_frac=0.4)
    rnn_KL = _neural_dry("results/grounding", "forward_kl", stationary_frac=0.4)

    fig, axes = plt.subplots(2, 2, figsize=(13, 9))
    NEU, KER = "#1f77b4", "#d62728"

    # --- VAE regime: sharpening ---
    ax = axes[0, 0]
    for val, c, lab in [(1.0, NEU, "neutral (τ=1)"), (0.8, KER, "sharpened (τ=0.8)")]:
        g, y = _mean_traj(sh, "temperature", val, "heterozygosity")
        ax.plot(g, y, "-o", color=c, ms=3, label=lab)
    ax.axhline(Hstar_sh, ls=":", color="gray", lw=1, label="$H^*$")
    ax.axhline(vae_H, ls="--", color="#2ca02c", lw=1.3, label=f"real VAE (no real data): {vae_H:.2f}")
    ax.set(xlabel="generation", ylabel="heterozygosity $H$",
           title="VAE regime ($n$=6000, $K$=30): neutral drift is inert;\nsharpening collapses (like the VAE)")
    ax.legend(frameon=False, fontsize=8)

    ax = axes[0, 1]
    for val, c, lab in [(1.0, NEU, "neutral (τ=1)"), (0.8, KER, "sharpened (τ=0.8)")]:
        g, y = _mean_traj(sh, "temperature", val, "support_size")
        ax.plot(g, y, "-o", color=c, ms=3, label=lab)
    ax.axhline(vae_sup, ls="--", color="#2ca02c", lw=1.3, label=f"real VAE (no real data): {vae_sup:.0f}")
    ax.set(xlabel="generation", ylabel="distinct modes alive",
           title="Support: neutral holds ~all; sharpening → 1 mode")
    ax.legend(frameon=False, fontsize=8)

    # --- RNN regime: smoothing ---
    ax = axes[1, 0]
    for val, c, lab in [(0.0, NEU, "neutral (u=0)"), (0.006, KER, "smoothed (u=0.006)")]:
        g, y = _mean_traj(sm, "reset", val, "heterozygosity")
        ax.plot(g, y, "-", color=c, lw=1.8, label=lab)
    ax.axhline(Hstar_sm, ls=":", color="gray", lw=1, label="$H^*$")
    ax.axhline(rnn_H, ls="--", color="#2ca02c", lw=1.3, label=f"real RNN (no real data): {rnn_H:.2f}")
    ax.set(xlabel="generation", ylabel="heterozygosity $H$",
           title="RNN regime ($n$=200, $K$=256): neutral → 0;\nsmoothing floors $H$ (like the RNN)")
    ax.legend(frameon=False, fontsize=8)

    ax = axes[1, 1]
    for val, c, lab in [(0.0, NEU, "neutral (u=0)"), (0.006, KER, "smoothed (u=0.006)")]:
        g, y = _mean_traj(sm, "reset", val, "forward_kl")
        ax.plot(g, y, "-", color=c, lw=1.8, label=lab)
    ax.axhline(rnn_KL, ls="--", color="#2ca02c", lw=1.3, label=f"real RNN (no real data): {rnn_KL:.1f}")
    ax.set(xlabel="generation", ylabel=r"forward-KL $D(p^*\Vert p)$",
           title="Forward-KL: neutral diverges; smoothing plateaus\n(overshoots RNN → prior is truth-like, not uniform)")
    ax.legend(frameon=False, fontsize=8)

    fig.tight_layout()
    for d in ("results/kernel_sharpen", "results/kernel_smooth"):
        letter_axes(fig)
        savefig(fig, d, "kernel")


if __name__ == "__main__":
    main()