"""
Publication figures for Paper 6 (IJGT).
Three figures, each carrying one pillar of the result. All data computed from the
same model primitives as the replication package. Grayscale-legible, vector PDF.
"""
import numpy as np
from scipy.optimize import brentq, minimize_scalar
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator

# ---- model primitives (identical to replication) ----
ALPHA, BETA, C, DMIN = 1.0, 0.8, 3.0, 0.15
gL, gH = 0.20, 0.95
def phi(x): return (1 - x) ** 2
def y_star(x, d):
    p = phi(x); A = ALPHA * x + BETA
    return p * A / (2 * BETA * p + d)
def K_star(x, d):
    p = phi(x); A = ALPHA * x + BETA
    return A * (BETA * p + d) / (2 * BETA * p + d)
def dK_dx(x, d):
    p = phi(x); A = ALPHA * x + BETA; Q = 2 * BETA * p + d
    R = (BETA * p + d) / Q
    return ALPHA * R + A * (2 * BETA * d * (1 - x) / Q**2)
def dev_FOC(x, g, d):
    K = K_star(x, d)
    return g * K * dK_dx(x, d) - C * x
def solve_dev(g, d):
    """Fast figure solver. The replication file performs the all-roots global check;
    this plotting routine uses the stable first-root branch for speed."""
    try:
        return brentq(lambda x: dev_FOC(x, g, d), 1e-9, 1 - 1e-9)
    except Exception:
        return 1e-6
def W(g, d):
    x = solve_dev(g, d); return K_star(x, d) ** 2 / 2 - C / 2 * x ** 2 - d / 2 * y_star(x, d) ** 2
def VD(g, d):
    x = solve_dev(g, d); return g / 2 * K_star(x, d) ** 2 - C / 2 * x ** 2
def wedge(g, dbar): return VD(g, dbar) - VD(g, DMIN)
_ARCH_CACHE = {}
_SCREEN_CACHE = {}
def arch_obj(dbar, N=220):
    key=(round(float(dbar),6),N)
    if key not in _ARCH_CACHE:
        gs = np.linspace(gL, gH, N)
        _ARCH_CACHE[key] = np.trapezoid([W(g, dbar) for g in gs], gs) / (gH - gL)
    return _ARCH_CACHE[key]
def screen_obj(dbar, lam, tilde, N=120):
    slo = wedge(tilde, dbar)
    glo = np.linspace(gL, tilde, N); ghi = np.linspace(tilde, gH, N)
    lo = np.trapezoid([W(g, DMIN) - (1 + lam) * slo for g in glo], glo)
    hi = np.trapezoid([W(g, dbar) for g in ghi], ghi)
    return (lo + hi) / (gH - gL)
def best_screen(dbar, lam):
    # The paper proves the screening objective is decreasing in the threshold.
    # Use the first non-empty grid point above the endpoint; the endpoint itself is architecture-only.
    key=(round(float(dbar),6),round(float(lam),6))
    if key not in _SCREEN_CACHE:
        t = gL + 1e-3
        _SCREEN_CACHE[key] = (screen_obj(dbar, lam, t, N=160), t)
    return _SCREEN_CACHE[key]

# ---- editorial style ----
mpl.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['DejaVu Serif'],
    'mathtext.fontset': 'dejavuserif',
    'font.size': 10.5,
    'axes.linewidth': 0.8,
    'axes.edgecolor': '#222222',
    'axes.labelcolor': '#111111',
    'xtick.color': '#222222', 'ytick.color': '#222222',
    'xtick.direction': 'out', 'ytick.direction': 'out',
    'axes.grid': False,
    'legend.frameon': False,
    'figure.dpi': 150,
})
INK = '#1a1a1a'
ACCENT = '#8c1515'     # a restrained oxblood — distinct from the AI-default terracotta
STEEL = '#31688e'      # cool counterpoint
SHADE = '#c9c9c9'
GRIDC = '#e6e6e6'

def style_axes(ax):
    for s in ['top', 'right']:
        ax.spines[s].set_visible(False)
    ax.spines['left'].set_color(INK); ax.spines['bottom'].set_color(INK)
    ax.tick_params(length=3.5, width=0.8)

