"""Figures from solved scenarios (results/scenarios/<name>_paths.npz). Skips scenarios not yet solved.
Outputs results/figures/fig03..fig09 (pdf + png)."""
import glob
import json
import os
import sys

import matplotlib
import numpy as np

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

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from agi_ge.analysis import campbell_shiller_decomposition, valuation_ratios  # noqa: E402
from agi_ge.solve import load_paths  # noqa: E402

ROOT = os.path.join(os.path.dirname(__file__), "..")
SC = os.path.join(ROOT, "results", "scenarios")
FIG = os.path.join(ROOT, "results", "figures")
os.makedirs(FIG, exist_ok=True)
plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "legend.frameon": False})
C = {"baseline": "#1f3b73", "no_ai": "#7f8c8d", "rapid_agi": "#c0392b", "fast_moore_cheap_energy": "#e67e22",
     "fast_moore_energy_constrained": "#8e44ad", "complements_leontief_like": "#16a085", "complements_cobb_douglas": "#2980b9",
     "substitutes": "#d35400", "ai_monopoly": "#2c3e50", "citizen_fund_20": "#27ae60", "land_tax_50": "#f39c12",
     "compute_tax_30": "#c0392b", "capital_tax_20": "#8e44ad", "foreign_ai_50": "#7f8c8d", "low_risk_aversion": "#e67e22",
     "high_risk_aversion": "#16a085", "scarce_land": "#8e44ad", "severe_chip_depreciation": "#2c3e50",
     "reinstatement": "#27ae60", "moore_ends_2040": "#f39c12", "moore_never_ends": "#c0392b"}
LAB = {"baseline": "baseline", "no_ai": "AI frozen (2026 capability)", "rapid_agi": "rapid AGI (2032)",
       "fast_moore_cheap_energy": "fast Moore, cheap energy", "fast_moore_energy_constrained": "fast Moore, energy-constrained",
       "complements_leontief_like": r"complements ($\sigma=0.4$)", "complements_cobb_douglas": r"Cobb-Douglas ($\sigma=1$)",
       "substitutes": r"substitutes ($\sigma=1.5$)", "ai_monopoly": "AI monopoly (markup 1.5)", "citizen_fund_20": "citizen fund 20%",
       "land_tax_50": "land tax 50%", "compute_tax_30": "compute tax 30%", "capital_tax_20": "capital-income tax 20%",
       "foreign_ai_50": "50% foreign-owned compute", "low_risk_aversion": r"$\gamma=1$", "high_risk_aversion": r"$\gamma=4$",
       "scarce_land": "scarce land", "severe_chip_depreciation": r"$\delta_Q=0.45$", "reinstatement": "task 4 never automated",
       "moore_ends_2040": "Moore's law ends 2040", "moore_never_ends": "Moore's law never ends"}


def load_all():
    out = {}
    for f in sorted(glob.glob(os.path.join(SC, "*_paths.npz"))):
        name = os.path.basename(f).replace("_paths.npz", "")
        out[name] = load_paths(f)
    return out


def save(fig, name):
    for ext in ("pdf", "png"):
        fig.savefig(os.path.join(FIG, f"{name}.{ext}"), dpi=200, bbox_inches="tight")
    plt.close(fig)


def yr_mask(P, y0=2026, y1=2100):
    return (P["years"] >= y0) & (P["years"] <= y1)


