"""The MNIST analogue of ``generation_loop.run_generative_lineage``.

Identical Wright-Fisher generational step — drift (``n`` samples from the parent model) +
grounding (``m`` fresh real samples, ``g = m/(n+m)``) + refit — but observations are **images**:
the parent VAE *generates* the drift images, and grounding *draws real MNIST images* from the
per-mode pools (``MnistSampler``) instead of rendering token sequences. The frozen classifier
oracle reads each generation's mode distribution, logged with the *same* metric schema as every
other tier (``neural.evaluate.measure_metrics``), so the MNIST curves overlay the synthetic ones.

The oracle, sampler, and ``p*`` are built once by the experiment runner and passed in (training
the classifier and indexing the image pools is expensive and shared across all arms/replicates).
"""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np
import pandas as pd

from knowledge.config import MetricsCfg, _sub
from knowledge.step import allocate_m, structured_multinomial
from knowledge.truth import TrueDist, uniform_init

from .config import MnistCfg, ModelCfg, NeuralDynamicsCfg
from .evaluate import measure_metrics
from .mnist_data import MnistSampler
from .mnist_vae import ConvVAEGenerator
from .oracle import Oracle


def _make_mnist_model(model_cfg: ModelCfg, cfg: MnistCfg, oracle: Oracle):
    """Construct the image generative model for the MNIST tier."""
    if model_cfg.kind == "convvae":
        return ConvVAEGenerator(cfg, model_cfg, oracle)
    raise ValueError(f"unknown MNIST model kind {model_cfg.kind!r} (expected convvae)")


def run_mnist_lineage(cfg: Mapping[str, Any], seed: int, oracle: Oracle,
                      sampler: MnistSampler, td: TrueDist) -> pd.DataFrame:
    """Run one MNIST lineage and return per-generation metrics (same schema as Layer 1).

    Args:
        cfg (Mapping): Resolved config with ``mnist``, ``model``, ``dynamics``, ``generations``,
            and optional ``metrics`` blocks.
        seed (int): Replicate seed (statistically reproducible for the VAE).
        oracle (Oracle): Prebuilt frozen classifier oracle.
        sampler (MnistSampler): Prebuilt per-mode real-image pools.
        td (TrueDist): ``p*``, tail mask, regions over the ``K`` modes.

    Returns:
        pd.DataFrame: One row per generation 0..T with the standard metric columns.
    """
    mcfg = _sub(cfg["mnist"], MnistCfg)
    model_cfg = _sub(cfg["model"], ModelCfg)
    dyn_raw = dict(cfg.get("dynamics", {}))
    from knowledge.config import GroundingCfg, RemintCfg
    dynamics = NeuralDynamicsCfg(
        n=dyn_raw.get("n", NeuralDynamicsCfg.n),
        grounding=_sub(dyn_raw.get("grounding", {}), GroundingCfg),
        remint=_sub(dyn_raw.get("remint", {}), RemintCfg),
    )
    metrics = _sub(cfg.get("metrics", {}), MetricsCfg)
    generations = int(cfg.get("generations", 20))

    p_star = td.p_star
    tail_mask, regions, R = td.tail_mask, td.regions, mcfg.R
    n = dynamics.n
    grounding = dynamics.grounding
    exercised = np.asarray(grounding.exercised) if grounding.exercised is not None else None
    m_vector = allocate_m(grounding.m, R, grounding.policy, exercised)

    rng = np.random.default_rng(seed)
    p0 = uniform_init(mcfg.K) if mcfg.init == "uniform" else p_star.copy()

    model = _make_mnist_model(model_cfg, mcfg, oracle)

    def real_images(counts: np.ndarray) -> np.ndarray:
        return sampler.draw(counts, rng)

    # Generation 0: train on n real images drawn from p0 (the loop builds gen-0, not the VAE).
    X0 = real_images(rng.multinomial(n, p0))
    model.fit(X0, rng)

    rows: list[dict] = []

    def record(t: int) -> None:
        row = {"generation": t}
        row.update(measure_metrics(model.mode_distribution(rng), p_star, tail_mask,
                                   regions, R, metrics))
        rows.append(row)

    record(0)
    for t in range(1, generations + 1):
        X_syn = model.sample(n, rng)                      # drift: n images from the parent
        if m_vector is not None:                          # immigration: m real images
            counts_real = structured_multinomial(m_vector, p_star, regions, grounding.policy, rng)
            pool = np.concatenate([X_syn, real_images(counts_real)], axis=0)
        else:
            pool = X_syn
        pupil = _make_mnist_model(model_cfg, mcfg, oracle)
        pupil.fit(pool, rng)
        model = pupil
        record(t)

    return pd.DataFrame(rows)