# =====================================================================
# FIGURE 1 — Welfare monotonicity and the necessity of d_min > 0
# =====================================================================
def figure1():
    fig, (axL, axR) = plt.subplots(1, 2, figsize=(9.2, 3.9), gridspec_kw={'width_ratios': [1.35, 1]})

    # LEFT: W(gamma; d) vs d for several types, dmin shaded
    ds = np.linspace(0.02, 6.0, 240)
    types = [0.20, 0.50, 0.80, 0.95]
    shades = [STEEL, '#5a7fa0', ACCENT, INK]
    for g, col in zip(types, shades):
        Wd = [W(g, d) for d in ds]
        axL.plot(ds, Wd, color=col, lw=1.9, label=fr'$\gamma={g:.2f}$')
    axL.axvspan(0.0, DMIN, color=SHADE, alpha=0.5, lw=0)
    axL.axvline(DMIN, color=ACCENT, lw=1.0, ls=(0, (4, 2)))
    axL.annotate(r'$d_{\min}=0.15$', xy=(DMIN, 0.02),
                 xytext=(DMIN + 0.45, 0.020), fontsize=9.5, color=ACCENT)
    axL.annotate('non-monotone\nregion below the floor', xy=(0.09, 0.14),
                 xytext=(0.85, 0.145), fontsize=8.3, color='#555555',
                 arrowprops=dict(arrowstyle='->', color='#999999', lw=0.7))
    axL.set_xlabel(r'architecture instrument $d$')
    axL.set_ylabel(r'equilibrium welfare $W(\gamma;d)$')
    axL.set_xlim(0, 6); axL.set_ylim(0, None)
    axL.legend(loc='lower right', fontsize=9, handlelength=1.6, labelspacing=0.3)
    axL.xaxis.set_major_locator(MultipleLocator(1))
    style_axes(axL)
    axL.set_title('(a)  welfare rises in $d$ above the floor', loc='left', fontsize=10.5, color=INK, pad=8)

    # RIGHT: total dW/dd near zero, showing sign flip below dmin
    ds2 = np.linspace(0.02, 0.55, 200)
    for g, col in zip([0.20, 0.50, 0.95], [STEEL, '#5a7fa0', INK]):
        der = []
        for d in ds2:
            e = 0.01
            der.append((W(g, d + e) - W(g, d - e)) / (2 * e))
        axR.plot(ds2, der, color=col, lw=1.9, label=fr'$\gamma={g:.2f}$')
    axR.axhline(0, color='#888888', lw=0.8)
    axR.axvspan(0.0, DMIN, color=SHADE, alpha=0.5, lw=0)
    axR.axvline(DMIN, color=ACCENT, lw=1.0, ls=(0, (4, 2)))
    axR.set_xlabel(r'architecture instrument $d$')
    axR.set_ylabel(r'total derivative $dW/dd$')
    axR.set_xlim(0, 0.55)
    axR.legend(loc='lower right', fontsize=9, handlelength=1.6, labelspacing=0.3)
    style_axes(axR)
    axR.set_title(r'(b)  $dW/dd<0$ below $d_{\min}$', loc='left', fontsize=10.5, color=INK, pad=8)

    fig.tight_layout(w_pad=2.0)
    fig.savefig('fig1_welfare_monotonicity.pdf', bbox_inches='tight')
    fig.savefig('fig1_welfare_monotonicity.png', bbox_inches='tight', dpi=200)
    plt.close(fig)
    print("fig1 done")

# =====================================================================
# FIGURE 2 — Why screening collapses: the forced-subsidy deadweight
# =====================================================================
def figure2():
    fig, (axL, axR) = plt.subplots(1, 2, figsize=(9.2, 3.9))
    dbar = 2.0

    # LEFT: utility wedge increasing in gamma => low types must be paid
    gs = np.linspace(gL, gH, 200)
    wv = [wedge(g, dbar) for g in gs]
    axL.plot(gs, wv, color=ACCENT, lw=2.1)
    axL.fill_between(gs, 0, wv, color=ACCENT, alpha=0.08, lw=0)
    # annotate a representative split
    tstar = 0.55
    axL.axvline(tstar, color=INK, lw=0.9, ls=(0, (3, 2)))
    ws = wedge(tstar, dbar)
    axL.plot([tstar], [ws], 'o', color=INK, ms=4.5)
    axL.annotate(r'$s_{\mathrm{lo}}=\Delta_V(\tilde\gamma)$' + '\nforced subsidy\nto low types',
                 xy=(tstar, ws), xytext=(0.24, 0.15), fontsize=8.6, color=INK,
                 arrowprops=dict(arrowstyle='->', color=INK, lw=0.8))
    axL.set_xlabel(r'developer type $\gamma$')
    axL.set_ylabel(r'utility wedge $\Delta_V(\gamma)=V_D(\gamma;\bar d)-V_D(\gamma;d_{\min})$')
    axL.set_xlim(gL, gH); axL.set_ylim(0, None)
    style_axes(axL)
    axL.set_title(r'(a)  the wedge rises in $\gamma$', loc='left', fontsize=10.5, color=INK, pad=8)

    # RIGHT: designer objective Omega(tilde) - strictly decreasing => optimum at gamma_L
    ts = np.linspace(gL + 0.005, gH - 0.005, 120)
    for lam, col, ls in [(0.3, STEEL, '-'), (1.0, ACCENT, '-'), (2.0, INK, '-')]:
        om = [screen_obj(dbar, lam, t) for t in ts]
        axR.plot(ts, om, color=col, lw=1.9, ls=ls, label=fr'$\lambda={lam:.1f}$')
    a = arch_obj(dbar)
    axR.axhline(a, color='#888888', lw=0.9, ls=(0, (5, 2)))
    ymin = min([screen_obj(dbar, 2.0, t) for t in ts]) - 0.03
    axR.set_ylim(ymin, a + 0.045)
    axR.text(0.62, a + 0.014, 'architecture-only', fontsize=8.4, color='#444444', ha='left')
    # mark that the analytical optimum is the left endpoint (architecture-only)
    axR.plot([gL], [a], 'o', color=INK, ms=5, zorder=5, clip_on=False)
    axR.set_xlabel(r'screening threshold $\tilde\gamma$')
    axR.set_ylabel(r"designer objective $\Omega(\tilde\gamma)=\mathbb{E}[W-(1+\lambda)s]$")
    axR.set_xlim(gL, gH)
    axR.legend(loc='lower left', fontsize=9, handlelength=1.6, labelspacing=0.3)
    style_axes(axR)
    axR.set_title(r'(b)  screening lowers the objective', loc='left', fontsize=10.5, color=INK, pad=8)

    fig.tight_layout(w_pad=2.4)
    fig.savefig('fig2_screening_collapse.pdf', bbox_inches='tight')
    fig.savefig('fig2_screening_collapse.png', bbox_inches='tight', dpi=200)
    plt.close(fig)
    print("fig2 done")