def fig_baseline(P, P0=None):
    m = yr_mask(P); y = P["years"][m]
    fig, ax = plt.subplots(3, 3, figsize=(10.5, 8))
    gf = P["gf"]
    ax[0, 0].plot(y, (P["GDP"] / P["GDP"][0])[m], color=C["baseline"], label="baseline")
    if P0 is not None:
        ax[0, 0].plot(y, (P0["GDP"] / P0["GDP"][0])[yr_mask(P0)], color=C["no_ai"], ls="--", label="AI frozen")
    ax[0, 0].set_yscale("log"); ax[0, 0].set_title("GDP (2026 = 1)"); ax[0, 0].legend()
    ax[0, 1].plot(y, P["labor_share"][m], color="#c0392b", label="labor")
    ax[0, 1].plot(y, (P["R_Q"] * P["K_Q"][:-1] / P["GDP"])[m], color=C["baseline"], label="compute")
    ax[0, 1].plot(y, (P["R_N"] * P["K_N"][:-1] / P["GDP"])[m], color="#7f8c8d", label="structures")
    ax[0, 1].plot(y, ((P["RTp"] + P["RTHp"]) / P["GDP"])[m], color="#27ae60", label="land")
    ax[0, 1].plot(y, (P["PE"] * P["E"] / P["GDP"])[m], color="#e67e22", label="energy")
    ax[0, 1].set_title("Income shares of GDP"); ax[0, 1].legend(fontsize=7)
    ax[0, 2].plot(y, (P["wL"] * gf / (P["wL"][0]))[m], color="#c0392b", label="$w_L$ per worker")
    ax[0, 2].plot(y, (P["wH"] * gf / (P["wH"][0]))[m], color=C["baseline"], label="$w_H$ per worker")
    ax[0, 2].plot(y, gf[m], color="#7f8c8d", ls=":", label="trend $(1+g_Z)^t$")
    ax[0, 2].set_title("Real wages (2026 = 1)"); ax[0, 2].legend(fontsize=7)
    ax[1, 0].plot(y, 100 * P["r"][1:][m], color=C["baseline"])
    ax[1, 0].set_title("Real return on wealth $r_{t+1}$ (%)")
    ax[1, 1].plot(y, P["K_Q"][:-1][m], color=C["baseline"], label="$K^Q$ (compute units)")
    ax[1, 1].set_yscale("log"); ax[1, 1].set_title("Compute capacity (log)")
    ax2 = ax[1, 1].twinx(); ax2.plot(y, P["A"][m], color="#c0392b", ls="--"); ax2.set_yscale("log"); ax2.set_ylabel("price of compute $A$ (dashed)", color="#c0392b")
    ax[1, 2].plot(y, (P["EQ"] / P["E"])[m], color="#e67e22", label="share of energy to compute")
    ax[1, 2].plot(y, P["util"][m], color="#7f8c8d", ls="--", label="utilization")
    ax[1, 2].plot(y, (P["PE"] / P["PE"][0])[m], color=C["baseline"], ls=":", label="energy price (2026=1)")
    ax[1, 2].set_title("Energy"); ax[1, 2].legend(fontsize=7)
    Li = P["Li"]; Qi = P["Qi"]; al = P["alpha"]
    comp_share_task = (Qi / al) / (Li + Qi / al + 1e-300)
    for i, nm in enumerate(["routine", "cognitive-L", "cognitive-H", "expert"]):
        ax[2, 0].plot(y, comp_share_task[m, i], label=nm)
    ax[2, 0].set_title("Compute share of task input"); ax[2, 0].legend(fontsize=7)
    ax[2, 1].plot(y, (P["inv_goods"] / P["GDP"])[m], color=C["baseline"], label="total investment / GDP")
    ax[2, 1].plot(y, ((P["pQ"] * P["I_Q"] + P["adj_Q"]) / P["GDP"])[m], color="#c0392b", label="compute investment / GDP")
    ax[2, 1].set_title("Investment"); ax[2, 1].legend(fontsize=7)
    ax[2, 2].plot(y, (P["C_W"] / P["C_W"][0] / gf)[m], color="#c0392b", label="workers (detrended)")
    ax[2, 2].plot(y, (P["C_K"] / P["C_K"][0] / gf)[m], color=C["baseline"], label="capitalists (detrended)")
    ax[2, 2].set_yscale("log"); ax[2, 2].set_title("Consumption by class (2026 = 1)"); ax[2, 2].legend(fontsize=7)
    for a in ax.ravel():
        a.set_xlim(2026, 2100)
    fig.tight_layout()
    save(fig, "fig03_baseline")


