"""E14 figure — mating systems: monogamy vs promiscuity (mate-pool breadth) across ruggedness.

Three panels, each vs mate-pool breadth (log x: 0.03 = monogamous/structured -> 1.0 = promiscuous/
panmictic), one line per landscape ruggedness K:

(A) best fitness / global optimum — the *champion*. On smooth landscapes (low K) it is maximised by
    wide breadth; as ruggedness rises the peak shifts to an INTERMEDIATE breadth (full promiscuity
    prematurely converges below it) — the mating-system image of E9's "optimal recombination rate
    shrinks with ruggedness".
(B) mean fitness / global optimum — the *typical* individual. Monotonically favoured by breadth at
    every K: panmixia lifts the whole population toward a good consensus.
(C) diversity (mean normalised pairwise Hamming) — monotonically DESTROYED by breadth at every K
    (promiscuity homogenises), the reservoir largest under monogamy and on rugged landscapes.

The tension between (A)/(C) is the result: promiscuity maximises the typical model and kills diversity;
on rugged landscapes the best model needs preserved diversity, so an intermediate breadth wins.

Usage:  python figures/plot_figS13_mating_breadth.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).parent))
from _figlib import load_bundle, savefig, letter_axes  # noqa: E402


def main() -> None:
    df, _ = load_bundle("results/figS13_mating_breadth")
    last = df[df["generation"] == df["generation"].max()].copy()
    last["best_n"] = last["best_fitness"] / last["global_opt"]
    last["mean_n"] = last["mean_fitness"] / last["global_opt"]
    Ks = sorted(last["K"].unique())
    cmap = plt.get_cmap("viridis")
    colors = {K: cmap(i / max(1, len(Ks) - 1)) for i, K in enumerate(Ks)}

    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    panels = [
        ("best_n", "best fitness / global optimum",
         "Best model peaks at intermediate breadth on rugged\nlandscapes (the peak shifts left as $K$ rises)"),
        ("mean_n", "mean fitness / global optimum",
         "Population mean rises monotonically with breadth\n(panmixia lifts the whole population)"),
        ("diversity", "diversity (mean pairwise Hamming)",
         "Standing diversity falls monotonically with breadth\n(promiscuity homogenises; monogamy preserves)"),
    ]
    for ax, (col, ylab, title) in zip(axes, panels):
        for K in Ks:
            g = (last[last["K"] == K].groupby("breadth")[col]
                 .agg(["mean", "sem"]).reset_index())
            ax.errorbar(g["breadth"], g["mean"], yerr=1.96 * g["sem"].fillna(0.0),
                        marker="o", lw=1.8, capsize=2, color=colors[K], label=f"K={K}")
        ax.set_xscale("log")
        ax.set(xlabel="mate-pool breadth  (monogamous ← → promiscuous)", ylabel=ylab)
        ax.set_title(title, fontsize=9)
        ax.legend(title="ruggedness", frameon=False, fontsize=8)

    fig.tight_layout()
    letter_axes(fig)
    savefig(fig, "results/figS13_mating_breadth", "figS13_mating_breadth")


if __name__ == "__main__":
    main()