"""Illustration of Propositions 2-5 (limit economy with one scarce complement): price of compute, wages and
factor shares as the compute endowment grows, for sigma in {0.5, 1, 1.5}; plus the energy-bounded case.
Output: results/figures/fig02_regimes.{pdf,png}, results/data/regime_map.csv"""
import os
import sys

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from agi_ge.ces import rho_of_sigma  # noqa: E402
from agi_ge.tasks import Technology  # noqa: E402

ROOT = os.path.join(os.path.dirname(__file__), "..")
NU = [0.3, 0.25, 0.25, 0.2]
SKILL = [0, 0, 1, 1]
ALPHA = np.array([1.0, 2.0, 4.0, 8.0])
H = [0.6, 0.4]
S = 1.0          # the scarce complement (enters as K_N with nu_N = 0.4)
Qg = np.logspace(-1, 6, 100)
rows = []
for sigma in (0.5, 1.0, 1.5):
    tech = Technology(nu=np.array(NU), gamma=np.zeros(4), skill=np.array(SKILL), rho_B=-1.0,
                      nu_N=0.4, rho_N=rho_of_sigma(sigma), nu_T=0.0, nu_E=0.0, rho_T=0.0)
    x0 = None
    for Q in Qg:
        s = tech.solve(Q, H, KN=S, T=0.0, EY=0.0, Z=1.0, alpha=ALPHA, x0=x0)
        x0 = s["x"]
        wb = float(np.sum(s["w"] * np.array(H)))
        rows.append(dict(sigma=sigma, Q=Q, Y=s["Y"], A=s["A"], wL=s["w"][0], wH=s["w"][1],
                         labor_share=wb / s["Y"], compute_share=s["A"] * Q / s["Y"], S_share=s["RN"] * S / s["Y"]))
df = pd.DataFrame(rows)
df.to_csv(os.path.join(ROOT, "results", "data", "regime_map.csv"), index=False)

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False})
cols = {0.5: "#c0392b", 1.0: "#7f8c8d", 1.5: "#1f3b73"}
fig, ax = plt.subplots(1, 4, figsize=(10.5, 2.9))
for sigma, d in df.groupby("sigma"):
    lab = rf"$\sigma={sigma}$"
    ax[0].plot(d.Q, d.Y, color=cols[sigma], label=lab)
    ax[1].plot(d.Q, d.A, color=cols[sigma], label=lab)
    ax[2].plot(d.Q, d.wH, color=cols[sigma], label=lab)
    ax[3].plot(d.Q, d.compute_share, color=cols[sigma], label=lab)
    ax[3].plot(d.Q, d.S_share, color=cols[sigma], ls="--")
ax[0].set_yscale("log"); ax[0].set_title("Output $Y$")
ax[1].set_yscale("log"); ax[1].set_title("Price of compute $A$")
ax[2].set_yscale("log"); ax[2].set_title("Wage $w_H = A\\cdot\\mathrm{CEU}(H)$")
ax[3].set_title("Shares: compute (solid), complement $S$ (dashed)")
for a in ax:
    a.set_xscale("log"); a.set_xlabel("compute $Q$ (complement $S$ fixed)")
ax[0].legend(frameon=False, fontsize=8)
fig.tight_layout()
for ext in ("pdf", "png"):
    fig.savefig(os.path.join(ROOT, "results", "figures", f"fig02_regimes.{ext}"), dpi=200)
print(df.groupby("sigma").tail(1)[["sigma", "Q", "Y", "A", "wH", "compute_share", "S_share", "labor_share"]].to_string())