"""SIMULATED — factorial experiments + ablations for H3. Usage: python -m wp35.simulations.run_experiments [--quick]
Design (pre-declared here before running): N x network x bandwidth x plan-entropy x shock x regime x seeds, plus an
admin-capacity factor for the two hierarchical regimes. Outputs results/simulated/runs_main.parquet, runs_ablation.parquet.
"""
from __future__ import annotations
import itertools, sys, time
from multiprocessing import Pool
from pathlib import Path
import pandas as pd
from wp35.simulations.model import Params, simulate

ROOT = Path(__file__).resolve().parents[3]; OUT = ROOT / "results" / "simulated"; OUT.mkdir(parents=True, exist_ok=True)
SEEDS = list(range(8))
REGIMES = ["none", "norms", "records", "central", "federated"]


def main_design(quick=False):
    Ns = [50, 200, 800] if not quick else [50, 200]
    cells = []
    for N, net, bw, zipf, shock, reg, seed in itertools.product(Ns, ["lattice", "smallworld", "random"], [2, 8], [0.0, 1.0], ["none", "distribution", "network"], REGIMES, SEEDS if not quick else SEEDS[:2]):
        base = dict(N=N, network=net, bandwidth=bw, zipf=zipf, shock_type=shock, shock_period=20 if shock != "none" else -1, regime=reg, seed=seed, T=40)
        if reg in ("central", "federated"):
            for ratio in [0.5, 2.0]:
                cells.append(Params(**base, admin_capacity_ratio=ratio))
        else:
            cells.append(Params(**base, admin_capacity_ratio=2.0))
    return cells


def ablation_design(quick=False):
    cells = []
    seeds = SEEDS if not quick else SEEDS[:2]
    base = dict(N=200, network="lattice", bandwidth=8, zipf=1.0, T=40)
    for seed in seeds:
        for reg in REGIMES:
            r = 2.0
            cells.append(Params(**base, regime=reg, seed=seed, admin_capacity_ratio=r, shock_type="none", shock_period=-1))  # reference
        # records without verification
        cells.append(Params(**base, regime="records", seed=seed, verify_records=False, admin_capacity_ratio=2.0))
        # code adaptation speed after a distribution shock
        for ad in [0, 10, 1000]:
            cells.append(Params(**base, regime="norms", seed=seed, shock_type="distribution", shock_period=20, adapt_delay=ad, shock_kl_target=2.0, admin_capacity_ratio=2.0))
        # horizon (break-even, Theorem 2) and learning cost
        for T in [10, 40, 160]:
            for lb in [0.0, 200.0, 1000.0]:
                for reg in ["none", "norms"]:
                    cells.append(Params(**{**base, "T": T}, regime=reg, seed=seed, learn_bits=lb, admin_capacity_ratio=2.0))
        # break-even with a large, concentrated plan space (high coding gain)
        for T in [10, 40, 160]:
            for lb in [200.0, 1000.0]:
                for reg in ["none", "norms"]:
                    cells.append(Params(**{**base, "T": T, "zipf": 1.5}, regime=reg, seed=seed, s=256, learn_bits=lb, admin_capacity_ratio=2.0))
        # heterogeneity of information (eps) for none vs norms
        for eps in [0.05, 0.25, 0.6]:
            for reg in ["none", "norms"]:
                cells.append(Params(**base, regime=reg, seed=seed, eps=eps, admin_capacity_ratio=2.0))
        # admin capacity sweep (phase diagram, Theorem 5) and branching
        for ratio in [0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 4.0]:
            for reg in ["central", "federated"]:
                cells.append(Params(**base, regime=reg, seed=seed, admin_capacity_ratio=ratio))
        for b in [4, 16]:
            cells.append(Params(**base, regime="federated", seed=seed, branching=b, admin_capacity_ratio=1.0))
        # plan-set size / entropy sweep (coding gain, Theorem 1)
        for s_, z in [(16, 0.0), (64, 0.0), (256, 0.0), (64, 0.5), (64, 1.5), (64, 2.5)]:
            for reg in ["none", "norms"]:
                cells.append(Params(**{**base, "zipf": z}, regime=reg, seed=seed, s=s_, admin_capacity_ratio=2.0))
    return cells


def _run(P):
    return simulate(P)


def run(cells, path, n_jobs=12):
    t0 = time.time()
    with Pool(n_jobs) as pool:
        rows = pool.map(_run, cells, chunksize=16)
    df = pd.DataFrame(rows); df.to_parquet(path, index=False)
    print(f"{path.name}: {len(df)} runs in {time.time() - t0:.1f}s")
    return df


def main():
    quick = "--quick" in sys.argv
    main_cells = main_design(quick); abl = ablation_design(quick)
    print("main design cells:", len(main_cells), "| ablation cells:", len(abl))
    run(main_cells, OUT / ("runs_main_quick.parquet" if quick else "runs_main.parquet"))
    run(abl, OUT / ("runs_ablation_quick.parquet" if quick else "runs_ablation.parquet"))


if __name__ == "__main__":
    main()