"""04_figures.py -- publication figures for the AI-disclosure / returns study."""
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl

mpl.rcParams.update({"font.size":11,"axes.spines.top":False,"axes.spines.right":False,
                     "figure.dpi":140,"savefig.dpi":200,"axes.grid":True,"grid.alpha":.25,
                     "font.family":"DejaVu Sans"})
D="data/processed"; F="figures"
df=pd.read_parquet(f"{D}/analysis_panel.parquet")
RES=json.load(open("results/tables/results.json"))
NAVY="#1b3a6b"; RED="#b22222"; GREY="#8a8a8a"; GREEN="#2e7d32"

# ---- Fig 1: forward return by AI quintile, full / pre / post ----
def qmeans(d, sig):
    out={}
    for q in range(1,6):
        out[q]=[]
    for yr,g in d.groupby("fiscal_year"):
        g=g.dropna(subset=[sig,"fwd_ret"])
        if len(g)<25: continue
        g=g.copy(); g["q"]=pd.qcut(g[sig],5,labels=False,duplicates="drop")+1
        for q in range(1,6): out[q].append(g[g.q==q].fwd_ret.mean())
    return [np.mean(out[q]) for q in range(1,6)]

fig,axes=plt.subplots(1,3,figsize=(13,4.2),sharey=True)
panels=[("Full sample (2010-2024)",df),("Pre-ChatGPT (2010-2021)",df[df.fiscal_year<2022]),
        ("Post-ChatGPT (2022-2024)",df[df.fiscal_year>=2022])]
for ax,(title,d) in zip(axes,panels):
    v=qmeans(d,"z_AI_COMPOS")
    bars=ax.bar(range(1,6),[x*100 for x in v],color=[GREY,GREY,GREY,GREY,NAVY],edgecolor="k",lw=.5)
    ax.set_title(title,fontsize=11); ax.set_xlabel("AI-disclosure quintile (1=low, 5=high)")
    ax.set_xticks(range(1,6))
    sp=(v[4]-v[0])*100
    ax.annotate(f"Q5−Q1 = {sp:+.1f} pp",xy=(.5,.93),xycoords="axes fraction",ha="center",
                fontsize=10,fontweight="bold",color=RED if sp>3 else "k")
axes[0].set_ylabel("Mean forward 12-month return (%)")
fig.suptitle("AI-Disclosure Intensity and Subsequent S&P 500 Returns",fontsize=13,fontweight="bold",y=1.02)
fig.tight_layout(); fig.savefig(f"{F}/fig1_quintile_returns.png",bbox_inches="tight"); plt.close()

# ---- Fig 2: cumulative long-short (AI composite) ----
fig,ax=plt.subplots(figsize=(11,5))
for k,col,lab in [("AI_COMPOS",NAVY,"AI opportunity (adoption+infra)"),("AI_RISK",GREEN,"AI disruption-risk salience")]:
    ts=pd.read_parquet(f"{D}/ls_ts_{k}.parquet").sort_values("date")
    ts["cum"]=(1+ts["ls"]).cumprod()
    ax.plot(ts["date"],ts["cum"],color=col,lw=1.8,label=lab)
ax.axhline(1,color="k",lw=.7)
ax.axvline(pd.Timestamp("2022-11-30"),color=RED,ls="--",lw=1.3)
ax.annotate("ChatGPT\n(Nov 2022)",xy=(pd.Timestamp("2022-11-30"),ax.get_ylim()[1]*.6),
            color=RED,fontsize=9,ha="right")
ax.set_title("Cumulative Value of $1 in the High-minus-Low AI Disclosure Portfolio\n(equal-weight, annually rebalanced, self-financing)",fontweight="bold")
ax.set_ylabel("Growth of $1 (long-short)"); ax.legend(frameon=False,loc="upper left")
fig.tight_layout(); fig.savefig(f"{F}/fig2_cumulative_ls.png",bbox_inches="tight"); plt.close()

