"""Stage C torch-model tests (skipped when torch is absent).

Small, fast sign checks — the neural tiers are statistically reproducible and directional,
not exact, so these assert the *sign* of each effect (blueprint 3.5): gen-0 fidelity, dry
collapse, and grounding arresting it. They gate the RNN before the N-series experiments.
"""

from __future__ import annotations

import numpy as np
import pytest

pytest.importorskip("torch")

from knowledge.metrics import forward_kl, heterozygosity  # noqa: E402
from neural.config import ModelCfg, SyntheticCfg  # noqa: E402
from neural.generation_loop import run_generative_lineage  # noqa: E402
from neural.models import make_model  # noqa: E402
from neural.oracle import ExactOracle  # noqa: E402
from neural.synthetic import make_mode_truth  # noqa: E402

_SYN = dict(K=256, R=1, zipf_s=1.3, init="truth", style_len=3, style_vocab=5, id_base=2,
            tail_threshold=1e-3)
# hidden/epochs high enough that the RNN sharpens (an underfit RNN smooths and resists
# collapse); with n=200 K=256 the dry lineage collapses robustly across seeds.
_MODEL = dict(kind="rnn", hidden=128, embed=24, epochs=25, lr=2e-3, batch_size=256, n_eval=10000)


def _lineage_cfg(g: float, n: int, gens: int) -> dict:
    m = 0 if g == 0 else round(n * g / (1 - g))
    return {
        "synthetic": dict(_SYN),
        "model": dict(_MODEL),
        "dynamics": {"n": n, "grounding": {"m": m, "policy": "proportional"}},
        "generations": gens,
    }


@pytest.mark.parametrize("kind", ["rnn", "mlp"])
def test_gen0_fidelity(kind):
    # A trained gen-0 model must recover p* (else "collapse" would be underfitting). Checked
    # for the RNN and MLP; the VAE does not clear this gate on the codeword task (see todo).
    syn = SyntheticCfg(**_SYN)
    td = make_mode_truth(syn)
    model = make_model(ModelCfg(**{**_MODEL, "kind": kind}), syn, ExactOracle(syn))
    model.initialise(td.p_star, np.random.default_rng(0))
    p_hat = model.mode_distribution(np.random.default_rng(1))
    assert forward_kl(td.p_star, p_hat, 1e-9) < 0.25     # close to truth
    assert (p_hat > 1e-9).sum() >= 0.9 * syn.K            # most modes represented


def test_rnn_dry_collapses_grounded_holds():
    # gens=20 gives clean dry-vs-grounded separation (KL ~2+ vs ~0.3); big margins survive
    # GPU non-determinism. Directional per blueprint 3.5.
    dry = run_generative_lineage(_lineage_cfg(0.0, 200, 25), seed=0)
    grd = run_generative_lineage(_lineage_cfg(0.05, 200, 25), seed=0)
    assert dry["heterozygosity"].iloc[-1] < dry["heterozygosity"].iloc[0] - 0.10
    assert dry["forward_kl"].iloc[-1] > 1.5                       # tail forgotten
    assert grd["forward_kl"].iloc[-1] < dry["forward_kl"].iloc[-1]   # grounding closer to truth
    assert grd["heterozygosity"].iloc[-1] > dry["heterozygosity"].iloc[-1]