"""Mating systems — monogamy vs promiscuity as mate-pool breadth (E14).

The society experiments (E8–E11) assumed **panmixia**: every offspring is recombined from parents
sampled across the *whole* population. But biology's mating systems span a continuum from **monogamy**
(each individual mates within a narrow, local circle) to **promiscuity** (mates drawn freely from the
whole population), and population genetics says the choice is consequential. Wide gene flow spreads a
beneficial allele across the population fast but **homogenises** it; restricted gene flow (population
structure / *isolation by distance*) keeps demes distinct so several fitness peaks can be explored in
parallel — Wright's *shifting balance*.

Here the mating system is one scalar: mate-pool **breadth** ``b``. Agents sit on a ring; an offspring's
second parent is drawn from a window of half-width ``≈ b·N/2`` around the focal parent. ``b→0`` =
**monogamous / structured** (local mating, isolation by distance); ``b=1`` = **promiscuous / panmictic**
(mate with anyone). Selection is **local** — an offspring competes only against the incumbent at its own
ring position — so restricted mating can actually sustain distinct demes rather than being washed out by
global truncation.

Crossed with landscape ruggedness ``K`` (Kauffman NK epistasis), this is the mating-system image of the
E9 design rule. Prediction: **promiscuity wins on additive/smooth landscapes** (one peak — spread the
single good direction fastest), while **structured/monogamous mating wins on rugged/epistatic
landscapes** (many peaks — diversity must be preserved to explore basins that recombination can later
combine). Falsifier: the best mating system is independent of ruggedness (no crossover).
"""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np
import pandas as pd

from .genotype import bits_to_index, crossover, hill_climb, nk_fitness


def _diversity(pop_bits: np.ndarray) -> float:
    """Mean normalised pairwise Hamming distance over the population (0 = clonal, 1 = maximal)."""
    N, L = pop_bits.shape
    if N < 2:
        return 0.0
    match = (pop_bits[:, None, :] == pop_bits[None, :, :]).sum(axis=2)   # (N, N) locus agreements
    ham = L - match                                                      # pairwise Hamming distances
    return float(ham.sum() / (N * (N - 1)) / L)                          # mean over ordered pairs, /L


def _distinct_peaks(pop_bits: np.ndarray, fitness: np.ndarray, L: int) -> int:
    """Number of distinct local optima the population occupies (hill-climb each agent to its basin)."""
    return len({hill_climb(fitness, L, bits_to_index(b)) for b in pop_bits})


def run_mating_system(cfg: Mapping[str, Any], seed: int) -> pd.DataFrame:
    """Run one mating-system lineage; return per-generation metrics.

    Args:
        cfg (Mapping): Config with a ``mating`` block (``L`` loci, ``K`` landscape ruggedness, ``N``
            population, ``breadth`` mate-pool breadth ``b∈[0,1]``, ``recomb_rate`` crossover rate,
            ``mu`` per-locus mutation) and ``generations``.
        seed (int): Replicate seed; the landscape and the run are a pure function of it.

    Returns:
        pd.DataFrame: One row per generation with ``best_fitness`` (real), ``mean_fitness`` (real),
        ``diversity`` (mean normalised pairwise Hamming), ``distinct_peaks`` (local optima occupied),
        and ``global_opt``.
    """
    ms = cfg["mating"]
    L, K, N = int(ms["L"]), int(ms["K"]), int(ms["N"])
    b = float(ms.get("breadth", 1.0))
    rate = float(ms.get("recomb_rate", 0.5))
    mu = float(ms.get("mu", 0.01))
    generations = int(cfg.get("generations", 100))

    fitness = nk_fitness(L, K, seed)                     # reality
    global_opt = float(fitness.max())
    rng = np.random.default_rng(seed)

    # Population on a ring: position i is fixed ring slot i (so structure persists across generations).
    pop = rng.integers(0, 2, size=(N, L)).astype(np.int8)
    half = max(1, int(round(b * N / 2)))                 # mate-window half-width; b=1 -> whole ring

    def fit_of(bits: np.ndarray) -> float:
        return float(fitness[bits_to_index(bits)])

    rows: list[dict] = []

    def record(t: int) -> None:
        tf = np.array([fit_of(g) for g in pop])
        rows.append({
            "generation": t,
            "best_fitness": float(tf.max()),
            "mean_fitness": float(tf.mean()),
            "diversity": _diversity(pop),
            "distinct_peaks": _distinct_peaks(pop, fitness, L),
            "global_opt": global_opt,
        })

    record(0)
    for t in range(1, generations + 1):
        new = pop.copy()
        for i in range(N):
            # Second parent from a ring window of half-width `half` around i (isolation by distance).
            offset = 0
            while offset == 0:
                offset = int(rng.integers(-half, half + 1))
            j = (i + offset) % N
            child = crossover(np.stack([pop[i], pop[j]]), rate, rng)
            flip = rng.random(L) < mu
            child = np.where(flip, 1 - child, child).astype(pop.dtype)
            # Local selection: the child replaces the incumbent at i only if strictly fitter.
            if fit_of(child) > fit_of(pop[i]):
                new[i] = child
        pop = new
        record(t)

    return pd.DataFrame(rows)