"""MNIST-tier tests (skipped when torchvision is absent).

Gate the real-image confirmation tier: the oracle measures modes accurately, thickness binning
and mode assignment are well-formed, the sampler draws the requested modes, the VAE clears a
floor-aware gen-0 fidelity check, and a short dry lineage collapses more than a grounded one.
These are sign checks (the tier is confirmation-only), and they reuse the cached classifier so
they stay fast.
"""

from __future__ import annotations

import numpy as np
import pytest

pytest.importorskip("torchvision")

from inheritance.metrics import forward_kl  # noqa: E402
from neural.config import MnistCfg, ModelCfg, OracleCfg  # noqa: E402
from neural.mnist_data import (  # noqa: E402
    MnistSampler, assign_modes, fit_thickness_thresholds, make_mnist_truth, thickness_bin,
)
from neural.mnist_loop import run_mnist_lineage  # noqa: E402
from neural.mnist_oracle import build_oracle, confusion_summary  # noqa: E402
from neural.mnist_vae import ConvVAEGenerator  # noqa: E402


@pytest.fixture(scope="module")
def world():
    """Load MNIST + build the oracle/sampler/truth once (cached CNN -> fast)."""
    from neural.mnist_data import load_mnist

    cfg = MnistCfg()
    data = load_mnist(cfg.data_root)
    td = make_mnist_truth(cfg)
    oracle, cuts, _ = build_oracle(cfg, OracleCfg(), data, seed=1)
    modes = assign_modes(data.train_x, data.train_y, cuts, cfg)
    sampler = MnistSampler(data.train_x, modes, cfg.K)
    return cfg, data, td, oracle, cuts, sampler


def test_thickness_bins_and_mode_bijection(world):
    cfg, data, td, oracle, cuts, sampler = world
    sub = data.train_x[:2000]
    lab = data.train_y[:2000]
    bins = thickness_bin(sub, lab, cuts)
    assert bins.min() >= 0 and bins.max() < cfg.style_bins            # bins in range
    modes = assign_modes(sub, lab, cuts, cfg)
    assert np.array_equal(modes, lab * cfg.style_bins + bins)          # exact bijection
    assert modes.min() >= 0 and modes.max() < cfg.K


def test_sampler_draws_requested_modes(world):
    cfg, data, td, oracle, cuts, sampler = world
    counts = np.zeros(cfg.K, dtype=int)
    counts[5] = 30            # mode 5 -> class 5//style_bins
    counts[17] = 10           # mode 17 -> class 17//style_bins
    imgs = sampler.draw(counts, np.random.default_rng(0))
    assert imgs.shape == (40, 1, 28, 28)
    # The sampler draws from the true per-mode pools, so the oracle should recover the two
    # requested classes for the large majority of the draw (allowing the ~2% oracle error).
    pred_classes = oracle.classify(imgs) // cfg.style_bins
    want = {5 // cfg.style_bins, 17 // cfg.style_bins}
    assert np.mean(np.isin(pred_classes, list(want))) > 0.9


def test_oracle_mode_accuracy_high(world):
    cfg, data, td, oracle, cuts, sampler = world
    conf = confusion_summary(oracle, data, cuts, cfg)
    assert conf["mode_accuracy"] > 0.95           # solid measurement-noise floor
    assert conf["class_accuracy"] > 0.97


def test_vae_gen0_recovers_full_support(world):
    # gen-0 fidelity gate (floor-aware): the VAE must represent every mode (incl. the tail),
    # else later tail loss would be underfitting, not collapse. Frequencies over-smooth (KL~0.5).
    cfg, data, td, oracle, cuts, sampler = world
    rng = np.random.default_rng(0)
    mcfg = ModelCfg(kind="convvae", latent=32, epochs=15, n_eval=8000)
    X0 = sampler.draw(rng.multinomial(6000, td.p_star), rng)
    vae = ConvVAEGenerator(cfg, mcfg, oracle)
    vae.fit(X0, rng)
    p_hat = vae.mode_distribution(rng)
    assert float(np.mean(p_hat[td.tail_mask] > 1e-9)) > 0.8      # tail represented
    assert forward_kl(td.p_star, p_hat, 1e-9) < 1.2             # loose, floor-aware


def test_dry_collapses_more_than_grounded(world):
    # The N1+N2 signs on real images: dry end-KL > grounded end-KL. Tiny config for speed.
    cfg, data, td, oracle, cuts, sampler = world
    base = {
        "mnist": cfg.__dict__,
        "model": ModelCfg(kind="convvae", latent=32, epochs=20, n_eval=6000).__dict__,
        "generations": 8,
        "metrics": {"kl_floor": 1e-9, "support_eps": 1e-9},
    }

    def run(g):
        n = 6000
        m = 0 if g == 0 else round(n * g / (1 - g))
        c = {**base, "dynamics": {"n": n, "grounding": {"m": m, "policy": "proportional"}}}
        return run_mnist_lineage(c, seed=0, oracle=oracle, sampler=sampler, td=td)

    dry, grd = run(0.0), run(0.1)
    assert dry["forward_kl"].iloc[-1] > dry["forward_kl"].iloc[0] + 1.0    # dry collapses
    assert grd["forward_kl"].iloc[-1] < dry["forward_kl"].iloc[-1]         # grounding arrests it
    assert grd["support_size"].iloc[-1] > dry["support_size"].iloc[-1]     # keeps more modes