"""Experiment runner: sweep a parameter grid x replicates, write reproducible artifacts.

One YAML per experiment (blueprint 2.7 / 4). ``run_experiment`` expands the declared
sweep grid, runs ``run_lineage`` for every (combo, replicate) at seeds derived from a
single master seed, and returns long-form results. ``save_artifacts`` writes the output
contract: ``results.parquet`` + ``resolved_config.yaml`` + ``manifest.json``. Figures are
regenerated separately from ``results.parquet`` alone.

CLI:  python -m knowledge.experiment configs/layer1/E2.yaml
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import itertools
import json
import subprocess
import sys
from importlib.metadata import version
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import yaml

from .lineage import run_lineage
from .seeding import spawn_seeds
from .teachers import make_correlated_teachers
from .truth import make_true_distribution

# Keys that make up a single-lineage configuration (everything else is experiment-level).
_LINEAGE_KEYS = ("truth", "dynamics", "generations", "metrics")


def _set_by_path(d: dict, path: str, value: Any) -> None:
    """Set a nested dict value by a dotted path, creating intermediate dicts."""
    keys = path.split(".")
    cur = d
    for k in keys[:-1]:
        cur = cur.setdefault(k, {})
    cur[keys[-1]] = value


def _apply_param(lineage_cfg: dict, param: str, value: Any) -> dict:
    """Apply one swept parameter to a lineage config; return label columns to record.

    Special-cases ``g`` (the grounding fraction): converts to an integer real-sample
    budget ``m = round(n*g/(1-g))`` (so g=0 -> m=0), records both ``g`` and ``m``.
    Any other param is a dotted path into the lineage config.
    """
    if param == "g":
        n = lineage_cfg["dynamics"]["n"]
        g = float(value)
        m = 0 if g <= 0.0 else int(round(n * g / (1.0 - g)))
        lineage_cfg["dynamics"].setdefault("grounding", {})["m"] = m
        return {"g": g, "m": m}
    if param == "arm":
        # A named arm bundling several overrides applied together (e.g. E6 varies
        # grounding + remint settings jointly). value = {name, set: {dotted.path: v}}.
        for path, v in value.get("set", {}).items():
            _set_by_path(lineage_cfg, path, v)
        return {"arm": value["name"]}
    _set_by_path(lineage_cfg, param, value)
    return {param.split(".")[-1]: value}


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

    ``sweep`` is a list of ``{param, values}`` entries; the Cartesian product is taken
    (so E4 can sweep K_T x rho). A missing/empty sweep yields a single run.

    Returns:
        list[tuple[dict, dict]]: One (label-columns, lineage-config) pair per grid point.
    """
    base = {k: copy.deepcopy(cfg[k]) for k in _LINEAGE_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 the full experiment: every grid point x every replicate.

    Replicate seeds are derived once from the master seed and *reused across grid points*,
    so comparisons across sweep values are paired (shared drift noise) — variance
    reduction, and a pure function of the master seed.

    Args:
        cfg (dict): Parsed experiment YAML (experiment, seed, n_replicates, truth,
            dynamics, generations, metrics, sweep, output).

    Returns:
        pd.DataFrame: Long-form results; 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, lineage_cfg in combos:
        for rep, ss in enumerate(seeds):
            df = run_lineage(lineage_cfg, ss)
            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 run_coverage(cfg: dict) -> pd.DataFrame:
    """E4 runner: multi-teacher recombination coverage (blueprint 2.5-E4 / 2.7.1).

    A single distillation step, not a lineage. For each (K_T, rho) grid point and
    replicate: build K_T correlated teachers (marginal retention q, pairwise
    correlation rho), then measure two coverages of the tail:

    * ``union_coverage`` — fraction of tail items retained by >=1 teacher (the
      construction-level union U(K_T, rho, q); must match the closed form). This is the
      recombination *supply*.
    * ``surviving_mean`` / ``surviving_max`` — fraction of tail items that survive the
      pupil's size-n resampling (+ optional grounding m) under two recombination
      operators: ``mean`` (blueprint mean-mixture distillation) and ``max`` (union-
      preserving model-merge, à la M2N2). Under ``mean`` the union gain is diluted by
      1/K_T and (in the rare-tail linear regime) is exactly cancelled — expected pupil
      tail mass is conserved at q·(tail mass) regardless of K_T, so surviving is flat.
      Under ``max`` each item keeps its strongest teacher, so surviving rises with K_T
      and with decorrelation. The gap ``union - surviving`` is the tail recombination
      supplied but drift/dilution re-erased.

    Matched budget: the pupil draws n samples total from the combined teachers
    (equivalently n/K_T each), so more teachers != more data.
    """
    truth, cov = cfg["truth"], cfg["coverage"]
    td = make_true_distribution(
        truth["K"], truth["R"], truth["tail"], truth["tail_frac"], truth["zipf_s"], 0,
        tail_threshold=truth["tail_threshold"],
    )
    tail_idx = np.flatnonzero(td.tail_mask)
    n, q = int(cov["n"]), float(cov["q"])
    retain_thresh = 1e-8  # dropped tails sit at ~tail_floor (1e-9); retained at ~p*_j

    sweeps = cfg["sweep"]
    if isinstance(sweeps, dict):
        sweeps = [sweeps]
    params = [s["param"] for s in sweeps]
    value_lists = [list(s["values"]) for s in sweeps]
    seeds = spawn_seeds(int(cfg["seed"]), int(cfg["n_replicates"]))

    rows: list[dict] = []
    for combo in itertools.product(*value_lists):
        d = dict(zip(params, combo))
        K_T, rho = int(d["K_T"]), float(d["rho"])
        g = float(d.get("g", cov.get("g", 0.0)))          # g may be swept or fixed
        m = 0 if g <= 0.0 else int(round(n * g / (1.0 - g)))
        for rep, ss in enumerate(seeds):
            child = int(ss.generate_state(1)[0])
            teachers = np.asarray(make_correlated_teachers(
                td.p_star, td.tail_mask, K_T, rho, q, seed=child))
            retained = teachers[:, tail_idx] > retain_thresh          # (K_T, T)
            union = float(np.mean(retained.any(axis=0)))
            surviving = {}
            for offset, combine in ((1, teachers.mean), (2, teachers.max)):
                p = combine(axis=0)
                p = p / p.sum()
                rng = np.random.default_rng(child + offset)
                counts = rng.multinomial(n, p)
                if m > 0:
                    counts = counts + rng.multinomial(m, td.p_star)
                surviving[offset] = float(np.mean(counts[tail_idx] > 0))
            rows.append({
                "experiment": cfg["experiment"], "K_T": K_T, "rho": rho,
                "replicate": rep, "union_coverage": union,
                "surviving_mean": surviving[1], "surviving_max": surviving[2],
                "g": g, "q": q,
            })
    return pd.DataFrame(rows)


def _git_commit() -> str | None:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL, text=True
        ).strip()
    except Exception:
        return None


