"""E7 figure — the advantage of sex: recombination adapts faster than clonal reproduction.

The dynamic mechanism behind E8. A single population adapts from all-wrong toward a multi-locus
optimum under selection + drift + mutation. Beneficial alleles arise in different sub-lineages;
sexual recombination reassorts them into one genotype, while an asexual lineage suffers clonal
interference. The sexual lineage climbs faster — the classical advantage of sex (an honest *speed*
advantage; both eventually plateau near the optimum in this tractable regime).

Two panels: (A) mean-fitness adaptation curves, asexual vs sexual, over generations; (B) linkage
disequilibrium over generations — asexual holds beneficial alleles in disequilibrium (scattered
across genotypes) while sexual drives it to ~0 (assembled), the mechanism of the speed gap.

Usage:  python figures/plot_E7.py [results/E7]
"""

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


def main(results_dir: str = "results/E7") -> None:
    df, cfg = load_bundle(results_dir)
    L = cfg["genotype"]["L"]
    arms = [(0.0, "#7f7f7f", "asexual (clonal)"), (1.0, "#d62728", "sexual (recombining)")]

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

    ax = axes[0]
    for rate, c, lab in arms:
        sub = df[df["recomb_rate"] == rate]
        g, m, ci = mean_ci(sub, "generation", "mean_fitness")
        ax.plot(g, m, "-", color=c, lw=1.8, label=lab)
        ax.fill_between(g, m - ci, m + ci, color=c, alpha=0.2)
    ax.axhline(L, ls=":", color="green", lw=1, label=f"optimum ($L$={L})")
    ax.set(xlabel="generation", ylabel="mean fitness (# correct loci)",
           title="Advantage of sex: recombination adapts faster\n(clonal interference slows the asexual lineage)")
    ax.legend(frameon=False, fontsize=9)

    ax = axes[1]
    for rate, c, lab in arms:
        sub = df[df["recomb_rate"] == rate]
        g, m, ci = mean_ci(sub, "generation", "ld")
        ax.plot(g, m, "-", color=c, lw=1.8, label=lab)
        ax.fill_between(g, m - ci, m + ci, color=c, alpha=0.2)
    ax.set(xlabel="generation", ylabel="mean linkage disequilibrium |D|",
           title="Mechanism: asexual scatters beneficial alleles (LD>0);\nsexual assembles them (LD→0)")
    ax.legend(frameon=False, fontsize=9)

    fig.suptitle("E7 — the advantage of sex: recombination reassorts beneficial alleles that arose "
                 "in different lineages", y=1.02, fontsize=12)
    fig.tight_layout()
    savefig(fig, results_dir, "E7")


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