"""Learning-kernel tests (pure NumPy).

The kernel is an additive Layer-1 extension: identity by default (so the neutral Wright-Fisher
core and every scientific-validation test are unchanged), a sharpening knob that ADDS collapse
where neutral drift is inert, and a smoothing knob that supplies a diversity FLOOR where neutral
drift would collapse to zero. These assert exactly those three behaviours.
"""

from __future__ import annotations

import numpy as np

from inheritance.kernel import LearningKernelCfg, apply_kernel
from inheritance.lineage import run_lineage
from inheritance.metrics import heterozygosity


def test_apply_kernel_identity_is_noop():
    p = np.array([0.5, 0.3, 0.15, 0.05])
    assert np.array_equal(apply_kernel(p, LearningKernelCfg()), p)


def test_apply_kernel_reset_mixes_toward_uniform():
    p = np.array([1.0, 0.0, 0.0, 0.0])
    out = apply_kernel(p, LearningKernelCfg(reset=0.2))
    assert np.isclose(out.sum(), 1.0)
    assert np.allclose(out, [0.8 + 0.2 / 4, 0.05, 0.05, 0.05])   # keeps dead modes alive
    assert (out > 0).all()


def test_apply_kernel_sharpen_and_floor_prune():
    p = np.array([0.6, 0.3, 0.09, 0.01])
    sharp = apply_kernel(p, LearningKernelCfg(temperature=0.5))    # p^2, renormalised
    assert sharp[0] > p[0] and sharp[-1] < p[-1]                   # mass concentrates
    floored = apply_kernel(p, LearningKernelCfg(floor=0.05))
    assert floored[-1] == 0.0 and np.isclose(floored.sum(), 1.0)   # weak mode dropped


def test_kernel_identity_leaves_lineage_unchanged():
    base = {"truth": {"K": 64, "init": "truth"}, "dynamics": {"n": 200},
            "generations": 20, "metrics": {"kl_floor": 1e-9, "support_eps": 1e-9}}
    withk = {**base, "dynamics": {"n": 200, "kernel": {"reset": 0.0, "temperature": 1.0}}}
    a = run_lineage(base, seed=3)["heterozygosity"].to_numpy()
    b = run_lineage(withk, seed=3)["heterozygosity"].to_numpy()
    assert np.array_equal(a, b)          # identity kernel == neutral Wright-Fisher, bitwise


def test_sharpening_adds_collapse_where_neutral_is_inert():
    # Large n vs small K: neutral drift barely collapses; sharpening drives it to ~1 mode.
    cfg = {"truth": {"K": 30, "tail": "zipf", "zipf_s": 1.5, "tail_threshold": 1e-2,
                     "init": "truth"},
           "dynamics": {"n": 6000, "grounding": {"m": 0}},
           "generations": 15, "metrics": {"kl_floor": 1e-9, "support_eps": 1e-9}}
    neutral = run_lineage(cfg, seed=0)
    sharp = run_lineage({**cfg, "dynamics": {**cfg["dynamics"],
                                             "kernel": {"temperature": 0.8}}}, seed=0)
    assert neutral["support_size"].iloc[-1] > 20          # neutral: ~all modes survive
    assert sharp["support_size"].iloc[-1] <= 3            # sharpening: collapse to a point
    assert sharp["heterozygosity"].iloc[-1] < 0.1


def test_smoothing_floors_diversity_where_neutral_collapses():
    # Small n vs large K: neutral drift drives H toward 0; smoothing holds a positive floor.
    cfg = {"truth": {"K": 256, "tail": "zipf", "zipf_s": 1.3, "tail_threshold": 1e-3,
                     "init": "truth"},
           "dynamics": {"n": 200, "grounding": {"m": 0}},
           "generations": 120, "metrics": {"kl_floor": 1e-9, "support_eps": 1e-9}}
    neutral = run_lineage(cfg, seed=0)
    smooth = run_lineage({**cfg, "dynamics": {**cfg["dynamics"],
                                              "kernel": {"reset": 0.006}}}, seed=0)
    assert neutral["heterozygosity"].iloc[-1] < 0.4      # neutral collapses
    assert smooth["heterozygosity"].iloc[-1] > 0.55      # smoothing floors H well above neutral