"""The dynamic Lamarckian society — the vertical claim (E11 / C3).

A finite population of ``N`` agents (genotypes of ``L`` biallelic loci) evolves on a Kauffman NK
landscape that *is* reality. The society climbs in real capability by composing the four operators the
whole study built toward — **grounding**, **directed recombination (sex)**, **quality-diversity
selection**, and mutation — and an ablation shows each is load-bearing.

The crux is what happens WITHOUT grounding. A plain genetic algorithm on true fitness would just
improve, so grounding must corrupt the *selection signal* to cause collapse. Here selection acts on a
**grounded score** ``g·true_fitness + (1−g)·conformity``, where conformity is agreement with the
population's own consensus (modal genotype). At ``g=0`` selection rewards fitting the crowd rather
than reality — self-consumption — and the society drifts to a fit-looking but actually-poor consensus,
losing diversity: the direct analogue of training on the majority of AI-generated outputs.

Four ablation arms, each breaking distinctly (only ``full`` avoids all three failures):
``full`` (climbs) · ``no_grounding`` (conformity collapse) · ``no_sex`` (stuck at local optima) ·
``no_diversity`` (collapses to one lineage, recombination starves).
"""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np
import pandas as pd

from .genotype import bits_to_index, crossover, genotype_bits, nk_fitness


def _consensus(pop_bits: np.ndarray) -> np.ndarray:
    """Population consensus genotype: the modal allele at each locus (majority vote)."""
    return (pop_bits.mean(axis=0) >= 0.5).astype(pop_bits.dtype)


def _conformity(pop_bits: np.ndarray, consensus: np.ndarray) -> np.ndarray:
    """Per-agent agreement with the consensus (fraction of loci matching the majority)."""
    return (pop_bits == consensus[None, :]).mean(axis=1)


def _novelty(pop_bits: np.ndarray) -> np.ndarray:
    """Per-agent novelty: mean Hamming distance to the rest of the population (diversity signal)."""
    N, L = pop_bits.shape
    if N < 2:
        return np.zeros(N)
    # pairwise Hamming via allele agreement: distance_ij = L - matches; mean over j != i.
    match = (pop_bits[:, None, :] == pop_bits[None, :, :]).sum(axis=2)   # (N, N) matches
    ham = L - match
    return (ham.sum(axis=1) / (N - 1)) / L                               # normalised to [0,1]


def _directed_offspring(pop_bits, fitness, n_off, rate, rng):
    """Directed sex: make ``n_off`` recombinants from the whole population, return them ranked-ready.

    Unbounded-parent crossover (the AI move); offspring selection happens in the survival step, so
    here we just generate the candidate offspring bit-matrix.
    """
    return np.stack([crossover(pop_bits, rate, rng) for _ in range(n_off)])


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

    Args:
        cfg (Mapping): Config with a ``society`` block (``L``, ``K`` landscape ruggedness, ``N``
            population, ``g`` grounding, ``mu`` mutation, ``novelty`` QD weight, ``n_off`` offspring
            pool, ``recomb_rate``, ``sex`` on/off, ``select`` in {``qd``, ``greedy``}) 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), ``consensus_fitness``,
        ``conformity_true_gap`` (mean conformity − mean true fitness; exposes the no-grounding
        collapse), and ``global_opt``.
    """
    soc = cfg["society"]
    L, K, N = int(soc["L"]), int(soc["K"]), int(soc["N"])
    g = float(soc.get("g", 1.0))
    mu = float(soc.get("mu", 0.02))
    novelty_w = float(soc.get("novelty", 0.0))
    n_off = int(soc.get("n_off", N))
    rate = float(soc.get("recomb_rate", 0.2))
    sex = bool(soc.get("sex", True))
    select = soc.get("select", "qd")
    generations = int(cfg.get("generations", 100))

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

    # Initialise a diverse population of random genotypes.
    pop = rng.integers(0, 2, size=(N, L)).astype(all_bits.dtype)

    def true_fit(bits):
        return np.array([fitness[bits_to_index(b)] for b in bits])

    rows: list[dict] = []

    def record(t: int) -> None:
        tf = true_fit(pop)
        cons = _consensus(pop)
        conf = _conformity(pop, cons)
        rows.append({
            "generation": t,
            "best_fitness": float(tf.max()),
            "mean_fitness": float(tf.mean()),
            "diversity": float(_novelty(pop).mean()),
            "consensus_fitness": float(fitness[bits_to_index(cons)]),
            "conformity_true_gap": float(conf.mean() - tf.mean()),
            "global_opt": global_opt,
        })

    record(0)
    for t in range(1, generations + 1):
        # (1) candidate pool = current population + directed offspring (sex) or mutated clones.
        if sex:
            offspring = _directed_offspring(pop, fitness, n_off, rate, rng)
        else:                                            # asexual: offspring are mutated copies
            idx = rng.integers(0, N, size=n_off)
            offspring = pop[idx].copy()
        # mutation on the offspring
        flip = rng.random(offspring.shape) < mu
        offspring = np.where(flip, 1 - offspring, offspring).astype(pop.dtype)
        pool = np.concatenate([pop, offspring], axis=0)

        # (2) grounded score: g*true_fitness + (1-g)*conformity (conformity vs the *current* consensus).
        cons = _consensus(pop)
        tf = true_fit(pool)
        conf = _conformity(pool, cons)
        score = g * tf + (1.0 - g) * conf

        # (3) survival: QD (score + novelty) keeps diverse high-scorers; greedy keeps top score only.
        if select == "qd" and novelty_w > 0.0:
            nov = _novelty(pool)
            merit = score + novelty_w * nov
        else:
            merit = score
        keep = np.argsort(merit)[-N:]                    # elitist truncation survival
        pop = pool[keep]
        record(t)

    return pd.DataFrame(rows)