"""Pilot grid: partial adjustment of the capital stock toward the paper's supply schedule.

Usage: python -I run_grid.py PKG_DIR OUT.json
Runs the three scenarios at long-run elasticities eps = 3 (the paper) and 6 (the Moll, Rachel and
Restrepo value the paper halves, p. 9 footnote) and capital half-lives of 0 (the paper's
instantaneous schedule), 6, 12, 24 and 48 months. Everything else stays at the paper's values.
Reads Table 3's rows in January 2030 (and January 2028) and yearly paths for the charts.
"""
import json
import math
import sys
from dataclasses import replace

sys.path.insert(0, sys.argv[1])
from econ_scenarios import EXTREME, MODEST, SUBSTANTIAL, Calibration, readout, simulate  # noqa: E402

SCENARIOS = (MODEST, SUBSTANTIAL, EXTREME)
EPS = (3.0, 6.0)
HALF_LIVES = (0.0, 6.0, 12.0, 24.0, 48.0)
KEYS = ("gdp", "gdp_growth", "wage_avg", "wage_C", "wage_N", "net_return", "capital", "labor_share",
        "employment_C", "unemployment_C", "unemployment", "wage_bill_C")
YEARS = (2024.0, 2025.0, 2026.0, 2027.0, 2028.0, 2029.0, 2030.0)


def multiple(rep, cal):
    """GDP gain over the cognitive wage-bill loss, both in percent of no-AI GDP (the p. 35 ratio)."""
    loss = -cal.s_L * cal.share_C * rep["wage_bill_C"]
    return rep["gdp"] / loss if loss > 0 else math.nan


rows = []
for s in SCENARIOS:
    for eps in EPS:
        cal = replace(Calibration(), eps=eps)
        for hl in HALF_LIVES:
            sim = simulate(s, cal, capital_half_life=hl)
            rep = readout(sim)
            rep28 = readout(sim, t=2028.0)
            row = {"scenario": s.name, "eps": eps, "half_life_months": hl,
                   **{k: rep[k] for k in KEYS}, "multiple": multiple(rep, cal),
                   "jan2028": {k: rep28[k] for k in ("gdp", "unemployment_C", "labor_share")},
                   "path": {k: [readout(sim, t=y)[k] if y >= 2025.0 else None for y in YEARS]
                            for k in ("gdp", "unemployment_C", "labor_share", "capital")}}
            rows.append(row)

# Reference: the paper's own instantaneous eps = 1 column (Table 5), for comparison.
for s in (SUBSTANTIAL, EXTREME):
    cal = replace(Calibration(), eps=1.0)
    rep = readout(simulate(s, cal))
    rows.append({"scenario": s.name, "eps": 1.0, "half_life_months": 0.0, "reference": "Table 5",
                 **{k: rep[k] for k in KEYS}, "multiple": multiple(rep, cal)})

json.dump({"years": YEARS, "rows": rows}, open(sys.argv[2], "w"), indent=1)

print(f"{'scenario':12} {'eps':>4} {'half-life':>9} {'GDP':>6} {'growth':>6} {'LS':>5} {'uC':>5} "
      f"{'u':>5} {'wC':>6} {'wN':>6} {'wavg':>6} {'r':>5} {'K':>6} {'mult':>5}")
for r in rows:
    tag = " (Table 5)" if r.get("reference") else ""
    print(f"{r['scenario']:12} {r['eps']:4.0f} {r['half_life_months']:9.0f} {r['gdp']:6.2f} {r['gdp_growth']:6.2f} "
          f"{r['labor_share']:5.1f} {r['unemployment_C']:5.2f} {r['unemployment']:5.2f} {r['wage_C']:6.2f} "
          f"{r['wage_N']:6.2f} {r['wage_avg']:6.2f} {r['net_return']:5.2f} {r['capital']:6.1f} {r['multiple']:5.2f}{tag}")