"""The neural analogue of ``knowledge.lineage.run_lineage``.

Runs ``T`` generations of *train-a-model-on-the-previous-model's-samples*, the neural
image of the Wright-Fisher generational step. Each generation the pupil is trained on a
pool of (i) ``n`` observations drawn from the parent model (drift) and (ii) ``m`` fresh
observations drawn from the grounding reference (immigration, ``g = m/(n+m)``), then its
oracle-measured mode distribution is logged with the *same* metric schema Layer 1 uses.
Grounding structure (proportional / uniform / matched over regions) and the re-mint gate
reuse ``knowledge.step`` and mirror ``run_lineage`` exactly, so a histogram-model lineage
reproduces the analytic core and a neural-model lineage tests whether the same signs hold
in real weights.
"""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np
import pandas as pd

from knowledge.metrics import heterozygosity
from knowledge.step import allocate_m, structured_multinomial
from knowledge.truth import uniform_init

from .config import NeuralLineageCfg
from .evaluate import measure_metrics
from .models import make_model
from .oracle import ExactOracle
from .synthetic import id_codewords, make_mode_truth, render_modes


def _counts_to_observations(counts: np.ndarray, cfg, rng, codewords) -> np.ndarray:
    """Expand a per-mode count vector into rendered token sequences."""
    modes = np.repeat(np.arange(counts.size), counts)
    return render_modes(modes, cfg, rng, codewords)


def run_generative_lineage(cfg: Mapping[str, Any] | NeuralLineageCfg,
                           seed: int) -> pd.DataFrame:
    """Run one neural lineage and return per-generation metrics.

    Args:
        cfg (Mapping | NeuralLineageCfg): Resolved neural-lineage configuration.
        seed (int): Seed for this replicate; the run is a pure function of (cfg, seed) for
            the histogram model (statistically reproducible for torch models).

    Returns:
        pd.DataFrame: One row per generation 0..T with the same metric columns as
        ``knowledge.lineage.run_lineage``.
    """
    cfg = NeuralLineageCfg.from_dict(cfg)
    syn = cfg.synthetic
    td = make_mode_truth(syn)
    p_star_orig = td.p_star            # forward_kl is always vs the original truth
    regions = td.regions
    tail_mask = td.tail_mask
    R = syn.R

    rng = np.random.default_rng(seed)
    oracle = ExactOracle(syn)
    codewords = id_codewords(syn)

    # Initial distribution over modes (exact, like Layer 1).
    if syn.init == "uniform":
        p0 = uniform_init(syn.K)
    elif syn.init == "truth":
        p0 = p_star_orig.copy()
    else:
        raise ValueError(f"unknown init {syn.init!r} (expected uniform|truth)")

    # Grounding wiring (reused verbatim from Layer 1).
    grounding = cfg.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)
    p_star_eff = p_star_orig.copy()    # grounding reference; may be re-minted (N6)
    remint = cfg.dynamics.remint
    n = cfg.dynamics.n

    model = make_model(cfg.model, syn, oracle)
    model.initialise(p0, rng)

    rows: list[dict] = []

    def record(t: int, p: np.ndarray) -> None:
        row = {"generation": t}
        row.update(measure_metrics(p, p_star_orig, tail_mask, regions, R, cfg.metrics))
        rows.append(row)

    record(0, model.mode_distribution(rng))

    for t in range(1, cfg.generations + 1):
        X_syn = model.sample(n, rng)                      # drift: n from the parent
        if m_vector is not None:                          # immigration: m grounded samples
            counts_real = structured_multinomial(
                m_vector, p_star_eff, regions, grounding.policy, rng)
            X_real = _counts_to_observations(counts_real, syn, rng, codewords)
            pool = np.concatenate([X_syn, X_real], axis=0)
        else:
            pool = X_syn

        pupil = make_model(cfg.model, syn, oracle)
        pupil.fit(pool, rng)
        model = pupil
        p = model.mode_distribution(rng)

        if remint.enabled and remint.period and t % remint.period == 0:
            # Founder event: current distribution becomes the new grounding reference and
            # the original truth is discarded for grounding. Gated on diversity (N6).
            if remint.H_gate is None or heterozygosity(p) >= remint.H_gate:
                p_star_eff = p.copy()
        record(t, p)

    return pd.DataFrame(rows)