"""Correctness tests for the Layer 1.5 neural scaffold (Stage A).

Pure-NumPy checks (no torch): the synthetic grammar is lossless, the exact oracle has zero
error, the histogram model reduces to a mode-frequency estimator, and the generation loop
produces the Layer-1 row schema deterministically.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from neural.config import NeuralLineageCfg, SyntheticCfg
from neural.generation_loop import run_generative_lineage
from neural.models import HistogramModel, make_model
from neural.oracle import ExactOracle, measure_distribution
from neural.synthetic import id_codewords, make_mode_truth, render_modes, sample_synthetic


def _syn(**over) -> SyntheticCfg:
    base = dict(K=64, R=1, zipf_s=1.1, tail_threshold=1e-3, style_len=3, style_vocab=5,
                id_base=2)
    base.update(over)
    return SyntheticCfg(**base)


# --- synthetic grammar ------------------------------------------------------------------

def test_id_len_covers_all_modes():
    syn = _syn(K=100, id_base=2)
    assert syn.id_base ** syn.id_len >= syn.K
    assert syn.id_base ** (syn.id_len - 1) < syn.K


def test_seq_len_and_vocab():
    syn = _syn(K=64, id_base=2, style_len=3, style_vocab=5)
    assert syn.id_len == 6                      # 2**6 = 64
    assert syn.seq_len == syn.id_len + syn.style_len
    assert syn.vocab == max(syn.id_base, syn.style_vocab)


def test_codewords_are_unique_and_invertible():
    syn = _syn(K=64)
    cw = id_codewords(syn)
    assert cw.shape == (syn.K, syn.id_len)
    assert cw.max() < syn.id_base
    # each mode's codeword is distinct
    assert len({tuple(r) for r in cw}) == syn.K


def test_render_shapes_and_token_ranges():
    syn = _syn(K=32, style_len=4, style_vocab=7)
    rng = np.random.default_rng(0)
    modes = np.arange(syn.K)
    X = render_modes(modes, syn, rng)
    assert X.shape == (syn.K, syn.seq_len)
    assert X[:, : syn.id_len].max() < syn.id_base
    assert X[:, syn.id_len :].max() < syn.style_vocab


# --- exact oracle -----------------------------------------------------------------------

def test_exact_oracle_zero_error_on_all_modes():
    syn = _syn(K=100)
    rng = np.random.default_rng(1)
    modes = np.repeat(np.arange(syn.K), 5)      # every mode, many style draws
    X = render_modes(modes, syn, rng)
    recovered = ExactOracle(syn).classify(X)
    assert np.array_equal(recovered, modes)     # zero measurement error


def test_measure_distribution_recovers_frequencies():
    syn = _syn(K=16)
    rng = np.random.default_rng(2)
    p = np.array([0.5] + [0.5 / 15] * 15)
    X, _ = sample_synthetic(p, 200_000, syn, rng)
    p_hat = measure_distribution(X, ExactOracle(syn), syn.K)
    assert p_hat.shape == (syn.K,)
    assert np.isclose(p_hat.sum(), 1.0)
    assert abs(p_hat[0] - 0.5) < 0.01


# --- histogram model --------------------------------------------------------------------

def test_histogram_initialise_is_exact():
    syn = _syn(K=32)
    m = HistogramModel(syn, ExactOracle(syn))
    p0 = np.full(syn.K, 1.0 / syn.K)
    m.initialise(p0, np.random.default_rng(0))
    assert np.allclose(m.mode_distribution(np.random.default_rng(0)), p0)


def test_histogram_fit_then_sample_roundtrip():
    syn = _syn(K=16)
    rng = np.random.default_rng(3)
    m = HistogramModel(syn, ExactOracle(syn))
    p = np.array([0.4, 0.3, 0.2] + [0.1 / 13] * 13)
    X, _ = sample_synthetic(p, 100_000, syn, rng)
    m.fit(X, rng)
    drawn = m.sample(100_000, rng)
    p_hat = measure_distribution(drawn, ExactOracle(syn), syn.K)
    assert np.allclose(p_hat, m.mode_distribution(rng), atol=0.01)


def test_make_model_histogram():
    syn = _syn()
    from neural.config import ModelCfg
    model = make_model(ModelCfg(kind="histogram"), syn, ExactOracle(syn))
    assert isinstance(model, HistogramModel)


def test_make_model_rejects_unknown_kind():
    syn = _syn()
    from neural.config import ModelCfg
    with pytest.raises(ValueError):
        make_model(ModelCfg(kind="nope"), syn, ExactOracle(syn))


# --- mode truth reuses Layer 1 ----------------------------------------------------------

def test_mode_truth_is_layer1_truth():
    syn = _syn(K=100, R=10)
    td = make_mode_truth(syn)
    assert td.p_star.shape == (syn.K,)
    assert np.isclose(td.p_star.sum(), 1.0)
    assert td.tail_mask.dtype == bool
    assert len(np.unique(td.regions)) == syn.R


# --- generation loop: schema + determinism ---------------------------------------------

def _cfg(**over) -> dict:
    base = {
        "synthetic": {"K": 64, "R": 1, "zipf_s": 1.1, "init": "truth",
                      "style_len": 2, "style_vocab": 4, "id_base": 2},
        "model": {"kind": "histogram"},
        "dynamics": {"n": 200, "grounding": {"m": 0}},
        "generations": 5,
    }
    base.update(over)
    return base


def test_lineage_returns_layer1_schema():
    df = run_generative_lineage(_cfg(), seed=0)
    assert isinstance(df, pd.DataFrame)
    assert list(df["generation"]) == [0, 1, 2, 3, 4, 5]
    for col in ("heterozygosity", "forward_kl", "tail_mass", "support_size",
                "tail_frac_alive", "head_frac_alive", "tail_truth_mass_alive"):
        assert col in df.columns


def test_lineage_deterministic_given_seed():
    a = run_generative_lineage(_cfg(), seed=7)
    b = run_generative_lineage(_cfg(), seed=7)
    pd.testing.assert_frame_equal(a, b)


def test_lineage_h0_is_truth_heterozygosity():
    # init='truth' -> gen-0 H equals H* of the truth exactly (histogram is exact at gen 0)
    syn = SyntheticCfg(K=64, R=1, zipf_s=1.1, init="truth", style_len=2, style_vocab=4)
    td = make_mode_truth(syn)
    h_star = 1.0 - np.sum(td.p_star ** 2)
    df = run_generative_lineage(_cfg(), seed=1)
    assert abs(df.loc[df["generation"] == 0, "heterozygosity"].iloc[0] - h_star) < 1e-12


def test_dry_lineage_collapses():
    # m=0, small n -> heterozygosity must fall over generations (collapse)
    df = run_generative_lineage(_cfg(generations=40, dynamics={"n": 50,
                                "grounding": {"m": 0}}), seed=2)
    h = df["heterozygosity"].to_numpy()
    assert h[-1] < h[0] - 0.1