"""Neural (Layer 1.5) experiment runner: sweep a grid x replicates, write artifacts.

Mirrors ``inheritance.experiment`` and reuses its sweep-expansion primitives
(``_apply_param`` — including the ``g -> m`` conversion — and ``_set_by_path``), its
provenance helpers, and its output contract (``save_artifacts``). Only the per-run call and
the config key set differ: a neural run trains generative models rather than resampling a
frequency vector, and its config groups are ``synthetic``/``model``/``dynamics``/... .

CLI:  python -m neural.experiment configs/neural/bridge.yaml
"""

from __future__ import annotations

import argparse
import copy
import itertools
from pathlib import Path
from typing import Any

import pandas as pd
import yaml

from inheritance.experiment import _apply_param, save_artifacts
from inheritance.seeding import spawn_seeds

from .generation_loop import run_generative_lineage

# Config groups that make up a single neural lineage (everything else is experiment-level).
_NEURAL_KEYS = ("synthetic", "model", "dynamics", "generations", "metrics", "n_eval")

# Config groups for one MNIST lineage (the real-image tier; `mnist` replaces `synthetic`).
_MNIST_KEYS = ("mnist", "model", "dynamics", "generations", "metrics")

# Libraries recorded in the manifest on top of the Layer-1 core set (skipped if absent).
_EXTRA_LIBS = ("torch", "torchvision")


def expand_sweeps(cfg: dict) -> list[tuple[dict, dict]]:
    """Expand the sweep grid into (label, resolved_neural_cfg) pairs.

    Identical semantics to ``inheritance.experiment.expand_sweeps`` (Cartesian product of the
    declared ``{param, values}`` entries, reusing ``_apply_param`` for the ``g -> m`` and
    ``arm`` special cases) but assembling the base from the neural config groups.

    Returns:
        list[tuple[dict, dict]]: One (label-columns, neural-config) pair per grid point.
    """
    base = {k: copy.deepcopy(cfg[k]) for k in _NEURAL_KEYS if k in cfg}
    sweeps = cfg.get("sweep", [])
    if isinstance(sweeps, dict):
        sweeps = [sweeps]
    if not sweeps:
        return [({}, base)]
    params = [s["param"] for s in sweeps]
    value_lists = [list(s["values"]) for s in sweeps]
    combos: list[tuple[dict, dict]] = []
    for values in itertools.product(*value_lists):
        lin = copy.deepcopy(base)
        label: dict = {}
        for param, val in zip(params, values):
            label.update(_apply_param(lin, param, val))
        combos.append((label, lin))
    return combos


def run_experiment(cfg: dict) -> pd.DataFrame:
    """Run every grid point x every replicate; return long-form results.

    Replicate seeds are derived once from the master seed and reused across grid points, so
    comparisons across sweep values are paired (shared drift noise) — as in Layer 1.

    Args:
        cfg (dict): Parsed experiment YAML.

    Returns:
        pd.DataFrame: One row per (combo, replicate, generation).
    """
    name = cfg["experiment"]
    master = int(cfg["seed"])
    n_rep = int(cfg["n_replicates"])
    combos = expand_sweeps(cfg)
    seeds = spawn_seeds(master, n_rep)

    frames: list[pd.DataFrame] = []
    for label, neural_cfg in combos:
        for rep, ss in enumerate(seeds):
            df = run_generative_lineage(neural_cfg, int(ss.generate_state(1)[0]))
            for col, val in label.items():
                df[col] = val
            df["replicate"] = rep
            frames.append(df)
    out = pd.concat(frames, ignore_index=True)
    out.insert(0, "experiment", name)
    return out


def _expand_mnist(cfg: dict) -> list[tuple[dict, dict]]:
    """Expand the MNIST g-sweep into (label, resolved-lineage) pairs (reuses ``_apply_param``)."""
    base = {k: copy.deepcopy(cfg[k]) for k in _MNIST_KEYS if k in cfg}
    sweeps = cfg.get("sweep", [])
    if isinstance(sweeps, dict):
        sweeps = [sweeps]
    if not sweeps:
        return [({}, base)]
    params = [s["param"] for s in sweeps]
    value_lists = [list(s["values"]) for s in sweeps]
    combos: list[tuple[dict, dict]] = []
    for values in itertools.product(*value_lists):
        lin = copy.deepcopy(base)
        label: dict = {}
        for param, val in zip(params, values):
            label.update(_apply_param(lin, param, val))
        combos.append((label, lin))
    return combos


