"""E3 figure: region-matched grounding.

Shows that grounding must *overlap* the content it protects. At the same total budget,
uniform grounding spreads thin and lets the exercised region's tail collapse, while
matched grounding concentrates on that region and keeps its rare items alive (at the cost
of the regions it does not touch). Usage:  python figures/plot_E3.py [results/E3]

Metric: per-region tail-item survival. (Per-region *heterozygosity* is confounded by
region mass under matched grounding, so it is deliberately not used here.)
"""

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


def main(results_dir: str = "results/E3") -> None:
    df, cfg = load_bundle(results_dir)
    R = cfg["truth"]["R"]
    exercised = cfg["dynamics"]["grounding"]["exercised"]
    target = exercised[0]
    last = int(cfg["generations"] * 0.8)
    colors = {"uniform": "#d62728", "matched": "#1f77b4"}

    fig, axes = plt.subplots(1, 2, figsize=(12, 4.4))

    # Panel 1: tail survival of the target region over generations
    ax = axes[0]
    tcol = f"tailalive_region_{target}"
    for pol in ("uniform", "matched"):
        sub = df[df["policy"] == pol].groupby("generation")[tcol]
        mean = sub.mean()
        sem = sub.sem()
        ax.plot(mean.index, mean.values, color=colors[pol], label=pol)
        ax.fill_between(mean.index, mean - 1.96 * sem, mean + 1.96 * sem,
                        color=colors[pol], alpha=0.2)
    ax.set(xlabel="generation",
           ylabel=f"tail items alive in region {target}",
           title=f"Target region {target} (exercised): matched holds, uniform collapses")
    ax.legend(frameon=False)

    # Panel 2: stationary tail survival per region, uniform vs matched
    ax = axes[1]
    stat = df[df["generation"] >= last]
    regions = np.arange(R)
    width = 0.4
    for i, pol in enumerate(("uniform", "matched")):
        vals = [stat[stat["policy"] == pol][f"tailalive_region_{r}"].mean()
                for r in regions]
        ax.bar(regions + (i - 0.5) * width, vals, width,
               color=colors[pol], label=pol)
    ax.axvline(target, ls=":", color="gray", lw=1)
    ax.annotate("exercised", (target, ax.get_ylim()[1] * 0.9), fontsize=8,
                ha="center", color="gray")
    ax.set(xlabel="region", ylabel="stationary tail items alive",
           title="Uniform spreads thin; matched concentrates on the exercised region",
           xticks=regions)
    ax.legend(frameon=False)

    fig.suptitle("E3 — grounding must overlap the content it protects", y=1.02)
    fig.tight_layout()
    savefig(fig, results_dir, "E3")


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