def fig_valuation(P, P0=None):
    m = yr_mask(P); y = P["years"][m]
    v = valuation_ratios(P)
    fig, ax = plt.subplots(2, 3, figsize=(10.5, 5.6))
    for k, lab, col in [("N", "structures", "#7f8c8d"), ("Q", "compute", C["baseline"]), ("E", "energy capital", "#e67e22"), ("H", "housing structures", "#16a085")]:
        ax[0, 0].plot(y, v[f"V_over_GDP_{k}"][m], color=col, label=lab)
    ax[0, 0].plot(y, ((P["PT"] + P["PTH"]) / P["GDP"])[m], color="#27ae60", label="land (both)")
    ax[0, 0].set_title("Market value of claims / GDP"); ax[0, 0].legend(fontsize=7)
    ax[0, 1].plot(y, v["W_over_GDP"][m], color=C["baseline"], label="baseline")
    if P0 is not None:
        ax[0, 1].plot(y, valuation_ratios(P0)["W_over_GDP"][yr_mask(P0)], color=C["no_ai"], ls="--", label="AI frozen")
    ax[0, 1].set_title("Aggregate wealth / GDP"); ax[0, 1].legend(fontsize=7)
    ax[0, 2].plot(y, v["PE_Q"][m], color=C["baseline"], label="compute")
    ax[0, 2].plot(y, v["PE_N"][m], color="#7f8c8d", label="structures")
    ax[0, 2].set_title("Price-earnings ratios"); ax[0, 2].legend(fontsize=7)
    ax[1, 0].plot(y, v["land_PR"][m], color="#27ae60", label="production land")
    ax[1, 0].plot(y, v["house_PR"][m], color="#16a085", label="house (structure + land)")
    ax[1, 0].set_title("Price-rent ratios"); ax[1, 0].legend(fontsize=7)
    ax[1, 1].plot(y, v["q_over_p_N"][m], color="#7f8c8d", label="structures")
    ax[1, 1].plot(y, v["q_over_p_Q"][m], color=C["baseline"], label="compute")
    ax[1, 1].plot(y, v["q_over_p_H"][m], color="#16a085", label="housing")
    ax[1, 1].set_title("Tobin's $q$ (market / replacement cost)"); ax[1, 1].legend(fontsize=7)
    ax[1, 2].plot(y, v["house_land_share"][m], color="#27ae60", label="land share of house value")
    ax[1, 2].plot(y, v["wealth_comp"]["land"][m], color="#7f8c8d", ls="--", label="land share of wealth")
    ax[1, 2].plot(y, v["wealth_comp"]["Q"][m], color=C["baseline"], ls=":", label="compute share of wealth")
    ax[1, 2].set_title("Composition"); ax[1, 2].legend(fontsize=7)
    for a in ax.ravel():
        a.set_xlim(2026, 2100)
    fig.tight_layout()
    save(fig, "fig04_valuation")