def run_mnist_experiment(cfg: dict) -> tuple[pd.DataFrame, dict]:
    """Run every MNIST grid point x replicate; return (results, oracle-provenance manifest).

    The frozen classifier oracle, per-mode real-image pools, and ``p*`` are built **once** and
    shared across all arms/replicates (training the CNN and indexing the pools is expensive).
    The oracle's confusion matrix over the real test set is returned for the manifest as the
    measurement-noise floor.

    Args:
        cfg (dict): Parsed MNIST experiment YAML (``mnist``/``model``/``dynamics``/``oracle``
            blocks, a ``sweep`` over ``g``, ``seed``, ``n_replicates``).

    Returns:
        tuple[pd.DataFrame, dict]: Long-form results and the oracle-provenance manifest dict.
    """
    from inheritance.config import _sub

    from .config import MnistCfg, OracleCfg
    from .mnist_data import assign_modes, load_mnist, make_mnist_truth, MnistSampler
    from .mnist_loop import run_mnist_lineage
    from .mnist_oracle import build_oracle, confusion_summary

    mnist_cfg = _sub(cfg["mnist"], MnistCfg)
    oracle_cfg = _sub(cfg.get("oracle", {}), OracleCfg)
    master = int(cfg["seed"])
    n_rep = int(cfg["n_replicates"])

    data = load_mnist(mnist_cfg.data_root)
    td = make_mnist_truth(mnist_cfg)
    oracle, cuts, ckpt_hash = build_oracle(mnist_cfg, oracle_cfg, data, seed=master)
    modes = assign_modes(data.train_x, data.train_y, cuts, mnist_cfg)
    sampler = MnistSampler(data.train_x, modes, mnist_cfg.K)
    conf = confusion_summary(oracle, data, cuts, mnist_cfg)

    combos = _expand_mnist(cfg)
    seeds = spawn_seeds(master, n_rep)
    frames: list[pd.DataFrame] = []
    for label, lineage_cfg in combos:
        for rep, ss in enumerate(seeds):
            df = run_mnist_lineage(lineage_cfg, int(ss.generate_state(1)[0]), oracle, sampler, td)
            for col, val in label.items():
                df[col] = val
            df["replicate"] = rep
            frames.append(df)
    out = pd.concat(frames, ignore_index=True)
    out.insert(0, "experiment", cfg["experiment"])

    manifest = {
        "layer": "1.5", "tier": "mnist", "model_kind": cfg.get("model", {}).get("kind"),
        "oracle_ckpt_sha256": ckpt_hash,
        "oracle_mode_accuracy": conf["mode_accuracy"],
        "oracle_class_accuracy": conf["class_accuracy"],
        "confusion_matrix": conf["confusion"],
    }
    return out, manifest


def run_and_save(config_path: str | Path) -> Path:
    """Load a neural experiment YAML, run it, and write artifacts. Returns the output dir."""
    config_path = Path(config_path)
    cfg = yaml.safe_load(config_path.read_text())
    out_dir = Path(cfg.get("output", {}).get("dir", f"results/{cfg['experiment']}"))
    kind = cfg.get("kind", "gen_lineage")
    extra_manifest = {"layer": "1.5", "model_kind": cfg.get("model", {}).get("kind", "histogram")}
    if kind == "recombination":
        from .recombine import run_recombination  # the `recombination` experiment; lazy import
        df = run_recombination(cfg)
        grid = None
    elif kind == "gen_lineage":
        df = run_experiment(cfg)
        grid = [{"label": label, "neural_cfg": c} for label, c in expand_sweeps(cfg)]
    elif kind == "mnist_lineage":
        df, extra_manifest = run_mnist_experiment(cfg)   # oracle provenance + confusion matrix
        grid = [{"label": label, "lineage_cfg": c} for label, c in _expand_mnist(cfg)]
    elif kind == "speciation_real":
        from .speciation_real import run_speciation_real   # E13: real-weight speciation + Git Re-Basin
        df = run_speciation_real(cfg, int(cfg["seed"]))
        grid = None
        extra_manifest = {"layer": "1.5", "tier": "speciation_real"}
    else:
        raise ValueError(f"unknown neural experiment kind {kind!r}")
    save_artifacts(cfg, df, out_dir, extra_libs=_EXTRA_LIBS,
                   extra_manifest=extra_manifest, grid=grid)
    return out_dir


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Run a Layer-1.5 neural experiment from a YAML config.")
    parser.add_argument("config", help="Path to configs/neural/<name>.yaml")
    args = parser.parse_args(argv)
    out_dir = run_and_save(args.config)
    print(f"wrote artifacts to {out_dir}/")


if __name__ == "__main__":
    main()