"""E9 figure — landscape robustness: when recombination helps, and the outbreeding-depression risk.

The credibility test for the sexual metaphor. E8 used an additive landscape where recombination
trivially helps; here parents are local optima ("trained models") of a Kauffman NK landscape whose
ruggedness (epistasis) is tunable. Blindly recombining entangled models breaks co-adapted allele
blocks and offspring fall *below* the parents — outbreeding depression — worse the more rugged the
landscape and the higher the recombination rate. With selection (best offspring), a nonzero optimal
recombination rate re-emerges. Design rule: merge freely when skills are complementary; merge
sparingly (and always select) when they are entangled.

Two panels: (A) the risk — mean offspring fitness minus best-parent vs recombination rate, one curve
per ruggedness K (all ≤0, steeper as K grows); (B) with offspring selection — best-of-brood fitness
vs rate per K, showing an intermediate optimum on rugged landscapes. Reads only the bundle.

Usage:  python figures/plot_E9.py [results/E9]
"""

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/E9") -> None:
    df, _ = load_bundle(results_dir)
    Ks = sorted(df["K"].unique())
    rates = sorted(df["rate"].unique())
    colors = plt.cm.viridis(np.linspace(0, 0.85, len(Ks)))
    bp = df.groupby("K")["best_parent"].mean()

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

    # Panel A: the risk — mean offspring minus best parent vs rate, per K.
    ax = axes[0]
    for K, c in zip(Ks, colors):
        s = df[df["K"] == K].groupby("rate")["mean_offspring"].mean() - bp[K]
        ax.plot(s.index, s.values, "-o", color=c, ms=4, label=f"K={K}")
    ax.axhline(0, ls=":", color="gray", lw=1)
    ax.set(xlabel="recombination rate", ylabel="mean offspring − best parent",
           title="The risk: outbreeding depression\n(worse with ruggedness K and recombination rate)")
    ax.legend(frameon=False, fontsize=8, title="ruggedness")

    # Panel B: with selection — best offspring vs rate, per K (intermediate optimum on rugged).
    ax = axes[1]
    for K, c in zip(Ks, colors):
        s = df[df["K"] == K].groupby("rate")["best_offspring"].mean()
        ax.plot(s.index, s.values, "-o", color=c, ms=4, label=f"K={K}")
        ax.axhline(bp[K], ls=":", color=c, lw=0.8, alpha=0.6)
    ax.set(xlabel="recombination rate", ylabel="best-of-brood fitness (with selection)",
           title="With offspring selection, an optimal\nrecombination rate re-emerges (dotted = parents)")
    ax.legend(frameon=False, fontsize=8, title="ruggedness")

    fig.suptitle("E9 — landscape robustness: recombination helps when skills are complementary, but "
                 "blindly merging entangled models causes outbreeding depression", y=1.02, fontsize=11)
    fig.tight_layout()
    savefig(fig, results_dir, "E9")


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