"""Figures and numbers from the heterogeneous-agent runs (results/ha/*_ha.npz).
Outputs: results/figures/fig10_distribution.{pdf,png}, results/tables/tab07_distribution.tex, results/data/ha_numbers.json"""
import glob
import json
import os

import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = os.path.join(os.path.dirname(__file__), "..")
HA = os.path.join(ROOT, "results", "ha")
FIG = os.path.join(ROOT, "results", "figures")
TAB = os.path.join(ROOT, "results", "tables")
plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "legend.frameon": False})
C = {"baseline": "#1f3b73", "no_ai": "#7f8c8d", "citizen_fund_20": "#27ae60", "rapid_agi": "#c0392b"}
LAB = {"baseline": "baseline", "no_ai": "AI frozen", "citizen_fund_20": "citizen fund 20%", "rapid_agi": "rapid AGI"}

runs = {}
for f in sorted(glob.glob(os.path.join(HA, "*_ha.npz"))):
    n = os.path.basename(f).replace("_ha.npz", "")
    runs[n] = dict(np.load(f))
    runs[n]["meta"] = json.load(open(f.replace(".npz", ".json")))
print("HA runs:", list(runs))
if runs:
    fig, ax = plt.subplots(2, 3, figsize=(10.5, 5.6))
    for n, d in runs.items():
        y = d["years"]; m = y <= 2100
        kw = dict(color=C.get(n, "k"), label=LAB.get(n, n))
        ax[0, 0].plot(y[m], d["gini"][m], **kw)
        ax[0, 1].plot(y[m], d["top10"][m], **kw)
        ax[0, 2].plot(y[m], d["top1"][m], **kw)
        ax[1, 0].plot(y[m], d["bottom50"][m], **kw)
        ax[1, 1].plot(y[m], 100 * d["r"][1:len(y) + 1][m], **kw)
        ax[1, 1].plot(y[m], 100 * d["r_twoclass"][1:len(y) + 1][m], color=C.get(n, "k"), ls=":", lw=0.8)
        ax[1, 2].plot(y[m], (d["X"] / d["GDP"][:len(y)])[m], **kw)
    ax[0, 0].set_title("Wealth Gini"); ax[0, 1].set_title("Top-10% wealth share"); ax[0, 2].set_title("Top-1% wealth share")
    ax[1, 0].set_title("Bottom-50% wealth share"); ax[1, 1].set_title("Real return: HA (solid) vs two-class (dotted), %")
    ax[1, 2].set_title("Household expenditure / GDP")
    ax[0, 0].legend(fontsize=7)
    for a in ax.ravel():
        a.set_xlim(2026, 2100)
    fig.tight_layout()
    for ext in ("pdf", "png"):
        fig.savefig(os.path.join(FIG, f"fig10_distribution.{ext}"), dpi=200, bbox_inches="tight")
    years = (2026, 2030, 2040, 2050, 2075, 2100)
    with open(os.path.join(TAB, "tab07_distribution.tex"), "w") as f:
        f.write("\\begin{tabular}{@{}ll" + "r" * len(years) + "@{}}\n\\toprule\nScenario & Statistic & " + " & ".join(str(y) for y in years) + "\\\\\n\\midrule\n")
        for n, d in runs.items():
            for stat, lab in [("gini", "Gini"), ("top10", "Top 10\\%"), ("top1", "Top 1\\%"), ("bottom50", "Bottom 50\\%")]:
                vals = []
                for yy in years:
                    i = np.where(d["years"] == yy)[0]
                    vals.append(f"{d[stat][i[0]]:.3f}" if len(i) else "--")
                f.write(f"{LAB.get(n, n) if stat == 'gini' else ''} & {lab} & " + " & ".join(vals) + "\\\\\n")
            vals = []
            for yy in years:
                i = np.where(d["years"] == yy)[0]
                vals.append(f"{100 * d['r'][i[0] + 1]:.1f}" if len(i) else "--")
            f.write(f" & $r$ (\\%) & " + " & ".join(vals) + "\\\\\n\\midrule\n")
        f.write("\\bottomrule\n\\end{tabular}\n")
    nums = {}
    if "baseline" in runs:
        d = runs["baseline"]; i0 = 0
        nums["GiniBase"] = f"{d['gini'][0]:.2f}"; nums["TopTenBase"] = f"{d['top10'][0]:.2f}"; nums["TopOneBase"] = f"{d['top1'][0]:.2f}"
        nums["BottomFiftyBase"] = f"{d['bottom50'][0]:.3f}"; nums["betaHigh"] = f"{d['meta']['beta_high']:.4f}"
        for yy in (2050, 2100):
            i = int(np.where(d["years"] == yy)[0][0])
            yw = {2050: "TwoZeroFiveZero", 2100: "TwoOneZeroZero"}[yy]
            nums[f"Gini{yw}"] = f"{d['gini'][i]:.2f}"
            nums[f"TopTen{yw}"] = f"{d['top10'][i]:.2f}"
            nums[f"TopOne{yw}"] = f"{d['top1'][i]:.2f}"
            nums[f"rHA{yw}"] = f"{100 * d['r'][i + 1]:.1f}"
            nums[f"rTC{yw}"] = f"{100 * d['r_twoclass'][i + 1]:.1f}"
        nums["HAconverged"] = "yes" if d["meta"]["converged"] else "no"
        nums["HAmaxGap"] = f"{d['meta']['hist'][-1]:.0e}"
    with open(os.path.join(ROOT, "results", "data", "ha_numbers.json"), "w") as f:
        json.dump(nums, f, indent=1)
    print(nums)