"""`mnist_collapse` figure — collapse and grounding-rescue on REAL MNIST images.

External validity for the neural tier: a convolutional VAE retrained each generation on its own
generated digits collapses — rare (class, thickness) modes die, forward-KL to the Zipf truth
climbs, support shrinks — while a grounded arm (a fraction of fresh real MNIST images each
generation) holds the tail. Modes are read by a frozen CNN oracle whose mode accuracy (the
measurement-noise floor) is annotated from the run manifest. Signs, not magnitudes.

Four panels, dry (g=0) vs grounded, mean ± 95% CI across replicates: (A) forward-KL trajectories;
(B) support size (distinct modes alive); (C) tail truth-mass alive; (D) heterozygosity. Reads the
committed bundle (parquet) + manifest.json only.

Usage:  python figures/plot_mnist.py [results/mnist_collapse]
"""

from __future__ import annotations

import json
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 neural.config import MnistCfg  # noqa: E402
from neural.mnist_data import make_mnist_truth  # noqa: E402


def _traj(df, g, col):
    """Return (generations, mean, 95% half-width) of ``col`` for arm ``g`` across replicates."""
    sub = df[df["g"] == g]
    grp = sub.groupby("generation")[col]
    gens = np.array(sorted(sub["generation"].unique()))
    return gens, grp.mean().to_numpy(), 1.96 * grp.sem().to_numpy()


def main(results_dir: str = "results/mnist_collapse") -> None:
    df, cfg = load_bundle(results_dir)
    syn = MnistCfg(**cfg["mnist"])
    H_star = heterozygosity(make_mnist_truth(syn).p_star)
    manifest = json.loads((Path(results_dir) / "manifest.json").read_text())
    oracle_acc = manifest.get("oracle_mode_accuracy", float("nan"))

    g_dry, g_wet = min(df["g"].unique()), max(df["g"].unique())
    arms = [(g_dry, "#d62728", f"no real data (g={g_dry:g})"), (g_wet, "#2ca02c", f"grounded (g={g_wet:g})")]

    fig, axes = plt.subplots(2, 2, figsize=(13, 9))

    def panel(ax, col, title, ylabel, hline=None):
        for g, c, lab in arms:
            gens, m, ci = _traj(df, g, col)
            ax.plot(gens, m, "-o", color=c, ms=3, label=lab)
            ax.fill_between(gens, m - ci, m + ci, color=c, alpha=0.2)
        if hline is not None:
            ax.axhline(hline[0], ls=":", color="gray", lw=1, label=hline[1])
        ax.set(xlabel="generation", ylabel=ylabel, title=title)
        ax.legend(frameon=False, fontsize=9)

    panel(axes[0, 0], "forward_kl", "Without real data forward-KL climbs; grounding holds it",
          r"forward-KL $D(p^*\Vert\hat p)$")
    panel(axes[0, 1], "support_size", f"Support collapses (of K={syn.K} modes)",
          "distinct modes alive", hline=(syn.K, f"$K$={syn.K}"))
    panel(axes[1, 0], "tail_truth_mass_alive", "Rare tail dies without real data, held by grounding",
          "tail truth-mass alive")
    panel(axes[1, 1], "heterozygosity", "Diversity collapses without real data, held by grounding",
          "heterozygosity $H$", hline=(H_star, "$H^*$"))

    fig.tight_layout()
    letter_axes(fig)
    savefig(fig, results_dir, "mnist_collapse")


if __name__ == "__main__":
    main(*sys.argv[1:])