"""
03_analysis.py  -- Empirical asset-pricing analysis
Does the SEMANTIC INTENSITY of AI disclosure in 10-K filings explain the
cross-section of subsequent S&P 500 stock returns?

Signals (from fiscal year Y, section-weighted cosine similarity to AI anchors):
  AI_ADOPT  = ai_adoption_wmean       (internal AI use)
  AI_INFRA  = ai_infrastructure_wmean (AI compute / capex exposure)
  AI_RISK   = ai_disruption_risk_wmean(AI product-disruption risk)
  AI_COMPOS = z(AI_ADOPT)+z(AI_INFRA) (standardized opportunity composite)

Outcome: forward 12-month return, July(Y+1)->June(Y+2)  [strictly tradeable].

Methods: (1) cross-sectional z-scoring within year; (2) quintile portfolio
sorts (EW & VW, raw & market-adjusted) with Newey-West t-stats; (3) long-short
spread regressed on FF5+Momentum to test for alpha; (4) two-way fixed-effects
panel regressions with double-clustered SE; (5) Fama-MacBeth; (6) a
post-ChatGPT (Nov-2022) structural-break interaction.
"""
import json, warnings
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
from statsmodels.stats.sandwich_covariance import cov_hac

warnings.filterwarnings("ignore")
D = "data/processed"; T = "results/tables"
m = pd.read_parquet(f"{D}/merged_panel.parquet")
ret = pd.read_parquet(f"{D}/monthly_returns.parquet")
ff = pd.read_parquet(f"{D}/ff5_monthly.parquet")
mom = pd.read_parquet(f"{D}/mom_monthly.parquet")
RES = {}

SIGNALS = {"AI_ADOPT": "ai_adoption_wmean",
           "AI_INFRA": "ai_infrastructure_wmean",
           "AI_RISK":  "ai_disruption_risk_wmean"}

df = m[m.fwd_ret.notna()].copy()

# winsorize returns & controls at 1/99 within sample for regressions
def wins(s, lo=.01, hi=.99):
    ql, qh = s.quantile(lo), s.quantile(hi)
    return s.clip(ql, qh)
df["fwd_w"] = wins(df["fwd_ret"])
df["fwd_madj_w"] = wins(df["fwd_ret_mktadj"])
df["mom_w"] = wins(df["mom_12m"])
df["vol_w"] = wins(df["vol_12m"])

# cross-sectional z-scores within formation year
def zscore(g):
    return (g - g.mean()) / g.std(ddof=0)
for k, col in SIGNALS.items():
    df[f"z_{k}"] = df.groupby("fiscal_year")[col].transform(zscore)
df["z_size"] = df.groupby("fiscal_year")["log_mktcap"].transform(zscore)
df["z_mom"]  = df.groupby("fiscal_year")["mom_w"].transform(zscore)
df["z_vol"]  = df.groupby("fiscal_year")["vol_w"].transform(zscore)
df["z_AI_COMPOS"] = df[["z_AI_ADOPT", "z_AI_INFRA"]].mean(axis=1)

# ============================================================
# 1. SUMMARY STATISTICS
# ============================================================
summ_cols = list(SIGNALS.values()) + ["fwd_ret", "fwd_ret_mktadj", "mom_12m", "vol_12m", "log_mktcap"]
summary = df[summ_cols].describe(percentiles=[.05,.25,.5,.75,.95]).T
summary.to_csv(f"{T}/t1_summary_stats.csv")
RES["n_obs"] = int(len(df))
RES["n_firms"] = int(df.ticker.nunique())
RES["yrs"] = f"{int(df.fiscal_year.min())}-{int(df.fiscal_year.max())}"
print("N obs", len(df), "| firms", df.ticker.nunique(), "| formation yrs", RES["yrs"])

# correlation among signals & with outcome
corr = df[["z_AI_ADOPT","z_AI_INFRA","z_AI_RISK","z_size","z_mom","fwd_ret"]].corr()
corr.to_csv(f"{T}/t1b_correlations.csv")

# ============================================================
# 2. QUINTILE PORTFOLIO SORTS
# ============================================================
def newey_mean_t(x, lags=3):
    x = x.dropna().values
    if len(x) < 3: return np.nan, np.nan
    X = np.ones((len(x),1))
    res = sm.OLS(x, X).fit()
    se = np.sqrt(cov_hac(res, nlags=lags)[0,0])
    return x.mean(), x.mean()/se