# =====================================================================
# FIGURE 3 — Non-empty-screen dominance margin over (dbar, lambda)
# =====================================================================
def figure3():
    fig, (axL, axR) = plt.subplots(1, 2, figsize=(9.2, 3.9), gridspec_kw={'width_ratios': [1, 1.15]})

    dbars = np.array([1, 1.5, 2, 3, 5, 8, 12, 20], dtype=float)
    lams = [0.3, 0.5, 1.0, 2.0]
    cols = [STEEL, '#5a7fa0', ACCENT, INK]

    # LEFT: margin vs dbar for each lambda
    for lam, col in zip(lams, cols):
        marg = []
        for db in dbars:
            a = arch_obj(db); s, _ = best_screen(db, lam)
            marg.append((a - s) * 1e4)  # in units of 1e-4
        axL.plot(dbars, marg, color=col, lw=1.9, marker='o', ms=3.5, label=fr'$\lambda={lam:.1f}$')
    axL.set_xlabel(r'instrument bound $\bar d$')
    axL.set_ylabel(r'non-empty-screen loss $\times 10^{4}$')
    axL.set_xscale('log')
    axL.set_xticks([1, 2, 5, 10, 20]); axL.set_xticklabels(['1', '2', '5', '10', '20'])
    axL.axhline(0, color='#888888', lw=0.8)
    axL.legend(loc='upper left', fontsize=9, handlelength=1.6, labelspacing=0.3)
    style_axes(axL)
    axL.set_title('(a)  non-empty-screen loss is positive', loc='left', fontsize=10.5, color=INK, pad=8)

    # RIGHT: heatmap of margin over (dbar, lambda) grid
    dbg = np.linspace(1, 20, 12)
    lmg = np.linspace(0.1, 2.5, 12)
    M = np.zeros((len(lmg), len(dbg)))
    for i, lam in enumerate(lmg):
        for j, db in enumerate(dbg):
            a = arch_obj(db, N=160); s, _ = best_screen(db, lam)
            M[i, j] = (a - s) * 1e4
    im = axR.pcolormesh(dbg, lmg, M, shading='gouraud', cmap='RdGy_r')
    cs = axR.contour(dbg, lmg, M, levels=6, colors='white', linewidths=0.6, alpha=0.7)
    axR.clabel(cs, inline=True, fontsize=7, fmt='%.0f')
    cb = fig.colorbar(im, ax=axR, pad=0.02)
    cb.set_label(r'non-empty-screen loss $\times 10^{4}$', fontsize=9)
    cb.ax.tick_params(labelsize=8)
    axR.set_xlabel(r'instrument bound $\bar d$')
    axR.set_ylabel(r'transfer cost $\lambda$')
    axR.set_title('(b)  architecture wins against separation', loc='left', fontsize=10.5, color=INK, pad=8)
    for s in ['top', 'right']:
        axR.spines[s].set_visible(False)

    fig.tight_layout(w_pad=2.2)
    fig.savefig('fig3_dominance_margin.pdf', bbox_inches='tight')
    fig.savefig('fig3_dominance_margin.png', bbox_inches='tight', dpi=200)
    plt.close(fig)
    print("fig3 done")

if __name__ == '__main__':
    figure1()
    figure2()
    figure3()
    print("all figures generated")
