"""Tests for E12 model speciation — the BDM construction, the snowball, and the isolation falsifiers."""

from __future__ import annotations

import numpy as np
import pytest

from inheritance.speciation import _bdm_point, _nk_point, run_speciation


def test_bdm_parents_carry_no_incompatibility():
    # The BDM construction's defining property: each derived allele is benign on its OWN parent's
    # background (disjoint substitutions), so a parent's fitness is purely additive even at rho=1.
    rng = np.random.default_rng(0)
    r = _bdm_point(L=20, d=10, rho=1.0, s=5.0, beta=1.0, rate=0.5, n_off=200, rng=rng)
    assert r["parent_fitness"] == pytest.approx(1.0 * (10 // 2))    # beta * (d/2), no penalty


def test_bdm_snowball_is_superlinear_in_divergence():
    # Orr-Turelli: # incompatibilities ~ (d/2)^2, so doubling divergence ~quadruples them.
    def mean_ndmi(d, reps=40):
        return np.mean([_bdm_point(24, d, 0.5, 1.0, 1.0, 0.5, 50,
                                   np.random.default_rng(i))["n_dmi"] for i in range(reps)])
    ratio = mean_ndmi(12) / max(mean_ndmi(6), 1e-9)
    assert ratio > 3.0                                             # ~4x (quadratic), well above linear (2x)


def test_bdm_no_epistasis_means_no_isolation():
    rng = np.random.default_rng(1)
    r = _bdm_point(L=20, d=20, rho=0.0, s=1.0, beta=1.0, rate=0.5, n_off=300, rng=rng)
    assert r["n_dmi"] == 0 and r["isolation"] == 0.0               # no BDMIs -> hybrids always viable


def test_bdm_isolation_rises_with_epistasis_density():
    # At fixed high divergence, denser epistasis (rho) -> more reproductive isolation.
    def iso(rho):
        return np.mean([_bdm_point(20, 20, rho, 1.0, 1.0, 0.5, 300,
                                   np.random.default_rng(i))["isolation"] for i in range(8)])
    assert iso(0.5) > iso(0.1)


def test_nk_additive_landscape_has_no_isolation():
    # K=0 is a single-peak additive landscape: parents hill-climb to the same optimum (divergence 0),
    # and recombination cannot produce outbreeding depression.
    r = _nk_point(L=12, K=0, landscape_seed=3, rate=0.5, n_pairs=20, n_off=50,
                  rng=np.random.default_rng(0))
    assert r["divergence"] == pytest.approx(0.0) and r["outbreeding_depression"] == pytest.approx(0.0)


def test_nk_ruggedness_increases_outbreeding_depression():
    def od(K):
        return _nk_point(12, K, 3, 0.5, 30, 60, np.random.default_rng(0))["outbreeding_depression"]
    assert od(8) > od(0)                                           # rugged landscapes punish recombination


def test_run_speciation_is_deterministic():
    cfg = {"seed": 7, "n_replicates": 3,
           "speciation": {"landscape": "bdm", "L": 12, "rho": [0.3], "divergences": [0, 4, 8],
                          "s": 1.0, "beta": 1.0, "recomb_rate": 0.5, "n_offspring": 80}}
    a = run_speciation(cfg, 7)
    b = run_speciation(cfg, 7)
    assert a.equals(b)                                             # pure function of the resolved config + seed