def fig_scenarios(allP, names, fname, title):
    names = [n for n in names if n in allP]
    if len(names) < 2:
        return
    fig, ax = plt.subplots(2, 3, figsize=(10.5, 5.6))
    for n in names:
        P = allP[n]; m = yr_mask(P); y = P["years"][m]; v = valuation_ratios(P); gf = P["gf"]
        kw = dict(color=C.get(n, "k"), label=LAB.get(n, n), lw=1.3 if n == "baseline" else 1.0)
        ax[0, 0].plot(y, (P["GDP"] / P["GDP"][0] / gf)[m], **kw)
        ax[0, 1].plot(y, P["labor_share"][m], **kw)
        ax[0, 2].plot(y, (P["wL"] * gf / P["wL"][0])[m], **kw)
        ax[1, 0].plot(y, 100 * P["r"][1:][m], **kw)
        ax[1, 1].plot(y, v["W_over_GDP"][m], **kw)
        ax[1, 2].plot(y, v["land_PR"][m], **kw)
    ax[0, 0].set_title("GDP / trend (2026 = 1)"); ax[0, 0].set_yscale("log")
    ax[0, 1].set_title("Labor share")
    ax[0, 2].set_title("Wage of skill L (2026 = 1)"); ax[0, 2].set_yscale("log")
    ax[1, 0].set_title("Real return (%)")
    ax[1, 1].set_title("Wealth / GDP")
    ax[1, 2].set_title("Land price-rent ratio")
    ax[0, 0].legend(fontsize=7)
    for a in ax.ravel():
        a.set_xlim(2026, 2100)
    fig.suptitle(title, fontsize=10)
    fig.tight_layout()
    save(fig, fname)


def fig_decomposition(allP):
    if "baseline" not in allP or "no_ai" not in allP:
        return
    P1, P0 = allP["baseline"], allP["no_ai"]
    claims = [("land", "production land"), ("housing_land", "residential land"), ("N", "structures"), ("Q", "compute"), ("H", "housing structures")]
    rows = []
    for c, lab in claims:
        try:
            d = campbell_shiller_decomposition(P1, P0, claim=c)
            rows.append((lab, d["cash_flow"], d["discount_rate"], d["total"]))
        except Exception as e:  # noqa: BLE001
            print("decomp failed", c, e)
    if not rows:
        return
    fig, ax = plt.subplots(figsize=(6.5, 3.2))
    x = np.arange(len(rows)); w = 0.35
    ax.bar(x - w / 2, [r[1] for r in rows], w, color="#27ae60", label="cash-flow channel")
    ax.bar(x + w / 2, [r[2] for r in rows], w, color="#c0392b", label="discount-rate channel")
    ax.plot(x, [r[3] for r in rows], "ko", label="total log change in value (2026)")
    ax.axhline(0, color="k", lw=0.5)
    ax.set_xticks(x); ax.set_xticklabels([r[0] for r in rows], fontsize=8)
    ax.set_ylabel("log(value, baseline / AI frozen)")
    ax.legend(fontsize=7)
    fig.tight_layout()
    save(fig, "fig07_decomposition")
    with open(os.path.join(ROOT, "results", "data", "decomposition_2026.json"), "w") as f:
        json.dump([dict(claim=r[0], cash_flow=r[1], discount_rate=r[2], total=r[3]) for r in rows], f, indent=1)


if __name__ == "__main__":
    allP = load_all()
    print("scenarios available:", list(allP))
    if "baseline" in allP:
        fig_baseline(allP["baseline"], allP.get("no_ai"))
        fig_valuation(allP["baseline"], allP.get("no_ai"))
    fig_scenarios(allP, ["baseline", "no_ai", "rapid_agi", "fast_moore_cheap_energy", "fast_moore_energy_constrained", "moore_never_ends"],
                  "fig05_technology_scenarios", "Technology scenarios")
    fig_scenarios(allP, ["baseline", "complements_leontief_like", "complements_cobb_douglas", "substitutes", "scarce_land", "severe_chip_depreciation"],
                  "fig06_structure_scenarios", "Substitutability and scarcity")
    fig_decomposition(allP)
    fig_scenarios(allP, ["baseline", "citizen_fund_20", "land_tax_50", "compute_tax_30", "capital_tax_20", "ai_monopoly", "foreign_ai_50"],
                  "fig08_policy_scenarios", "Ownership, competition and policy")
    fig_scenarios(allP, ["baseline", "low_risk_aversion", "high_risk_aversion", "reinstatement", "moore_ends_2040"],
                  "fig09_preference_scenarios", "Preferences, reinstatement and the end of Moore's law")
    print("done")