"""E6 figure: the re-minting gate and irreversibility.

Re-minting freezes the current distribution as the new grounding reference and discards
the original truth. Re-minting a collapsed lineage locks in the collapse: KL to the
original truth diverges, because the lost original tails can no longer be grounded.
Gating re-mint on diversity refuses to re-mint while collapsed and keeps KL bounded;
re-minting a healthy lineage is harmless. Usage:  python figures/plot_figS3_rebaselining.py [results/figS3_rebaselining]
"""

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

STYLE = {
    "healthy_remint":     ("#2ca02c", "re-mint while healthy (H high)"),
    "collapsed_remint":   ("#d62728", "re-mint while collapsed (ungated)"),
    "collapsed_gated":    ("#1f77b4", "collapsed + diversity gate"),
    "collapsed_noremint": ("#7f7f7f", "collapsed, no re-mint (baseline)"),
}


def main(results_dir: str = "results/figS3_rebaselining") -> None:
    df, cfg = load_bundle(results_dir)
    period = cfg["dynamics"]["remint"]["period"]
    G = cfg["generations"]
    remint_gens = list(range(period, G + 1, period))

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

    # Panel 1: forward KL to the ORIGINAL truth
    ax = axes[0]
    for arm, (c, lab) in STYLE.items():
        s = df[df["arm"] == arm].groupby("generation")["forward_kl"]
        mean, sem = s.mean(), s.sem()
        ax.plot(mean.index, mean.values, color=c, label=lab)
        ax.fill_between(mean.index, mean - 1.96 * sem, mean + 1.96 * sem,
                        color=c, alpha=0.15)
    for g in remint_gens:
        ax.axvline(g, ls=":", color="k", lw=0.8, alpha=0.5)
    ax.set(xlabel="generation", ylabel=r"forward KL to original truth",
           title="Re-minting while collapsed locks in divergence")
    ax.legend(frameon=False, fontsize=8)

    # Panel 2: heterozygosity (which arms are collapsed; gate reads this)
    ax = axes[1]
    gate = None
    for arm, (c, lab) in STYLE.items():
        s = df[df["arm"] == arm].groupby("generation")["heterozygosity"].mean()
        ax.plot(s.index, s.values, color=c, label=lab)
    # draw the gate threshold used by the gated arm
    for v in cfg["sweep"][0]["values"]:
        if v["name"] == "collapsed_gated":
            gate = v["set"].get("dynamics.remint.H_gate")
    if gate is not None:
        ax.axhline(gate, ls="--", color="k", lw=1)
        ax.annotate(f"gate H={gate}", (G * 0.02, gate + 0.02), fontsize=8)
    for g in remint_gens:
        ax.axvline(g, ls=":", color="k", lw=0.8, alpha=0.5)
    ax.set(xlabel="generation", ylabel="heterozygosity $H$",
           title="Diversity at re-mint time (the gate reads this)")
    ax.legend(frameon=False, fontsize=8)

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


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