"""E12 figure — model speciation: the merge-compatibility limit of the sexual society.

Three panels, reading only the committed bundles. (A) BDM: mean recombinant (hybrid) fitness vs
parental divergence, one line per epistasis density rho, against the rising parent fitness — the
compatible -> outbreeding-depression -> hybrid-inviability trajectory, peaking then crashing sooner the
denser the epistasis. (B) BDM: the reproductive-isolation rate (fraction of hybrids below the ancestor)
vs divergence — the isolation cliff, moving to lower divergence as epistasis density rises. (C) NK: the
epistasis wedge — as landscape ruggedness K grows, recombining two adapted local-optimum parents flips
from a gain to outbreeding depression.

Usage:  python figures/plot_E12.py
"""

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 _agg(df, keys, value):
    g = df.groupby(keys)[value].agg(["mean", "std", "count"]).reset_index()
    g["se"] = g["std"] / np.sqrt(g["count"].clip(lower=1))
    return g


def main() -> None:
    bdm, _ = load_bundle("results/E12")
    nk, _ = load_bundle("results/E12_nk")
    rhos = sorted(bdm["rho"].unique())
    colors = plt.cm.viridis(np.linspace(0.15, 0.85, len(rhos)))

    fig, axes = plt.subplots(1, 3, figsize=(16, 5))

    # Panel A: hybrid fitness vs divergence, per epistasis density, + parent fitness.
    ax = axes[0]
    par = _agg(bdm, "divergence", "parent_fitness")
    ax.plot(par["divergence"], par["mean"], "k--", lw=1.6, label="parent fitness")
    for rho, c in zip(rhos, colors):
        g = _agg(bdm[bdm["rho"] == rho], "divergence", "offspring_fitness")
        ax.plot(g["divergence"], g["mean"], "-o", color=c, lw=2, label=f"hybrid, ρ={rho}")
        ax.fill_between(g["divergence"], g["mean"] - g["se"], g["mean"] + g["se"], color=c, alpha=0.15)
    ax.axhline(0, color="#999", lw=0.8, ls=":")
    ax.set(xlabel="parental divergence  (substitutions $d$)", ylabel="fitness",
           title="Hybrid fitness collapses as lineages diverge\n(compatible → outbreeding depression → inviability)")
    ax.legend(frameon=False, fontsize=8)

    # Panel B: reproductive-isolation rate vs divergence, per epistasis density.
    ax = axes[1]
    for rho, c in zip(rhos, colors):
        g = _agg(bdm[bdm["rho"] == rho], "divergence", "isolation")
        ax.plot(g["divergence"], g["mean"], "-o", color=c, lw=2, label=f"ρ={rho}")
    ax.set(xlabel="parental divergence  (substitutions $d$)", ylabel="reproductive isolation\n(P hybrid inviable)",
           ylim=(-0.02, 1.02),
           title="The isolation cliff moves to lower divergence\nas epistasis density rises")
    ax.legend(frameon=False, fontsize=9, title="epistasis density")

    # Panel C: NK epistasis wedge — recombination gain vs ruggedness K.
    ax = axes[2]
    g = _agg(nk, "K", "offspring_minus_parent")
    ax.axhline(0, color="#999", lw=0.8, ls=":")
    ax.plot(g["K"], g["mean"], "-o", color="#d62728", lw=2)
    ax.fill_between(g["K"], g["mean"] - g["se"], g["mean"] + g["se"], color="#d62728", alpha=0.15)
    ax.set(xlabel="landscape ruggedness $K$ (epistasis)", ylabel="recombination gain\n(hybrid − worse parent)",
           title="Epistasis wedge: recombining adapted parents\nflips from gain to loss as ruggedness grows")

    fig.suptitle("E12 — model speciation: when two diverged models are too incompatible to merge",
                 y=1.02, fontsize=13)
    fig.tight_layout()
    savefig(fig, "results/E12", "E12")


if __name__ == "__main__":
    main()