def sort_table(signal_z, weight=None, outcome="fwd_ret"):
    rows = []
    ann_means = {q: [] for q in range(1,6)}
    ann_ls = []
    for yr, g in df.groupby("fiscal_year"):
        g = g.dropna(subset=[signal_z, outcome])
        if len(g) < 25: continue
        g = g.copy()
        g["q"] = pd.qcut(g[signal_z], 5, labels=False, duplicates="drop") + 1
        for q in range(1,6):
            sub = g[g.q==q]
            if sub.empty: continue
            if weight is None:
                ann_means[q].append(sub[outcome].mean())
            else:
                w = sub[weight].clip(lower=0).fillna(0)
                ann_means[q].append(np.average(sub[outcome], weights=w) if w.sum()>0 else sub[outcome].mean())
        if ann_means[5] and ann_means[1]:
            ann_ls.append(ann_means[5][-1] - ann_means[1][-1])
    out = {}
    for q in range(1,6):
        mu, t = newey_mean_t(pd.Series(ann_means[q]))
        out[f"Q{q}"] = mu
    mu_ls, t_ls = newey_mean_t(pd.Series(ann_ls))
    out["Q5_minus_Q1"] = mu_ls
    out["t_stat"] = t_ls
    out["n_years"] = len(ann_ls)
    return out

sorts = {}
for k in ["AI_ADOPT","AI_INFRA","AI_RISK","AI_COMPOS"]:
    sorts[f"{k}_EW"] = sort_table(f"z_{k}", None, "fwd_ret")
    sorts[f"{k}_VW"] = sort_table(f"z_{k}", "mktcap", "fwd_ret")
    sorts[f"{k}_EW_madj"] = sort_table(f"z_{k}", None, "fwd_ret_mktadj")
sort_df = pd.DataFrame(sorts).T
sort_df.to_csv(f"{T}/t2_portfolio_sorts.csv")
RES["sorts"] = sort_df.round(4).to_dict("index")
print("\n=== Quintile sorts (forward 12m return) ===")
print(sort_df.round(4).to_string())

# ============================================================
# 3. LONG-SHORT MONTHLY TIME SERIES + FACTOR ALPHA
# ============================================================
ff = ff.merge(mom[["date","MOM"]], on="date", how="left")
ff["date"] = pd.to_datetime(ff["date"]).dt.to_period("M").dt.to_timestamp("M")
ff_ix = ff.set_index("date")
ret = ret.copy(); ret["date"] = pd.to_datetime(ret["date"]).dt.to_period("M").dt.to_timestamp("M")
ret_by = {t: d.set_index("date")["ret"] for t, d in ret.groupby("ticker")}

def ls_monthly(signal_z, n_groups=5):
    """Annually-rebalanced EW quintile portfolios, monthly returns."""
    monthly = {}  # date -> dict(q->list of rets)
    for yr, g in df.groupby("fiscal_year"):
        g = g.dropna(subset=[signal_z])
        if len(g) < 25: continue
        g = g.copy(); g["q"] = pd.qcut(g[signal_z], n_groups, labels=False, duplicates="drop")+1
        hold = pd.date_range(pd.Timestamp(yr+1,7,31), pd.Timestamp(yr+2,6,30), freq="ME")
        for q in [1, n_groups]:
            tks = g[g.q==q].ticker.tolist()
            for dt in hold:
                r = [ret_by[t].get(dt, np.nan) for t in tks if t in ret_by]
                r = [x for x in r if pd.notna(x)]
                if r:
                    monthly.setdefault(dt, {}).setdefault(q, []).append(np.mean(r))
    recs = []
    for dt, qd in sorted(monthly.items()):
        if 1 in qd and n_groups in qd:
            recs.append({"date": dt, "lo": np.mean(qd[1]), "hi": np.mean(qd[n_groups])})
    ts = pd.DataFrame(recs)
    ts["ls"] = ts["hi"] - ts["lo"]
    return ts

