"""E2 figure: the grounding phase boundary (headline), publication-honest.

Four panels: (A) H trajectories (g=0 slides to 0, g>0 plateau); (B) the phase boundary —
stationary H vs g tracking the exact H_eq, with an operational g* (where H first reaches
0.95·H*) and its bootstrap CI, and g=0 marked as a finite-time artifact; (C) tail coverage
by item-count vs truth-mass — both stay low, the deep tail is largely unrescuable at
feasible grounding; (D) per-rarity-band survival — the m·p*_i≳1 threshold made visible
(deep bands lag, motivating E4/E6). 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.analysis import critical_grounding, reduce_to_stationary  # noqa: E402
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"]
    td = make_true_distribution(cfg["truth"]["K"], 1, "zipf", cfg["truth"]["tail_frac"],
                                cfg["truth"]["zipf_s"], 0,
                                tail_threshold=cfg["truth"]["tail_threshold"])
    H_star = heterozygosity(td.p_star)

    def H_eq(m):
        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)
    stat = df[df["generation"] >= last]

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

    # Panel A: H trajectories, one line per g
    ax = axes[0, 0]
    colors = plt.cm.viridis(np.linspace(0, 0.9, len(g_values)))
    for g, c in zip(g_values, colors):
        s = df[df["g"] == g].groupby("generation")["heterozygosity"].mean()
        ax.plot(s.index, s.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 B: stationary H vs g + exact H_eq + operational g* with bootstrap CI
    ax = axes[0, 1]
    st = reduce_to_stationary(df[df["generation"] >= last], value_col="heterozygosity",
                              replicate_col="replicate", last_frac=1.0)
    gg, Hm, Hci = mean_ci(stat, "g", "heterozygosity")
    m_of_g = stat.groupby("g")["m"].first().to_numpy()
    # g=0 marked hollow (finite-time artifact: true H_eq(0)=0)
    nz = gg > 0
    ax.errorbar(gg[nz], Hm[nz], yerr=Hci[nz], fmt="o", color="#1f77b4", capsize=3,
                label="simulation", zorder=3)
    ax.plot(gg[~nz], Hm[~nz], "o", mfc="white", mec="#1f77b4", zorder=3)
    ax.annotate("g=0: pre-convergence\n(true $H_{eq}=0$)", (gg[~nz][0], Hm[~nz][0]),
                textcoords="offset points", xytext=(12, -4), fontsize=7, color="gray")
    m_grid = np.linspace(0, m_of_g.max(), 400)
    ax.plot(m_grid / (n + m_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)")
    r = critical_grounding(st, H_star=H_star, frac=0.95, seed=7)
    ax.axhline(r["target_H"], ls=":", color="#d62728", lw=1)
    ax.axvspan(r["ci_low"], r["ci_high"], color="#d62728", alpha=0.15)
    ax.axvline(r["g_star"], color="#d62728", lw=1.2,
               label=f"$g^*$={r['g_star']:.3f} (95% CI [{r['ci_low']:.3f},{r['ci_high']:.3f}])")
    ax.set(xlabel="grounding fraction $g=m/(n+m)$", ylabel="stationary $H$",
           title=r"Grounding saturates by $g^\star\approx0.05$ (95% of $H^*$)")
    ax.legend(frameon=False, fontsize=8)

    # Panel C: tail coverage by item-count vs truth-mass (both low: deep tail unrescuable)
    ax = axes[1, 0]
    ig, Im, Ici = mean_ci(stat, "g", "tail_frac_alive")
    mg, Mm, Mci = mean_ci(stat, "g", "tail_truth_mass_alive")
    ax.errorbar(ig, Im, yerr=Ici, fmt="s-", color="#d62728", capsize=3,
                label="tail items alive (count)")
    ax.errorbar(mg, Mm, yerr=Mci, fmt="o-", color="#9467bd", capsize=3,
                label="tail truth-mass alive")
    ax.set(xlabel="grounding fraction $g$", ylabel="fraction of tail retained",
           title="Tail stays largely unrescued at feasible g\n(rises with g; motivates E4/E6)")
    ax.legend(frameon=False, fontsize=9)

    # Panel D: per-rarity-band survival across g (band 0 = rarest)
    ax = axes[1, 1]
    band_cols = [c for c in df.columns if c.startswith("band") and c.endswith("_alive")]
    band_cols = sorted(band_cols)
    band_colors = plt.cm.plasma(np.linspace(0.1, 0.85, len(band_cols)))
    for col, c in zip(band_cols, band_colors):
        s = stat.groupby("g")[col].mean()
        depth = col.replace("band", "").replace("_alive", "")
        lab = f"band {depth}" + (" (rarest)" if depth == "0" else
                                 " (shallowest)" if col == band_cols[-1] else "")
        ax.plot(s.index, s.values, "-o", color=c, ms=4, label=lab)
    ax.set(xlabel="grounding fraction $g$", ylabel="fraction of band alive",
           title=r"Per-rarity band: the $m\,p^*_i\gtrsim1$ threshold (deep lags)")
    ax.legend(frameon=False, fontsize=8)

    fig.suptitle("E2 — a critical grounding ratio $g^\\star \\ll 1$ rescues diversity; "
                 "the deep tail needs recombination", y=1.0, fontsize=13)
    fig.tight_layout()
    savefig(fig, results_dir, "E2")


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