"""E2 figure: the grounding phase boundary (headline).

Shows that a critical grounding fraction g* << 1 separates collapse from a healthy
plateau: H trajectories (g=0 slides to 0, g>0 plateau), and stationary H / tail mass vs
g with the exact analytic H_eq overlaid. Usage:  python figures/plot_E2.py [results/E2]
"""

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

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 main(results_dir: str = "results/E2") -> None:
    df, cfg = load_bundle(results_dir)
    n = cfg["dynamics"]["n"]
    K, zs = cfg["truth"]["K"], cfg["truth"]["zipf_s"]
    td = make_true_distribution(K, 1, "zipf", cfg["truth"]["tail_frac"], zs, 0,
                                tail_threshold=cfg["truth"]["tail_threshold"])
    H_star = heterozygosity(td.p_star)

    def H_eq(m):  # exact stationary heterozygosity (blueprint 2.4-3)
        m = np.asarray(m, dtype=float)
        return np.where(m <= 0, 0.0, H_star * m * (2 * n + m - 1) / (n + 2 * n * m + m * m))

    g_values = sorted(df["g"].unique())
    last = int(cfg["generations"] * 0.8)   # stationary window: final 20% of generations

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

    # Panel 1: H trajectories, one line per g
    ax = axes[0]
    colors = plt.cm.viridis(np.linspace(0, 0.9, len(g_values)))
    for g, c in zip(g_values, colors):
        sub = df[df["g"] == g].groupby("generation")["heterozygosity"].mean()
        ax.plot(sub.index, sub.values, color=c, label=f"g={g:g}")
    ax.axhline(H_star, ls=":", color="gray", lw=1)
    ax.set(xlabel="generation", ylabel="heterozygosity $H$",
           title="Trajectories: g=0 collapses, g>0 plateau")
    ax.legend(frameon=False, fontsize=8, ncol=2)

    # Panel 2: stationary H vs g, with exact H_eq overlay
    stat = df[df["generation"] >= last]
    gg, Hm, Hci = mean_ci(stat, "g", "heterozygosity")
    m_of_g = stat.groupby("g")["m"].first().to_numpy()
    ax = axes[1]
    ax.errorbar(gg, Hm, yerr=Hci, fmt="o", color="#1f77b4", capsize=3,
                label="simulation (stationary)", zorder=3)
    m_grid = np.linspace(0, m_of_g.max(), 400)
    g_grid = m_grid / (n + m_grid)
    ax.plot(g_grid, H_eq(m_grid), "k--", label=r"exact $H_{eq}$", zorder=2)
    ax.axhline(H_star, ls=":", color="gray", lw=1, label="$H^*$ (truth)")
    ax.set(xlabel="grounding fraction $g=m/(n+m)$", ylabel="stationary $H$",
           title=r"Phase boundary: $g^\star \ll 1$")
    ax.legend(frameon=False, fontsize=9)

    # Panel 3: stationary fraction of TAIL ITEMS still alive vs g. (Aggregate tail *mass*
    # is a drift martingale and near-constant, so it is a poor indicator; the fraction of
    # rare items kept alive is the honest, monotone measure of how much tail grounding
    # rescues.) Tail-item survival rises steeply with g even where H is already saturated.
    tg, Tm, Tci = mean_ci(stat, "g", "tail_frac_alive")
    ax = axes[2]
    ax.errorbar(tg, Tm, yerr=Tci, fmt="s", color="#d62728", capsize=3)
    ax.set(xlabel="grounding fraction $g$",
           ylabel="fraction of tail items alive",
           title="Grounding keeps rare items alive")

    fig.suptitle("E2 — a critical grounding ratio $g^\\star \\ll 1$ separates ratchet "
                 "from collapse", y=1.02)
    fig.tight_layout()
    savefig(fig, results_dir, "E2")


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