"""Heterogeneous-agent general equilibrium for a solved scenario.
Usage: python3 scripts/30_ha_equilibrium.py <scenario> [maxit] [lam]
Outputs: results/ha/<scenario>_ha.npz (r path, wealth stats by year, consumption stats), results/ha/<scenario>_ha.json"""
import json
import os
import sys
import time

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from agi_ge.ha_equilibrium import solve_ha_equilibrium  # noqa: E402
from agi_ge.params import Params  # noqa: E402
from agi_ge.transition import TwoClassModel  # noqa: E402

ROOT = os.path.join(os.path.dirname(__file__), "..")
name = sys.argv[1]
maxit = int(sys.argv[2]) if len(sys.argv) > 2 else 40
lam = float(sys.argv[3]) if len(sys.argv) > 3 and not sys.argv[3].startswith('--') else 1.0
SC = os.path.join(ROOT, "results", "scenarios")
p = Params.from_json(os.path.join(SC, f"{name}_params.json"))
u = np.load(os.path.join(SC, f"{name}_u.npy"))
m = TwoClassModel(p)
P = m.evaluate(u, full=True)
t0 = time.time()
ckpt = os.path.join(ROOT, "results", "ha", f"{name}_r_checkpoint.npy")
os.makedirs(os.path.join(ROOT, "results", "ha"), exist_ok=True)
r_init = np.load(ckpt) if (os.path.exists(ckpt) and "--resume" in sys.argv) else None
out, hb, pols, dist_hist = solve_ha_equilibrium(p, m, u, P, lam=lam, maxit=maxit, verbose=True, checkpoint=ckpt, r_init=r_init)
os.makedirs(os.path.join(ROOT, "results", "ha"), exist_ok=True)
stats = out["stats"]
years = np.array([s["year"] for s in stats])
Pp = out["paths"]
np.savez_compressed(os.path.join(ROOT, "results", "ha", f"{name}_ha.npz"), r=out["r"], Wd=out["Wd"], Ws=out["Ws"], years=years,
                    gini=np.array([s["gini"] for s in stats]), top1=np.array([s["top1"] for s in stats]),
                    top10=np.array([s["top10"] for s in stats]), bottom50=np.array([s["bottom50"] for s in stats]),
                    X=out["X"], r_twoclass=P["r"], W_twoclass=P["W"], GDP=Pp["GDP"], labor_share=Pp["labor_share"], wL=Pp["wL"], wH=Pp["wH"])
meta = dict(name=name, converged=bool(out["converged"]), hist=[float(h) for h in out["hist"]], beta_high=out["hcal"]["beta_high"],
            stats_2026=stats[0], seconds=time.time() - t0,
            stats_by_year={int(y): {k: float(v) for k, v in s.items() if k != "year"} for y, s in zip(years, stats) if y in (2026, 2030, 2040, 2050, 2075, 2100, 2150, 2200)})
with open(os.path.join(ROOT, "results", "ha", f"{name}_ha.json"), "w") as f:
    json.dump(meta, f, indent=1, default=float)
print(json.dumps({k: meta[k] for k in ("name", "converged", "beta_high", "seconds")}, default=float))
print("r (two-class vs HA) 2026/2050/2100:", [(float(P["r"][t + 1]), float(out["r"][t + 1])) for t in (0, 24, 74)])
for y in (2026, 2050, 2100, 2200):
    if y in meta["stats_by_year"]:
        print(y, meta["stats_by_year"][y])