"""Dump every monthly series for the runs the paper prints (Tables 3, 5, 6), from a given package.

Usage: python -I dump_series.py PKG_DIR OUT.json [--half-life MONTHS] [--eps-zero]
PKG_DIR holds an ``econ_scenarios`` package (the unmodified Econ-ARK copy or the pilot variant).
--half-life is passed only to the variant. --eps-zero runs every case at eps = 0 instead.
"""
import argparse
import inspect
import json
import math
import sys
from dataclasses import replace

parser = argparse.ArgumentParser()
parser.add_argument("pkg")
parser.add_argument("out")
parser.add_argument("--half-life", type=float, default=None)
parser.add_argument("--eps-zero", action="store_true")
args = parser.parse_args()
sys.path.insert(0, args.pkg)

from econ_scenarios import EXTREME, MODEST, SUBSTANTIAL, Calibration, simulate  # noqa: E402

CASES = [(s, {}) for s in (MODEST, SUBSTANTIAL, EXTREME)]
CASES += [(s, {"eps": e}) for s in (SUBSTANTIAL, EXTREME) for e in (1.0, 6.0, math.inf)]
CASES += [(s, {"xi": x}) for s, xs in ((SUBSTANTIAL, (0.75, 0.9)), (EXTREME, (0.0, 0.75, 0.9))) for x in xs]

kwargs = {}
if args.half_life is not None:
    if "capital_half_life" not in inspect.signature(simulate).parameters:
        raise SystemExit("this package has no capital_half_life switch")
    kwargs["capital_half_life"] = args.half_life

result = {}
for scenario, over in CASES:
    if args.eps_zero:
        if "eps" in over:
            continue
        over = {**over, "eps": 0.0}
    cal = replace(Calibration(), **over)
    sim = simulate(scenario, cal, **kwargs)
    key = scenario.name + "".join(f"|{k}={v}" for k, v in sorted(over.items()))
    result[key] = {name: [float(v) for v in values] for name, values in sim.series.items()}

with open(args.out, "w") as fh:
    json.dump(result, fh)
print(f"{len(result)} runs written to {args.out}")