"""E5 figure: quality-diversity vs greedy selection.

At matched grounding, greedy (directional) selection drives the lineage toward the
fittest items and collapses diversity, while quality-diversity selection (a novelty bonus
w_i ∝ f_i·p_i^{-alpha}) maintains a high stationary heterozygosity that rises with the
novelty exponent alpha. Usage:  python figures/plot_E5.py [results/E5]
"""

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


def main(results_dir: str = "results/E5") -> None:
    df, cfg = load_bundle(results_dir)
    last = int(cfg["generations"] * 0.8)

    def arm(mode, alpha=1.0):
        return df[(df["mode"] == mode) & (df["novelty_alpha"] == alpha)]

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

    # Panel 1: H trajectories
    ax = axes[0]
    series = [("greedy", 1.0, "#d62728", "greedy"),
              ("qd", 1.0, "#ff7f0e", "quality-diversity (α=1)"),
              ("qd", 2.0, "#1f77b4", "quality-diversity (α=2)"),
              ("none", 1.0, "#2ca02c", "none (grounding only)")]
    for mode, a, c, lab in series:
        s = arm(mode, a).groupby("generation")["heterozygosity"].mean()
        ax.plot(s.index, s.values, color=c, label=lab)
    ax.set(xlabel="generation", ylabel="heterozygosity $H$",
           title="Greedy collapses;\nquality-diversity maintains diversity")
    ax.legend(frameon=False, fontsize=8)

    # Panel 2: stationary H vs alpha for qd, with greedy/none reference lines
    ax = axes[1]
    qd = df[(df["mode"] == "qd") & (df["generation"] >= last)]
    st = qd.groupby("novelty_alpha")["heterozygosity"].agg(["mean", "sem"])
    ax.errorbar(st.index, st["mean"], yerr=1.96 * st["sem"], fmt="-o",
                color="#ff7f0e", capsize=3, label="quality-diversity")
    for mode, c in (("greedy", "#d62728"), ("none", "#2ca02c")):
        h = arm(mode, 1.0)
        h = h[h["generation"] >= last]["heterozygosity"].mean()
        ax.axhline(h, ls="--", color=c, label=f"{mode}")
    ax.set(xlabel=r"novelty exponent $\alpha$", ylabel="stationary $H$",
           title="Quality-diversity keeps $H$\nabove greedy for all α")
    ax.legend(frameon=False, fontsize=9)

    # Panel 3: stationary support size per arm
    ax = axes[2]
    arms = [("greedy", 1.0, "greedy"), ("qd", 0.5, "quality-diversity α=0.5"),
            ("qd", 1.0, "quality-diversity α=1"), ("qd", 2.0, "quality-diversity α=2"), ("none", 1.0, "none")]
    labels, vals, errs, colors = [], [], [], []
    palette = {"greedy": "#d62728", "qd": "#ff7f0e", "none": "#2ca02c"}
    for mode, a, lab in arms:
        s = arm(mode, a)
        s = s[s["generation"] >= last]["support_size"]
        labels.append(lab); vals.append(s.mean()); errs.append(1.96 * s.sem())
        colors.append(palette[mode])
    ax.bar(range(len(labels)), vals, yerr=errs, color=colors, capsize=3)
    ax.set(ylabel="stationary support size", title="Surviving items per arm",
           xticks=range(len(labels)))
    ax.set_xticklabels(labels, rotation=25, ha="right", fontsize=8)

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


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