def factor_alpha(ts, label):
    s = ts.merge(ff_ix.reset_index(), on="date", how="inner").dropna(subset=["ls","MktRF","SMB","HML","RMW","CMA","MOM"])
    y = s["ls"] - 0   # long-short is already self-financing (RF nets out)
    out = {}
    # CAPM
    X1 = sm.add_constant(s[["MktRF"]]); r1 = sm.OLS(y, X1).fit(cov_type="HAC", cov_kwds={"maxlags":3})
    out["capm_alpha_m"] = r1.params["const"]; out["capm_alpha_t"] = r1.tvalues["const"]
    # FF5+MOM
    X2 = sm.add_constant(s[["MktRF","SMB","HML","RMW","CMA","MOM"]]); r2 = sm.OLS(y, X2).fit(cov_type="HAC", cov_kwds={"maxlags":3})
    out["ff6_alpha_m"] = r2.params["const"]; out["ff6_alpha_t"] = r2.tvalues["const"]
    out["ff6_alpha_ann"] = (1+r2.params["const"])**12 - 1
    out["mean_ls_m"] = y.mean(); out["mean_ls_ann"] = (1+y.mean())**12-1
    out["n_months"] = len(s)
    out["betas"] = {k: round(r2.params[k],3) for k in ["MktRF","SMB","HML","RMW","CMA","MOM"]}
    RES.setdefault("ls_alpha", {})[label] = {k:(round(v,5) if isinstance(v,float) else v) for k,v in out.items()}
    print(f"\n[{label}] mean LS {out['mean_ls_ann']*100:.2f}%/yr | CAPM a={out['capm_alpha_m']*1200:.2f}%/yr (t={out['capm_alpha_t']:.2f}) "
          f"| FF6 a={out['ff6_alpha_ann']*100:.2f}%/yr (t={out['ff6_alpha_t']:.2f}) | n={out['n_months']}m")
    return ts, out

ls_results = {}
for k in ["AI_ADOPT","AI_INFRA","AI_COMPOS","AI_RISK"]:
    ts = ls_monthly(f"z_{k}")
    ts.to_parquet(f"{D}/ls_ts_{k}.parquet", index=False)
    factor_alpha(ts, k)

# ============================================================
# 4. TWO-WAY FIXED-EFFECTS PANEL REGRESSIONS (double-clustered SE)
# ============================================================
from linearmodels.panel import PanelOLS
pdf = df.dropna(subset=["fwd_w","z_AI_ADOPT","z_AI_INFRA","z_AI_RISK","z_size","z_mom","z_vol"]).copy()
pdf = pdf.set_index(["ticker","fiscal_year"])

def fe_reg(formula, label):
    res = PanelOLS.from_formula(formula + " + EntityEffects + TimeEffects", data=pdf,
                                drop_absorbed=True).fit(cov_type="clustered", cluster_entity=True, cluster_time=True)
    return res

panel_specs = {
  "M1_adopt":   "fwd_w ~ z_AI_ADOPT",
  "M2_adopt_ctrl": "fwd_w ~ z_AI_ADOPT + z_size + z_mom + z_vol",
  "M3_infra_ctrl": "fwd_w ~ z_AI_INFRA + z_size + z_mom + z_vol",
  "M4_risk_ctrl":  "fwd_w ~ z_AI_RISK + z_size + z_mom + z_vol",
  "M5_all":     "fwd_w ~ z_AI_ADOPT + z_AI_INFRA + z_AI_RISK + z_size + z_mom + z_vol",
}
print("\n=== Two-way FE panel (DV=forward 12m return, winsorized) ===")
panel_out = {}
for lab, f in panel_specs.items():
    r = fe_reg(f, lab)
    coefs = {}
    for v in r.params.index:
        coefs[v] = {"coef": round(r.params[v],4), "t": round(r.tstats[v],2), "p": round(r.pvalues[v],4)}
    panel_out[lab] = {"coefs": coefs, "n": int(r.nobs), "r2_within": round(r.rsquared_within,4)}
    key = [v for v in r.params.index if v.startswith("z_AI")]
    print(f"{lab}: " + ", ".join(f"{v}={r.params[v]:+.4f}(t={r.tstats[v]:.2f})" for v in key) + f"  [n={int(r.nobs)}]")
