"""E1 figure: reproduce collapse (null model).

Shows tail-first collapse under pure neutral drift: geometric H decay matching the
analytic law, tail items dying faster than head items, support -> 1 and forward-KL
diverging. Usage:  python figures/plot_E1.py [results/E1]
"""

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, mean_ci, savefig  # noqa: E402


def main(results_dir: str = "results/E1") -> None:
    df, cfg = load_bundle(results_dir)
    n = cfg["dynamics"]["n"]

    gens, Hmean, Hci = mean_ci(df, "generation", "heterozygosity")
    H0 = Hmean[0]
    analytic = H0 * (1.0 - 1.0 / n) ** gens

    m = df.groupby("generation").mean(numeric_only=True)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))

    # Panel 1: heterozygosity decay vs the analytic law
    ax = axes[0]
    ax.plot(gens, Hmean, color="#1f77b4", label="simulation (mean)")
    ax.fill_between(gens, Hmean - Hci, Hmean + Hci, color="#1f77b4", alpha=0.25)
    ax.plot(gens, analytic, "k--", label=r"$H_0(1-1/n)^t$")
    ax.set(xlabel="generation", ylabel="heterozygosity $H$",
           title=f"Geometric decay (n={n})")
    ax.legend(frameon=False)

    # Panel 2: tail-first — fraction of tail vs head items still alive
    ax = axes[1]
    ax.plot(m.index, m["tail_frac_alive"], color="#d62728", label="tail items alive")
    ax.plot(m.index, m["head_frac_alive"], color="#2ca02c", label="head items alive")
    ax.set(xlabel="generation", ylabel="fraction of items surviving",
           title="Tail dies first", yscale="log")
    ax.legend(frameon=False)

    # Panel 3: support collapse and KL divergence
    ax = axes[2]
    ax.plot(m.index, m["support_size"], color="#9467bd", label="support size")
    ax.set(xlabel="generation", ylabel="support size", yscale="log", title="Collapse")
    ax2 = ax.twinx()
    ax2.plot(m.index, m["forward_kl"], color="#ff7f0e", label="forward KL")
    ax2.set_ylabel(r"forward KL $D_{KL}(p^*\,\|\,p_t)$", color="#ff7f0e")
    ax.legend(loc="center right", frameon=False)

    fig.suptitle("E1 — distillation without grounding collapses, tail first", y=1.02)
    fig.tight_layout()
    savefig(fig, results_dir, "E1")


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