"""Stage B — the HARD GATE: the neural runner reproduces the Layer-1 analytic core.

With ``model.kind == "histogram"`` the neural generational step (train-on-parent's-samples
+ grounding) is *exactly* neutral Wright-Fisher drift with immigration. This module asserts
that the neural runner reproduces the two closed forms Layer 1 is validated against
(blueprint 2.4-1 neutral heterozygosity decay, 2.4-3 exact mutation-drift equilibrium) and
that its mean H-trajectory tracks ``inheritance.lineage.run_lineage`` directly. If any of
these fail the neural plumbing is wrong — no real network should be trained until they pass.
"""

from __future__ import annotations

import numpy as np
import pytest

from inheritance.lineage import run_lineage
from inheritance.seeding import spawn_seeds
from neural.generation_loop import run_generative_lineage
from neural.synthetic import make_mode_truth
from neural.config import SyntheticCfg


def theory_decay(H0: float, n: int, t: np.ndarray) -> np.ndarray:
    return H0 * (1.0 - 1.0 / n) ** np.asarray(t, dtype=float)


def theory_H_eq(n: int, m: int, H_star: float) -> float:
    return H_star * m * (2 * n + m - 1) / (n + 2 * n * m + m * m)


def _mean_H(cfg: dict, n_rep: int, master: int = 20260704) -> np.ndarray:
    """Mean heterozygosity trajectory over ``n_rep`` histogram-model replicates."""
    seeds = spawn_seeds(master, n_rep)
    Hs = [run_generative_lineage(cfg, int(s.generate_state(1)[0]))["heterozygosity"].to_numpy()
          for s in seeds]
    return np.mean(np.stack(Hs), axis=0)


# --- Pred. 1: neutral heterozygosity decay ----------------------------------------------

def test_bridge_neutral_decay_matches_theory():
    K, n, gens, reps = 50, 100, 20, 800
    cfg = {
        "synthetic": {"K": K, "R": 1, "zipf_s": 1.1, "init": "uniform",
                      "style_len": 2, "style_vocab": 4, "id_base": 2},
        "model": {"kind": "histogram"},
        "dynamics": {"n": n, "grounding": {"m": 0}},
        "generations": gens,
    }
    H_sim = _mean_H(cfg, reps)
    t = np.arange(gens + 1)
    H_theory = theory_decay(1.0 - 1.0 / K, n, t)
    rel_err = np.abs(H_sim - H_theory) / H_theory
    assert rel_err.max() < 0.03, f"max rel err {rel_err.max():.4f} exceeds 0.03"


# --- Pred. 3: exact mutation-drift equilibrium under grounding ---------------------------

def test_bridge_grounded_equilibrium_matches_theory():
    K, n, m, gens, reps = 80, 100, 8, 220, 300
    cfg = {
        "synthetic": {"K": K, "R": 1, "zipf_s": 1.1, "init": "uniform",
                      "style_len": 2, "style_vocab": 4, "id_base": 2},
        "model": {"kind": "histogram"},
        "dynamics": {"n": n, "grounding": {"m": m}},   # R=1 -> proportional immigration
        "generations": gens,
    }
    H_sim_traj = _mean_H(cfg, reps)
    H_sim = float(H_sim_traj[-60:].mean())             # stationary average
    td = make_mode_truth(SyntheticCfg(K=K, R=1, zipf_s=1.1))
    H_star = 1.0 - float(np.sum(td.p_star ** 2))
    H_eq = theory_H_eq(n, m, H_star)
    assert H_sim == pytest.approx(H_eq, rel=0.05), f"sim {H_sim:.4f} vs theory {H_eq:.4f}"


# --- Direct bridge: histogram lineage tracks Layer-1 run_lineage -------------------------

def test_bridge_tracks_layer1_trajectory():
    K, n, m, gens, reps = 60, 120, 6, 40, 400
    neural_cfg = {
        "synthetic": {"K": K, "R": 1, "zipf_s": 1.1, "init": "uniform",
                      "style_len": 2, "style_vocab": 4, "id_base": 2},
        "model": {"kind": "histogram"},
        "dynamics": {"n": n, "grounding": {"m": m}},
        "generations": gens,
    }
    layer1_cfg = {
        "truth": {"K": K, "R": 1, "zipf_s": 1.1, "init": "uniform"},
        "dynamics": {"n": n, "grounding": {"m": m, "policy": "proportional"}},
        "generations": gens,
    }
    seeds = spawn_seeds(20260704, reps)
    H_neural = np.mean(np.stack([
        run_generative_lineage(neural_cfg, int(s.generate_state(1)[0]))["heterozygosity"].to_numpy()
        for s in seeds]), axis=0)
    H_layer1 = np.mean(np.stack([
        run_lineage(layer1_cfg, int(s.generate_state(1)[0]))["heterozygosity"].to_numpy()
        for s in seeds]), axis=0)
    rel_err = np.abs(H_neural - H_layer1) / H_layer1
    assert rel_err.max() < 0.03, f"neural vs Layer-1 max rel err {rel_err.max():.4f}"