def _content_hash(path: Path) -> str:
    h = hashlib.sha256()
    h.update(path.read_bytes())
    return h.hexdigest()


def save_artifacts(cfg: dict, df: pd.DataFrame, out_dir: Path,
                   extra_libs: tuple[str, ...] = (),
                   extra_manifest: dict | None = None,
                   grid: list | None = None) -> None:
    """Write the reproducibility output contract (blueprint 2.7 / 4).

    Writes ``results.parquet``, ``resolved_config.yaml`` (the fully-expanded config), and
    ``manifest.json`` (library versions, master seed, git commit, content hash).

    Args:
        cfg (dict): The parsed experiment config.
        df (pd.DataFrame): The long-form results.
        out_dir (Path): Output directory.
        extra_libs (tuple[str, ...]): Extra library names to record versions for (e.g.
            ``torch``, ``torchvision`` for Layer 1.5). Missing libraries are skipped, so a
            caller can pass optional deps unconditionally.
        extra_manifest (dict | None): Extra key/value pairs to merge into the manifest
            (e.g. model architecture, oracle checkpoint hash, determinism flags).
        grid (list | None): Pre-expanded ``[{label, lineage_cfg}, ...]`` to record in the
            resolved config. If None, it is computed via ``expand_sweeps`` for the Layer-1
            ``lineage`` kind (a caller with a different schema, e.g. Layer 1.5, passes its
            own expanded grid here).
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    results_path = out_dir / "results.parquet"
    df.to_parquet(results_path, index=False)

    resolved = {
        "experiment": cfg["experiment"],
        "seed": cfg["seed"],
        "n_replicates": cfg["n_replicates"],
        "source_config": cfg,
    }
    if grid is not None:
        resolved["grid"] = grid
    elif cfg.get("kind", "lineage") == "lineage":
        resolved["grid"] = [
            {"label": label, "lineage_cfg": lineage_cfg}
            for label, lineage_cfg in expand_sweeps(cfg)
        ]
    (out_dir / "resolved_config.yaml").write_text(yaml.safe_dump(resolved, sort_keys=False))

    libraries: dict[str, str] = {}
    for lib in ("numpy", "scipy", "pandas", "pyarrow") + tuple(extra_libs):
        try:
            libraries[lib] = version(lib)
        except Exception:  # optional dep not installed -> omit rather than crash
            pass
    manifest = {
        "experiment": cfg["experiment"],
        "master_seed": cfg["seed"],
        "git_commit": _git_commit(),
        "python": sys.version.split()[0],
        "libraries": libraries,
        "rows": int(len(df)),
        "results_sha256": _content_hash(results_path),
    }
    if extra_manifest:
        manifest.update(extra_manifest)
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))


def run_and_save(config_path: str | Path) -> Path:
    """Load an 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']}"))
    df = run_coverage(cfg) if cfg.get("kind") == "coverage" else run_experiment(cfg)
    save_artifacts(cfg, df, out_dir)
    return out_dir


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


if __name__ == "__main__":
    main()