"""Model A experiments: delegation threshold, aggregate model noise, efficiency benchmark and wedges, compute market,
ownership, savings-channel negative result.  Outputs: results/model_a/*.csv and results/figures/figA_*.pdf.

Run:  python3 scripts/10_model_a.py
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'src'))
import matplotlib.pyplot as plt  # noqa: E402

from age.analysis import plotstyle as ps  # noqa: E402
from age.equilibrium import lq_game as lq  # noqa: E402
from age.equilibrium import lq_map as m  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'results' / 'model_a'
FIG = ROOT / 'results' / 'figures'
OUT.mkdir(parents=True, exist_ok=True)
FIG.mkdir(parents=True, exist_ok=True)
ps.use()

p = m.EconParams()
ss = m.steady_state(p)
pf = m.portfolio_lq_params(p, ss)
lqp = pf['lqp']
summary = dict(steady_state={k: float(v) for k, v in ss.items()}, portfolio_map={k: float(v) for k, v in pf.items() if k != 'lqp'})

# ------------------------------------------------------------------ A1. savings channel (negative result)
lin = m.linearized_household(p, ss)
summary['savings_channel'] = {k: float(v) for k, v in lin.items() if np.isscalar(v)}
pd.DataFrame(dict(t=np.arange(len(lin['a_ge'])), a_ge=lin['a_ge'], a_pe=lin['a_pe'])).to_csv(OUT / 'savings_news_response.csv', index=False)

# ------------------------------------------------------------------ A2. private value, delegation share, model noise vs D, M, tau_A
rows = []
for sigA in [0.08, 0.06, 0.04, 0.02]:
    for M in [1, 2, 4, 8, 1000]:
        pp = m.EconParams(sig_A_R=sigA, M=M)
        pfx = m.portfolio_lq_params(pp, ss)
        for D in np.linspace(0, 1, 41):
            mo = lq.moments(pfx['lqp'], D)
            rows.append(dict(sig_A=sigA, M=M, D=D, PV_goods=m.private_value_goods(pp, ss, pfx['lqp'], pfx, D),
                             V_noise=mo['V_noise'], V_track=mo['V_track'], V_disp=mo['V_disp'], Abar=mo['Abar'],
                             xi=mo['xi'], aA=mo['aA'], aH=mo['aH'],
                             W_rel=lq.welfare_decomposed(pfx['lqp'], D)['W_rel'] * pfx['Uxx_goods'],
                             wedge=lq.wedge(pfx['lqp'], D)))
dfA = pd.DataFrame(rows)
dfA.to_csv(OUT / 'game_moments.csv', index=False)

# delegation equilibrium vs model quality and compute price
rows = []
for sigA in np.linspace(0.02, 0.095, 16):
    for pQ in [0.2, 0.6, 1.0]:
        for qbar in [0.002, 0.05]:
            pp = m.EconParams(sig_A_R=sigA, qbar=qbar)
            pfx = m.portfolio_lq_params(pp, ss)
            rows.append(dict(sig_A=sigA, pQ=pQ, qbar=qbar, D=m.delegation_share(pp, ss, pfx['lqp'], pfx, pQ),
                             fee=pQ * qbar + pp.pE * pp.ebar))
dfD = pd.DataFrame(rows)
dfD.to_csv(OUT / 'delegation_share.csv', index=False)

# ------------------------------------------------------------------ A3. wedges: (a) non-pecuniary, (b) hand-to-mouth wage risk
rows = []
for D in np.linspace(0, 1, 41):
    for ups in [0.0, 0.5, 1.0, 1.5]:
        pfx = m.portfolio_lq_params(p, ss)
        lq_u = lq.LQParams(**{**pfx['lqp'].__dict__, 'upsilon': ups})
        rows.append(dict(D=D, upsilon=ups, wedge_goods=lq.wedge(lq_u, D) * pfx['Uxx_goods'],
                         PV_goods=lq.private_value(lq_u, D) * pfx['Uxx_goods']))
dfW = pd.DataFrame(rows)
# (b) hand-to-mouth workers: wage w ∝ K^alpha ; log-wage noise variance = alpha^2 V_noise ; CE loss = 1/2 gamma wbar Var
rows = []
for D in np.linspace(0.01, 1, 40):
    for M in [1, 2, 4, 8]:
        pfx = m.portfolio_lq_params(p, ss, M=M)
        h = 1e-4
        dV = (lq.moments(pfx['lqp'], min(D + h, 1))['V_noise'] - lq.moments(pfx['lqp'], max(D - h, 0))['V_noise']) / (min(D + h, 1) - max(D - h, 0))
        for htm in [0.25, 0.5]:
            wedge_b = htm * 0.5 * p.gamma * ss['w'] * (p.alpha ** 2) * dV      # goods per unit of D
            rows.append(dict(D=D, M=M, htm=htm, wedge_b_goods=wedge_b, PV_goods=m.private_value_goods(p, ss, pfx['lqp'], pfx, D),
                             tax_over_fee=wedge_b / ss['fee']))
dfB = pd.DataFrame(rows)
dfW.to_csv(OUT / 'wedge_nonpecuniary.csv', index=False)
dfB.to_csv(OUT / 'wedge_hand_to_mouth.csv', index=False)

# ------------------------------------------------------------------ A4. compute market: short run vs long run
rows = []
for KQ in [0.02, 0.05, 0.1]:
    ssk = m.steady_state(p, KQ=KQ)
    for qbar in np.linspace(0.0, 0.04, 41):
        pp = m.EconParams(qbar=qbar)
        pfx = m.portfolio_lq_params(pp, ssk)
        r = m.compute_market_short_run(pp, ssk, pfx['lqp'], pfx, KQ=KQ)
        rows.append(dict(KQ=KQ, qbar=qbar, regime='short run (fixed K_Q)', **{k: float(v) for k, v in r.items() if k != 'KQ'}))
        # long run: compute supplied at user cost -> production compute unaffected by delegation; K_Q expands
        pr_lr = pp.tasks.prices(0.0, ssk['K'], KQ, ssk['L'])
        D_lr = m.delegation_share(pp, ssk, pfx['lqp'], pfx, pr_lr['pQ'])
        rows.append(dict(KQ=KQ, qbar=qbar, regime='long run (elastic K_Q)', QP=KQ, D=D_lr, fee=pr_lr['pQ'] * qbar + pp.pE * pp.ebar,
                         KQ_total=KQ + qbar * D_lr, **{k: float(v) for k, v in pr_lr.items()}))
dfC = pd.DataFrame(rows)
dfC.to_csv(OUT / 'compute_market.csv', index=False)

# ------------------------------------------------------------------ A5. ownership: two-class income shares vs automation
rows = []
for KQ in np.linspace(0.0, 2.0, 41):
    ssk = m.steady_state(p, KQ=KQ)
    compute_income = ssk['share_Q'] * ssk['Y']
    for n_own in [0.01, 0.10, 1.0]:
        owners_share = (ssk['share_K'] * n_own + ssk['share_Q'])  # owners hold compute + their pro-rata share of K
        rows.append(dict(KQ=KQ, iota=ssk['iota'], share_L=ssk['share_L'], share_Q=ssk['share_Q'], Y=ssk['Y'], w=ssk['w'],
                         n_owners=n_own, owners_income_share=owners_share))
dfO = pd.DataFrame(rows)
dfO.to_csv(OUT / 'ownership_shares.csv', index=False)

# ------------------------------------------------------------------ A6. sensitivity heat-map: model-noise variance of wages over (sig_A, M)
grid_sA = np.linspace(0.02, 0.1, 17)
grid_M = [1, 2, 3, 4, 6, 8, 12, 16]
H = np.zeros((len(grid_sA), len(grid_M)))
Dgrid = np.zeros_like(H)
for i, sA in enumerate(grid_sA):
    for j, M in enumerate(grid_M):
        pp = m.EconParams(sig_A_R=sA, M=M)
        pfx = m.portfolio_lq_params(pp, ss)
        D = m.delegation_share(pp, ss, pfx['lqp'], pfx, 0.5)
        Dgrid[i, j] = D
        H[i, j] = p.alpha ** 2 * lq.moments(pfx['lqp'], D)['V_noise']      # Var(log wage) from model noise
pd.DataFrame(H, index=grid_sA, columns=grid_M).to_csv(OUT / 'heatmap_wage_noise.csv')
pd.DataFrame(Dgrid, index=grid_sA, columns=grid_M).to_csv(OUT / 'heatmap_D.csv')

json.dump(summary, open(OUT / 'summary.json', 'w'), indent=1, default=float)

# ================================================================== figures
# figA_pv: private value vs D for several model qualities (M=1) + model diversity
fig, axs = plt.subplots(1, 2, figsize=(7.2, 3.1))
for k, sigA in enumerate([0.08, 0.06, 0.04, 0.02]):
    d = dfA[(dfA.M == 1) & (dfA.sig_A == sigA)]
    axs[0].plot(d.D, 100 * d.PV_goods / ss['C'], color=ps.CAT[k], label=f'model noise s.d. {sigA:.2f}')
axs[0].set(xlabel='delegation share $D$', ylabel='private value of delegation (% of consumption)', title='(a) Private value of delegation')
axs[0].legend()
for k, M in enumerate([1, 2, 4, 8, 1000]):
    d = dfA[(dfA.M == M) & (dfA.sig_A == 0.04)]
    axs[1].plot(d.D, 1e4 * d.V_noise * p.alpha ** 2, color=ps.CAT[k], label=f'M = {M}' if M < 1000 else 'M → ∞')
axs[1].set(xlabel='delegation share $D$', ylabel='wage variance from model noise (×10⁴)', title='(b) Aggregate model noise ∝ $D^2/M$')
axs[1].legend()
ps.save(fig, FIG / 'figA_private_value_noise.pdf')

# figA_threshold: D* vs model quality for compute prices
fig, ax = plt.subplots(figsize=(5.2, 3.2))
for k, (pQ, qbar) in enumerate([(0.2, 0.05), (0.6, 0.05), (1.0, 0.05)]):
    d = dfD[(dfD.pQ == pQ) & (dfD.qbar == qbar)]
    ax.plot(d.sig_A, d.D, color=ps.CAT[k], label=f'compute rental $p_Q$ = {pQ}')
ax.set(xlabel='foundation-model noise s.d. (lower = better model)', ylabel='equilibrium delegation share $D^*$', title='Delegation threshold (H1)')
ax.invert_xaxis()
ax.legend()
ps.save(fig, FIG / 'figA_threshold.pdf')

# figA_wedge
fig, axs = plt.subplots(1, 2, figsize=(7.2, 3.1))
for k, ups in enumerate([0.0, 0.5, 1.0, 1.5]):
    d = dfW[dfW.upsilon == ups]
    axs[0].plot(d.D, 100 * d.wedge_goods / ss['C'], color=ps.CAT[k], label=f'υ = {ups}' + (' (CRS benchmark)' if ups == 1 else ''))
axs[0].axhline(0, color=ps.AXIS, lw=0.6)
axs[0].set(xlabel='delegation share $D$', ylabel='private − social marginal value (% of consumption)', title='(a) Wedge by curvature ratio υ')
axs[0].legend()
for k, M in enumerate([1, 2, 4, 8]):
    d = dfB[(dfB.M == M) & (dfB.htm == 0.5)]
    axs[1].plot(d.D, d.tax_over_fee, color=ps.CAT[k], label=f'M = {M}')
axs[1].set(xlabel='delegation share $D$', ylabel='Pigouvian delegation tax / fee', title='(b) Wage-risk wedge, 50% hand-to-mouth')
axs[1].legend()
ps.save(fig, FIG / 'figA_wedge.pdf')

# figA_compute: short vs long run
fig, axs = plt.subplots(1, 3, figsize=(9.0, 3.0))
for k, KQ in enumerate([0.02, 0.05, 0.1]):
    for ls, reg in [('-', 'short run (fixed K_Q)'), ('--', 'long run (elastic K_Q)')]:
        d = dfC[(dfC.KQ == KQ) & (dfC.regime == reg)]
        axs[0].plot(d.qbar, d.pQ, ls, color=ps.CAT[k], label=f'$K_Q$={KQ}, {reg}' if k == 0 or ls == '-' else None)
        axs[1].plot(d.qbar, d.iota, ls, color=ps.CAT[k])
        axs[2].plot(d.qbar, d.share_L, ls, color=ps.CAT[k])
axs[0].set(xlabel='inference intensity $\\bar q$', ylabel='compute rental $p_Q$', title='(a) Compute price')
axs[1].set(xlabel='inference intensity $\\bar q$', ylabel='automated task share $\\iota$', title='(b) Automation')
axs[2].set(xlabel='inference intensity $\\bar q$', ylabel='labour share', title='(c) Labour share (H6)')
axs[0].legend(fontsize=7)
ps.save(fig, FIG / 'figA_compute_market.pdf')

# figA_ownership
fig, ax = plt.subplots(figsize=(5.2, 3.2))
for k, n_own in enumerate([0.01, 0.10, 1.0]):
    d = dfO[dfO.n_owners == n_own]
    ax.plot(d.iota, d.owners_income_share, color=ps.CAT[k], label=f'owners = top {int(n_own*100)}%' if n_own < 1 else 'broad ownership')
ax.plot(dfO[dfO.n_owners == 1.0].iota, dfO[dfO.n_owners == 1.0].share_L, color=ps.TEXT2, lw=1, ls=':', label='labour share')
ax.set(xlabel='automated task share $\\iota$', ylabel='income share', title='Compute ownership and income concentration (Prop. 6)')
ax.legend()
ps.save(fig, FIG / 'figA_ownership.pdf')

# figA_heatmap
fig, ax = plt.subplots(figsize=(5.4, 3.4))
im = ax.imshow(1e4 * H, origin='lower', aspect='auto', cmap=ps.SEQ,
               extent=[-0.5, len(grid_M) - 0.5, grid_sA[0], grid_sA[-1]])
ax.set_xticks(range(len(grid_M))); ax.set_xticklabels(grid_M)
ax.set(xlabel='number of independent foundation models $M$', ylabel='model noise s.d.', title='Wage variance from model noise at $D^*$ (×10⁴)')
ax.grid(False)
fig.colorbar(im, ax=ax, shrink=0.85)
ps.save(fig, FIG / 'figA_heatmap.pdf')

# figA_savings (negative result): GE and PE news responses
fig, ax = plt.subplots(figsize=(5.2, 3.0))
ax.plot(np.arange(12), lin['a_ge'][:12], color=ps.CAT[0], label='general equilibrium')
ax.plot(np.arange(12), lin['a_pe'][:12], color=ps.CAT[1], label='partial equilibrium (prices fixed)')
ax.axhline(0, color=ps.AXIS, lw=0.6)
ax.set(xlabel='years after the news', ylabel='log capital response to a unit TFP news', title='Savings channel: a 1-log-point news moves savings by < 0.01')
ax.legend()
ps.save(fig, FIG / 'figA_savings_news.pdf')
print('Model A done:', summary['portfolio_map'], 'savings kappa0', lin['kappa0'])