"""`bridge` figure — the histogram bridge reproduces Layer-1 E2 exactly (the HARD GATE).

The histogram model reduces Layer 1.5 to Layer 1 (MLE histogram + multinomial resampling),
so running it through the *neural* runner must reproduce E2's grounding phase boundary and its
exact `H_eq` closed form. Recovering g*≈0.047 here (vs Layer-1's 0.048) is what licenses every
later trained-model result to be read against the analytic core.

Four panels: (A) H trajectories (g=0 collapses, g>0 plateau); (B) the phase boundary —
stationary H vs g on the exact `H_eq` curve, with the operational g* + bootstrap CI; (C) tail
survival vs g; (D) per-rarity-band survival (the m·p*_i≳1 threshold). Reads only the bundle.

Usage:  python figures/plot_bridge.py [results/bridge]
"""

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 inheritance.analysis import critical_grounding, reduce_to_stationary  # noqa: E402
from inheritance.metrics import heterozygosity  # noqa: E402
from neural.config import SyntheticCfg  # noqa: E402
from neural.synthetic import make_mode_truth  # noqa: E402


def main(results_dir: str = "results/bridge") -> None:
    df, cfg = load_bundle(results_dir)
    n = cfg["dynamics"]["n"]
    syn = SyntheticCfg(**cfg["synthetic"])
    H_star = heterozygosity(make_mode_truth(syn).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))
    colors = plt.cm.viridis(np.linspace(0, 0.9, len(g_values)))

    # Panel A: H trajectories, one line per g.
    ax = axes[0, 0]
    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 on the exact H_eq curve + operational g* with bootstrap CI.
    ax = axes[0, 1]
    st = reduce_to_stationary(stat, value_col="heterozygosity", last_frac=1.0)
    gg, Hm, Hci = mean_ci(stat, "g", "heterozygosity")
    m_of_g = stat.groupby("g")["m"].first().to_numpy()
    nz = gg > 0
    ax.errorbar(gg[nz], Hm[nz], yerr=Hci[nz], fmt="o", color="#1f77b4", capsize=3, zorder=3,
                label="neural histogram runner")
    ax.plot(gg[~nz], Hm[~nz], "o", mfc="white", mec="#1f77b4", zorder=3)
    m_grid = np.linspace(0, m_of_g.max(), 400)
    ax.plot(m_grid / (n + m_grid), H_eq(m_grid), "k--", zorder=2, label=r"exact $H_{eq}$ (Layer 1)")
    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.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"Bridge reproduces the exact $H_{eq}$ and $g^\star$")
    ax.legend(frameon=False, fontsize=8)

    # Panel C: tail survival vs g (item-count and truth-mass).
    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")
    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 survival rises with g (deep tail lags)")
    ax.legend(frameon=False, fontsize=9)

    # Panel D: per-rarity-band survival across g.
    ax = axes[1, 1]
    band_cols = sorted(c for c in df.columns if c.startswith("band") and c.endswith("_alive"))
    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}" + (" (deepest)" if col == band_cols[-1] else
                                 " (shallowest)" if depth == "0" 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")
    ax.legend(frameon=False, fontsize=8)

    fig.suptitle("bridge — the histogram model reproduces Layer-1 E2 through the neural runner "
                 f"($g^*$={r['g_star']:.3f} vs Layer-1 0.048; HARD GATE passed)", y=1.0, fontsize=13)
    fig.tight_layout()
    savefig(fig, results_dir, "bridge")


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