RES["panel_fe"] = panel_out

# ============================================================
# 5. FAMA-MACBETH
# ============================================================
def fama_macbeth(xvars, label):
    coefs = {v: [] for v in ["const"]+xvars}
    for yr, g in df.groupby("fiscal_year"):
        g = g.dropna(subset=["fwd_ret"]+xvars)
        if len(g) < 30: continue
        X = sm.add_constant(g[xvars]); y = g["fwd_ret"]
        r = sm.OLS(y, X).fit()
        for v in ["const"]+xvars:
            coefs[v].append(r.params.get(v, np.nan))
    out = {}
    for v in ["const"]+xvars:
        mu, t = newey_mean_t(pd.Series(coefs[v]))
        out[v] = {"mean_coef": round(mu,4), "t": round(t,2)}
    out["n_years"] = len(coefs["const"])
    RES.setdefault("fama_macbeth", {})[label] = out
    print(f"\n[FM {label}] " + ", ".join(f"{v}={out[v]['mean_coef']:+.4f}(t={out[v]['t']:.2f})" for v in xvars))
    return out

print("\n=== Fama-MacBeth (DV=forward 12m return) ===")
fama_macbeth(["z_AI_ADOPT","z_size","z_mom","z_vol"], "adopt")
fama_macbeth(["z_AI_INFRA","z_size","z_mom","z_vol"], "infra")
fama_macbeth(["z_AI_ADOPT","z_AI_INFRA","z_AI_RISK","z_size","z_mom","z_vol"], "all")

# ============================================================
# 6. POST-CHATGPT STRUCTURAL BREAK  (formation years >= 2022)
# ============================================================
df["post"] = (df.fiscal_year >= 2022).astype(int)
print("\n=== Post-ChatGPT interaction (pooled OLS, year-clustered) ===")
brk = {}
for k in ["AI_ADOPT","AI_INFRA","AI_COMPOS"]:
    d2 = df.dropna(subset=["fwd_w",f"z_{k}","z_size","z_mom"]).copy()
    r = smf.ols(f"fwd_w ~ z_{k} * post + z_size + z_mom", data=d2).fit(
        cov_type="cluster", cov_kwds={"groups": d2["fiscal_year"]})
    iv = f"z_{k}:post"
    brk[k] = {"pre_slope": round(r.params[f"z_{k}"],4),
              "interaction": round(r.params[iv],4),
              "interaction_t": round(r.tvalues[iv],2),
              "post_slope": round(r.params[f"z_{k}"]+r.params[iv],4)}
    print(f"{k}: pre={brk[k]['pre_slope']:+.4f} | post-2022 extra={brk[k]['interaction']:+.4f} (t={brk[k]['interaction_t']:.2f}) | post total={brk[k]['post_slope']:+.4f}")
RES["structural_break"] = brk

# pre/post quintile spread comparison for headline
def spread_subperiod(signal_z, yrs):
    sp=[]
    for yr,g in df[df.fiscal_year.isin(yrs)].groupby("fiscal_year"):
        g=g.dropna(subset=[signal_z,"fwd_ret"])
        if len(g)<25: continue
        g=g.copy(); g["q"]=pd.qcut(g[signal_z],5,labels=False,duplicates="drop")+1
        sp.append(g[g.q==5].fwd_ret.mean()-g[g.q==1].fwd_ret.mean())
    mu,t=newey_mean_t(pd.Series(sp)); return round(mu,4),round(t,2),len(sp)
RES["subperiod_spread"]={}
for k in ["AI_ADOPT","AI_INFRA","AI_COMPOS"]:
    pre=spread_subperiod(f"z_{k}", range(2010,2022))
    post=spread_subperiod(f"z_{k}", range(2022,2025))
    RES["subperiod_spread"][k]={"pre_2010_2021":pre,"post_2022_2024":post}
    print(f"{k} Q5-Q1: pre(10-21)={pre[0]:+.4f}(t={pre[1]}) | post(22-24)={post[0]:+.4f}(t={post[1]})")

with open(f"{T}/results.json","w") as fh:
    json.dump(RES, fh, indent=2, default=str)
df.to_parquet(f"{D}/analysis_panel.parquet")
print("\nSaved results.json and analysis_panel.parquet")