# ---- Fig 3: annual Q5-Q1 spread bars ----
def annual_spread(sig):
    rec=[]
    for yr,g in df.groupby("fiscal_year"):
        g=g.dropna(subset=[sig,"fwd_ret"])
        if len(g)<25: continue
        g=g.copy(); g["q"]=pd.qcut(g[sig],5,labels=False,duplicates="drop")+1
        rec.append((yr,(g[g.q==5].fwd_ret.mean()-g[g.q==1].fwd_ret.mean())*100))
    return pd.DataFrame(rec,columns=["year","sp"])
sp=annual_spread("z_AI_COMPOS")
fig,ax=plt.subplots(figsize=(11,4.6))
cols=[RED if y>=2022 else NAVY for y in sp.year]
ax.bar(sp.year,sp.sp,color=cols,edgecolor="k",lw=.5)
ax.axhline(0,color="k",lw=.8)
ax.set_title("Annual High-minus-Low AI-Disclosure Return Spread by Formation Year",fontweight="bold")
ax.set_ylabel("Q5 − Q1 forward 12m return (pp)"); ax.set_xlabel("Fiscal year of 10-K (signal)")
ax.set_xticks(sp.year); ax.tick_params(axis="x",rotation=45)
from matplotlib.patches import Patch
ax.legend(handles=[Patch(color=NAVY,label="Pre-ChatGPT (2010-2021)"),Patch(color=RED,label="Post-ChatGPT (2022-2024)")],frameon=False)
fig.tight_layout(); fig.savefig(f"{F}/fig3_annual_spread.png",bbox_inches="tight"); plt.close()

# ---- Fig 4: secular rise of AI language (context) ----
fig,ax=plt.subplots(figsize=(11,4.6))
for col,c,lab in [("ai_adoption_wmean",NAVY,"AI adoption"),("ai_infrastructure_wmean",RED,"AI infrastructure"),
                  ("ai_disruption_risk_wmean",GREEN,"AI disruption risk")]:
    g=df.groupby("fiscal_year")[col].mean()
    ax.plot(g.index,g.values,marker="o",color=c,label=lab,lw=1.8)
ax.axvline(2022,color="grey",ls="--",lw=1)
ax.set_title("Rising AI Language Across S&P 500 10-K Filings (cross-sectional mean similarity)",fontweight="bold")
ax.set_ylabel("Section-weighted cosine similarity to AI anchor"); ax.set_xlabel("Fiscal year")
ax.legend(frameon=False)
fig.tight_layout(); fig.savefig(f"{F}/fig4_ai_language_trend.png",bbox_inches="tight"); plt.close()

# ---- Fig 5: binscatter AI score vs forward return, pre vs post ----
fig,axes=plt.subplots(1,2,figsize=(12,4.6),sharey=True)
for ax,(title,d) in zip(axes,[("Pre-ChatGPT (2010-2021)",df[df.fiscal_year<2022]),
                              ("Post-ChatGPT (2022-2024)",df[df.fiscal_year>=2022])]):
    d=d.dropna(subset=["z_AI_COMPOS","fwd_ret"])
    d=d.copy(); d["bin"]=pd.qcut(d.z_AI_COMPOS,12,labels=False,duplicates="drop")
    b=d.groupby("bin").agg(x=("z_AI_COMPOS","mean"),y=("fwd_ret","mean"))
    ax.scatter(b.x,b.y*100,color=NAVY,s=45,zorder=3)
    z=np.polyfit(b.x,b.y*100,1); xs=np.linspace(b.x.min(),b.x.max(),50)
    ax.plot(xs,np.polyval(z,xs),color=RED,lw=2)
    ax.set_title(f"{title}\nslope={z[0]:.2f} pp / SD",fontsize=10)
    ax.set_xlabel("AI-disclosure composite (z-score)")
axes[0].set_ylabel("Mean forward 12m return (%)")
fig.suptitle("AI Disclosure vs. Subsequent Return: Binned Scatter",fontweight="bold",y=1.02)
fig.tight_layout(); fig.savefig(f"{F}/fig5_binscatter.png",bbox_inches="tight"); plt.close()

print("figures written:", [f"fig{i}" for i in range(1,6)])