"""r3k_15_analysis_v2.py — Stage D (FIRC-0002): core asset-pricing tests on the rebuilt panel (v2 of r3k_03_analysis.py).
Packages L2 (cluster B), L3 (cluster A), L4 (clusters E/M/N, inference), L7 (facets / forms, cluster H), L8 (SPID, cluster I).
Reads data/processed_r3k_v2/merged_panel_v2.parquet + monthly_returns_v2.parquet, the staged factors and filing_scores.
Writes results/r3k_v2/results.json, data/processed_r3k_v2/analysis_panel_v2.parquet, ls_ts_v2_*.parquet, vw_weights_diag.csv.
Differences from r3k_03 (tagged): [B] winsorization WITHIN formation year (fwd_w) with the author's pooled version kept (fwd_w_pooled);
every sort/spread/FM reported on raw | winsorized | median | strict-12 outcomes with plain t, NW t (n beside), 95 % CI (t and cohort
bootstrap) and the top-1 % contribution; [A] z_size from within-year winsorized log cap, VW weights on the corrected cap with the
single-firm weight diagnostic; [E] two-way (firm x year) clustered and wild-cluster-bootstrap inference for every pooled regression;
[H] AI Risk, PC1, equal-form and 10-K composites carried as signals; [I] SPID any-tech / any-AI binary contrasts with CI and MDE.
Run: OMP_NUM_THREADS=4 nice -n 10 .venv-review/bin/python -I scripts/r3k_15_analysis_v2.py
"""
import os, sys, json, time, warnings, numpy as np, pandas as pd, statsmodels.api as sm, statsmodels.formula.api as smf
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); import r3k_inference as I
from linearmodels.panel import PanelOLS
warnings.filterwarnings("ignore"); T0 = time.time()
D = "data/processed_r3k"; V = "data/processed_r3k_v2"; R = "data/raw_russell3000"; T = "results/r3k_v2"
m = pd.read_parquet(f"{V}/merged_panel_v2.parquet"); ret = pd.read_parquet(f"{V}/monthly_returns_v2.parquet")
ff = pd.read_parquet(f"{D}/ff5_monthly.parquet").merge(pd.read_parquet(f"{D}/mom_monthly.parquet"), on="date", how="left")
RES = {"script": "r3k_15_analysis_v2.py", "panel": "merged_panel_v2.parquet (r3k_14)"}
SIGNALS = {"AI_ADOPT": "ai_adoption_wmean", "AI_INFRA": "ai_infrastructure_wmean", "AI_RISK": "ai_disruption_risk_wmean"}
PRE, POST = range(2010, 2022), range(2022, 2025)
df = m[m.fwd_ret.notna() & (m.fiscal_year <= 2024)].copy()
def zscore(g): return (g - g.mean()) / g.std(ddof=0)
def wins_within(s): return s.clip(s.quantile(.01), s.quantile(.99))
def wins_pooled(s): return s.clip(s.quantile(.01), s.quantile(.99))
# [B] outcomes and controls
for c, w in [("fwd_ret", "fwd_w"), ("fwd_ret_mktadj", "fwd_madj_w"), ("mom_12m", "mom_w"), ("vol_12m", "vol_w"), ("log_mktcap", "logcap_w"), ("fwd_ret_12m_strict", "fwd_strict_w"), ("fwd_ret_6m", "fwd6_w"), ("fwd_ret_24m", "fwd24_w")]:
    df[w] = df.groupby("fiscal_year")[c].transform(wins_within)
df["fwd_w_pooled"] = wins_pooled(df["fwd_ret"])
for k, col in SIGNALS.items(): df[f"z_{k}"] = df.groupby("fiscal_year")[col].transform(zscore)
df["z_AI_COMPOS"] = df[["z_AI_ADOPT", "z_AI_INFRA"]].mean(axis=1)
# [H] PC1 of the three facets, computed before the sorts (cluster H), signed with adoption, re-standardised within year
Z = df[["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK"]].dropna(); Zc = Z - Z.mean(); U, S, Vt = np.linalg.svd(Zc.values, full_matrices=False); pc1 = Zc.values @ Vt[0]
if np.corrcoef(pc1, Z["z_AI_ADOPT"])[0, 1] < 0: pc1 = -pc1
df.loc[Z.index, "z_AI_PC1"] = pc1; df["z_AI_PC1"] = df.groupby("fiscal_year")["z_AI_PC1"].transform(zscore); RES["pca"] = {"evr": [round(float(x), 3) for x in (S ** 2) / np.sum(S ** 2)], "loadings": [round(float(x), 3) for x in Vt[0]]}
df["z_CYBER"] = df.groupby("fiscal_year")["cybersecurity_wmean"].transform(zscore)
# [H] form-specific composites (10-K / 10-Q / 8-K) from the staged filing-level scores; equal-form composite
fs = pd.read_parquet(f"{R}/filing_scores"); fl = pd.read_parquet(f"{R}/filings")[["accession_number", "form_type"]]
fs = fs.drop(columns=[c for c in ["form_type"] if c in fs.columns]).merge(fl, on="accession_number", how="left"); fs["fiscal_year"] = pd.to_datetime(fs["filing_date"]).dt.year
by = fs.groupby(["ticker", "fiscal_year", "form_type"]).agg(ad=("ai_adoption_wmean", "mean"), inf=("ai_infrastructure_wmean", "mean"), n=("accession_number", "count")).reset_index()
for form, tag in [("10-K", "10k"), ("10-Q", "10q"), ("8-K", "8k")]:
    b = by[by.form_type == form][["ticker", "fiscal_year", "ad", "inf", "n"]].rename(columns={"ad": f"ad_{tag}", "inf": f"inf_{tag}", "n": f"n_{tag}"})
    df = df.drop(columns=[c for c in df.columns if c in b.columns and c not in ("ticker", "fiscal_year")]).merge(b, on=["ticker", "fiscal_year"], how="left")
    df[f"z_AI_{tag}"] = (df.groupby("fiscal_year")[f"ad_{tag}"].transform(zscore) + df.groupby("fiscal_year")[f"inf_{tag}"].transform(zscore)) / 2
df["z_AI_formw"] = df[["z_AI_10k", "z_AI_10q", "z_AI_8k"]].mean(axis=1, skipna=False); df["z_AI_formw"] = df.groupby("fiscal_year")["z_AI_formw"].transform(zscore)
# SPID [I]
df["spid_tech_log"] = np.log1p(df["spid_mentions_tech"]); df["spid_ai_any"] = (df["spid_mentions_ai"] + df["spid_mentions_ai_tech_tag"] > 0).astype(float); df["spid_tech_any"] = (df["spid_mentions_tech"] > 0).astype(float)
df["z_SPID_TECH"] = df.groupby("fiscal_year")["spid_tech_log"].transform(zscore)
# [A] controls: z_size from within-year winsorized log cap
df["z_size"] = df.groupby("fiscal_year")["logcap_w"].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["post"] = (df.fiscal_year >= 2022).astype(int); df = df.sort_values(["cik", "fiscal_year"]).reset_index(drop=True)
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())}"; RES["n_by_year"] = df.groupby("fiscal_year").size().to_dict()
RES["n_sp500_firm_years"] = int(df.in_sp500.sum()); RES["n_large"] = int((df.size_group == "large").sum()); RES["n_without_cap"] = int(df.mktcap.isna().sum()); RES["n_strict12"] = int(df.fwd_ret_12m_strict.notna().sum())
print("N", len(df), "firms", df.ticker.nunique(), RES["yrs"], "| no cap", RES["n_without_cap"], flush=True)
summ_cols = list(SIGNALS.values()) + ["spid_mentions_tech", "spid_ai_any", "spid_tech_any", "fwd_ret", "fwd_w", "fwd_ret_mktadj", "mom_12m", "vol_12m", "log_mktcap", "mktcap", "total_chunks"]
df[summ_cols].describe(percentiles=[.01, .05, .25, .5, .75, .95, .99]).T.to_csv(f"{T}/t1_summary_stats.csv")
df[["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_AI_PC1", "z_SPID_TECH", "z_CYBER", "z_size", "z_mom", "fwd_ret"]].corr().to_csv(f"{T}/t1b_correlations.csv")
RES["corr"] = df[["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_SPID_TECH", "z_CYBER", "z_AI_PC1"]].corr().round(3).to_dict()

# ---------------- sort machinery ----------------
VWD = []
def spread_stats(ls, by_cohort=None):
    mu, t = I.nw_t(ls); out = {"Q5_minus_Q1": mu, "t_stat": t, "t_plain": I.plain_t(ls), "ci95": I.ci95(ls), "ci95_boot": I.block_boot_ci(ls), "n_years": len(ls)}
    if by_cohort is not None: out["by_cohort"] = {int(k): round(v, 4) for k, v in by_cohort.items()}
    return out
def sort_table(d, z, weight=None, oc="fwd_ret", stat="mean", minn=40, yrs=None, diag_label=None):
    ann = {q: [] for q in range(1, 6)}; ls = []; ls_trim = []; byc = {}
    for yr, g in d.groupby("fiscal_year"):
        if yrs is not None and yr not in yrs: continue
        g = g.dropna(subset=[z, oc]).copy()
        if len(g) < minn: continue
        g["q"] = pd.qcut(g[z].rank(method="first"), 5, labels=False) + 1; p99 = g[oc].quantile(.99); vals = {}; vals_trim = {}
        for q in range(1, 6):
            s = g[g.q == q]
            if s.empty: continue
            if weight is None:
                vals[q] = s[oc].median() if stat == "median" else s[oc].mean(); st = s[s[oc] <= p99]; vals_trim[q] = st[oc].mean() if len(st) else np.nan
            else:
                w = s[weight].clip(lower=0).fillna(0); vals[q] = np.average(s[oc], weights=w) if w.sum() > 0 else s[oc].mean(); vals_trim[q] = vals[q]
                if diag_label and w.sum() > 0:
                    ws = (w / w.sum()).sort_values(ascending=False); VWD.append({"sort": diag_label, "cohort": int(yr), "q": q, "n": len(s), "max_weight": round(float(ws.iloc[0]), 4), "max_ticker": s.loc[ws.index[0], "ticker"], "top10_weight": round(float(ws.iloc[:10].sum()), 4)})
            ann[q].append(vals[q])
        if 5 in vals and 1 in vals: ls.append(vals[5] - vals[1]); ls_trim.append(vals_trim[5] - vals_trim[1]); byc[int(yr)] = vals[5] - vals[1]
    out = {f"Q{q}": float(np.mean(ann[q])) if ann[q] else np.nan for q in range(1, 6)}; out.update(spread_stats(ls, byc)); out["spread_ex_top1pct"] = float(np.nanmean(ls_trim)) if ls_trim else np.nan; out["top1pct_contribution"] = out["Q5_minus_Q1"] - out["spread_ex_top1pct"]; return out
OUTS = [("", "fwd_ret", "mean"), ("_w", "fwd_w", "mean"), ("_wp", "fwd_w_pooled", "mean"), ("_med", "fwd_ret", "median"), ("_strict", "fwd_ret_12m_strict", "mean"), ("_strict_w", "fwd_strict_w", "mean")]
SIGL = ["AI_ADOPT", "AI_INFRA", "AI_RISK", "AI_COMPOS", "AI_PC1", "AI_formw", "AI_10k", "SPID_TECH", "CYBER"]
sorts = {}
for k in SIGL:
    for sfx, oc, st in OUTS:
        sorts[f"{k}_EW{sfx}"] = sort_table(df, f"z_{k}", None, oc, st)
        if sfx in ("", "_w"): sorts[f"{k}_VW{sfx}"] = sort_table(df, f"z_{k}", "mktcap", oc, st, diag_label=f"{k}_VW{sfx}"); sorts[f"{k}_VW_v1{sfx}"] = sort_table(df, f"z_{k}", "mktcap_v1", oc, st, diag_label=f"{k}_VW_v1{sfx}")
    sorts[f"{k}_EW_madj"] = sort_table(df, f"z_{k}", None, "fwd_ret_mktadj"); sorts[f"{k}_EW_madj_w"] = sort_table(df, f"z_{k}", None, "fwd_madj_w")
RES["sorts"] = {k: {kk: (round(vv, 5) if isinstance(vv, float) else vv) for kk, vv in v.items()} for k, v in sorts.items()}
pd.DataFrame(sorts).T.drop(columns=["by_cohort"]).to_csv(f"{T}/t2_portfolio_sorts.csv")
vwd = pd.DataFrame(VWD); vwd.to_csv(f"{T}/vw_weights_diag.csv", index=False)
RES["vw_diag"] = {"max_single_firm_weight_v2": float(vwd[vwd["sort"].str.contains("_VW$|_VW_w$", regex=True)].max_weight.max()), "max_single_firm_weight_v1": float(vwd[vwd["sort"].str.contains("_VW_v1")].max_weight.max()),
                  "n_cells_above_25pct_v2": int((vwd[vwd["sort"].str.contains("_VW$|_VW_w$", regex=True)].max_weight > 0.25).sum()), "n_cells_above_25pct_v1": int((vwd[vwd["sort"].str.contains("_VW_v1")].max_weight > 0.25).sum()),
                  "worst_cells_v2": vwd[vwd["sort"] == "AI_COMPOS_VW"].nlargest(5, "max_weight")[["cohort", "q", "max_weight", "max_ticker"]].to_dict("records"), "worst_cells_v1": vwd[vwd["sort"] == "AI_COMPOS_VW_v1"].nlargest(5, "max_weight")[["cohort", "q", "max_weight", "max_ticker"]].to_dict("records")}
for k in ["AI_COMPOS_EW", "AI_COMPOS_EW_w", "AI_COMPOS_EW_med", "AI_COMPOS_VW", "AI_COMPOS_VW_v1", "AI_COMPOS_EW_strict"]:
    s = sorts[k]; print(f"{k:22s} Q1 {s['Q1']:+.4f} Q5 {s['Q5']:+.4f} spread {s['Q5_minus_Q1']:+.4f} NW t {s['t_stat']:.2f} plain t {s['t_plain']:.2f} CI {np.round(s['ci95'], 4)} n {s['n_years']}", flush=True)
print("VW diag", {k: v for k, v in RES["vw_diag"].items() if not k.startswith("worst")}, flush=True)

# ---------------- monthly long-short + alphas ----------------
ff["date"] = pd.to_datetime(ff["date"]).dt.to_period("M").dt.to_timestamp("M"); ff_ix = ff.set_index("date")
ret["date"] = pd.to_datetime(ret["date"]).dt.to_period("M").dt.to_timestamp("M"); ret_by = {t: g.set_index("date")["ret"] for t, g in ret.groupby("ticker")}
def ls_monthly(d, z, n=5, minn=40):
    monthly = {}
    for yr, g in d.groupby("fiscal_year"):
        g = g.dropna(subset=[z]).copy()
        if len(g) < minn: continue
        g["q"] = pd.qcut(g[z].rank(method="first"), n, labels=False) + 1; hold = pd.date_range(pd.Timestamp(yr + 1, 7, 31), pd.Timestamp(yr + 2, 6, 30), freq="ME")
        for q in (1, n):
            tks = [t for t in g[g.q == q].ticker if t in ret_by]
            for dt in hold:
                r = [ret_by[t].get(dt, np.nan) for t in tks]; r = [x for x in r if pd.notna(x)]
                if len(r) >= 10: monthly.setdefault(dt, {}).setdefault(q, []).append(np.mean(r))
    recs = [{"date": dt, "lo": np.mean(v[1]), "hi": np.mean(v[n])} for dt, v in sorted(monthly.items()) if 1 in v and n in v]; ts = pd.DataFrame(recs); ts["ls"] = ts["hi"] - ts["lo"]; return ts
SR_TRIALS = []
def factor_alpha(ts, record=True):
    s = ts.merge(ff_ix.reset_index(), on="date", how="inner").dropna(subset=["ls", "MktRF", "SMB", "HML", "RMW", "CMA", "MOM"]); y = s["ls"]; out = {}
    if len(s) < 12: return {"n_months": int(len(s))}
    r1 = sm.OLS(y, sm.add_constant(s[["MktRF"]])).fit(cov_type="HAC", cov_kwds={"maxlags": 3}); r2 = sm.OLS(y, sm.add_constant(s[["MktRF", "SMB", "HML", "RMW", "CMA", "MOM"]])).fit(cov_type="HAC", cov_kwds={"maxlags": 3})
    mu, t = I.nw_t(y.values); sr = mu / y.std() * np.sqrt(12)
    out.update(mean_ls_ann_pct=round(mu * 1200, 2), hac_t=round(t, 2), plain_t=round(I.plain_t(y.values), 2), sharpe=round(sr, 2), skew=round(float(y.skew()), 3), kurt=round(float(y.kurt() + 3), 3), capm_alpha_ann_pct=round(r1.params["const"] * 1200, 2), capm_alpha_t=round(r1.tvalues["const"], 2),
               ff6_alpha_ann_pct=round(r2.params["const"] * 1200, 2), ff6_alpha_t=round(r2.tvalues["const"], 2), n_months=int(len(s)), betas={k: round(r2.params[k], 3) for k in ["MktRF", "SMB", "HML", "RMW", "CMA", "MOM"]})
    if record: SR_TRIALS.append(sr)
    return out
RES["ls_alpha"] = {}; POST_START = pd.Timestamp("2023-07-31")
for k in ["AI_ADOPT", "AI_INFRA", "AI_COMPOS", "AI_RISK", "AI_PC1", "AI_formw", "AI_10k", "SPID_TECH", "CYBER"]:
    ts = ls_monthly(df, f"z_{k}"); ts.to_parquet(f"{V}/ls_ts_v2_{k}.parquet", index=False)
    RES["ls_alpha"][k] = {"full": factor_alpha(ts), "pre": factor_alpha(ts[ts.date < POST_START]), "post": factor_alpha(ts[ts.date >= POST_START])}
    print(k, "LS full", RES["ls_alpha"][k]["full"].get("mean_ls_ann_pct"), "ff6a", RES["ls_alpha"][k]["full"].get("ff6_alpha_ann_pct"), "(t", RES["ls_alpha"][k]["full"].get("ff6_alpha_t"), ") | post ff6a", RES["ls_alpha"][k]["post"].get("ff6_alpha_ann_pct"), "(t", RES["ls_alpha"][k]["post"].get("ff6_alpha_t"), ")", flush=True)

# ---------------- two-way FE panel (fwd_w within-year) ----------------
pdf = df.dropna(subset=["fwd_w", "z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_size", "z_mom", "z_vol"]).set_index(["ticker", "fiscal_year"])
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", "M6_spid": "fwd_w ~ z_SPID_TECH + z_size + z_mom + z_vol", "M7_all_spid": "fwd_w ~ z_AI_COMPOS + z_SPID_TECH + z_size + z_mom + z_vol", "M8_compos": "fwd_w ~ z_AI_COMPOS + z_size + z_mom + z_vol"}
RES["panel_fe"] = {}
for lab, f in specs.items():
    r = PanelOLS.from_formula(f + " + EntityEffects + TimeEffects", data=pdf, drop_absorbed=True).fit(cov_type="clustered", cluster_entity=True, cluster_time=True)
    RES["panel_fe"][lab] = {"coefs": {v: {"coef": round(r.params[v], 4), "t": round(r.tstats[v], 2)} for v in r.params.index}, "n": int(r.nobs), "r2_within": round(r.rsquared_within, 4)}
print("FE M8", RES["panel_fe"]["M8_compos"]["coefs"]["z_AI_COMPOS"], flush=True)

# ---------------- Fama-MacBeth: raw | within-year winsorized | pooled-winsorized | strict ----------------
FB15 = I.fixed_b_cv(15, 3)
def fama_macbeth(d, xvars, minn=40, oc="fwd_ret"):
    coefs = {v: [] for v in ["const"] + xvars}; ns = []
    for yr, g in d.groupby("fiscal_year"):
        g = g.dropna(subset=[oc] + xvars)
        if len(g) < minn: continue
        r = sm.OLS(g[oc], sm.add_constant(g[xvars])).fit(); ns.append(int(r.nobs))
        for v in coefs: coefs[v].append(r.params.get(v, np.nan))
    out = {}
    for v in coefs:
        mu, t = I.nw_t(coefs[v]); out[v] = {"mean_coef": round(mu, 4), "t": round(t, 2), "t_plain": round(I.plain_t(coefs[v]), 2), "ci95": [round(x, 4) for x in I.ci95(coefs[v])], "ci95_boot": [round(x, 4) for x in I.block_boot_ci(coefs[v])], "fixed_b_cv05": FB15["cv"], "by_year": [round(x, 4) for x in coefs[v]]}
    out["n_years"] = len(coefs["const"]); out["n_mean_per_year"] = int(np.mean(ns)) if ns else 0; return out
RES["fama_macbeth"] = {}
for sfx, oc in [("", "fwd_ret"), ("_w", "fwd_w"), ("_wp", "fwd_w_pooled"), ("_strict", "fwd_ret_12m_strict")]:
    RES["fama_macbeth"]["compos" + sfx] = fama_macbeth(df, ["z_AI_COMPOS", "z_size", "z_mom", "z_vol"], oc=oc)
    if sfx in ("", "_w"):
        for lab, xv in [("adopt", ["z_AI_ADOPT"]), ("infra", ["z_AI_INFRA"]), ("risk", ["z_AI_RISK"]), ("all", ["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK"]), ("pc1", ["z_AI_PC1"]), ("formw", ["z_AI_formw"]), ("tenk", ["z_AI_10k"]), ("spid", ["z_SPID_TECH"]), ("compos_spid", ["z_AI_COMPOS", "z_SPID_TECH"]), ("cyber", ["z_CYBER"]), ("compos_cyber", ["z_AI_COMPOS", "z_CYBER"]), ("spid_any_ai", ["spid_ai_any"]), ("spid_any_tech", ["spid_tech_any"]), ("compos_spid_any", ["z_AI_COMPOS", "spid_ai_any"]), ("compos_nocontrols", [])]:
            RES["fama_macbeth"][lab + sfx] = fama_macbeth(df, (xv if lab != "compos_nocontrols" else ["z_AI_COMPOS"]) + (["z_size", "z_mom", "z_vol"] if lab != "compos_nocontrols" else []), oc=oc)
for k in ["compos", "compos_w", "compos_wp", "compos_strict", "all_w", "pc1_w", "formw_w", "tenk_w"]:
    v = RES["fama_macbeth"][k]; key = [x for x in v if x.startswith("z_AI")][0]; print("FM", k, key, v[key]["mean_coef"], "NW t", v[key]["t"], "plain t", v[key]["t_plain"], "CI", v[key]["ci95"], flush=True)

# ---------------- structural break regressions: firm / year / two-way / wild bootstrap ----------------
def break_reg(d, z, oc="fwd_w", boot=True):
    d2 = d.dropna(subset=[oc, z, "z_size", "z_mom"]).copy(); d2["zp"] = d2[z] * d2["post"]
    X = sm.add_constant(d2[[z, "post", "zp", "z_size", "z_mom"]].values); y = d2[oc].values; j = 3
    r = sm.OLS(y, X).fit(); gf = pd.factorize(d2.ticker)[0]; gy = pd.factorize(d2.fiscal_year)[0]
    se_f = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gf).cov_params())); se_y = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gy).cov_params())); se_tw = np.sqrt(np.diag(I.cgm_cov(r, gf, gy)))
    out = {"pre_slope": round(float(r.params[1]), 4), "interaction": round(float(r.params[j]), 4), "post_slope": round(float(r.params[1] + r.params[j]), 4), "interaction_t_firm": round(float(r.params[j] / se_f[j]), 2), "interaction_t_year": round(float(r.params[j] / se_y[j]), 2), "interaction_t_twoway": round(float(r.params[j] / se_tw[j]), 2),
           "pre_slope_t_firm": round(float(r.params[1] / se_f[1]), 2), "pre_slope_t_twoway": round(float(r.params[1] / se_tw[1]), 2), "interaction_ci95_firm": [round(float(r.params[j] - 1.96 * se_f[j]), 4), round(float(r.params[j] + 1.96 * se_f[j]), 4)], "n": int(r.nobs), "G_years": int(gy.max() + 1)}
    if boot:
        for wname in ["rademacher", "webb"]:
            b = I.wild_cluster_boot(y, X, gy, j, B=9999, weights=wname); out[f"interaction_p_wild_{wname}"] = b["p"]
    return out
RES["structural_break"] = {}
for k in ["AI_ADOPT", "AI_INFRA", "AI_COMPOS", "AI_RISK", "AI_PC1", "AI_formw", "SPID_TECH"]:
    RES["structural_break"][k] = break_reg(df, f"z_{k}"); RES["structural_break"][k + "_raw"] = break_reg(df, f"z_{k}", oc="fwd_ret", boot=False)
    print("break", k, {x: RES["structural_break"][k][x] for x in ["interaction", "interaction_t_firm", "interaction_t_twoway", "interaction_p_wild_rademacher", "interaction_p_wild_webb"]}, flush=True)

# ---------------- sub-period spreads (raw | winsorized | median), n beside every t ----------------
def spread_sub(d, z, yrs, oc="fwd_ret", stat="mean", minn=40):
    s = sort_table(d, z, None, oc, stat, minn, yrs=set(yrs)); return {"mean": round(s["Q5_minus_Q1"], 4) if s["n_years"] else None, "t": round(s["t_stat"], 2) if s["t_stat"] == s["t_stat"] else None, "t_plain": round(s["t_plain"], 2) if s["t_plain"] == s["t_plain"] else None, "ci95": s["ci95"], "ci95_boot": s["ci95_boot"], "n": s["n_years"], "by_cohort": s["by_cohort"], "top1pct_contribution": round(s["top1pct_contribution"], 4) if s["n_years"] else None}
RES["subperiod_spread"] = {}
for k in SIGL:
    RES["subperiod_spread"][k] = {}
    for sfx, oc, st in [("", "fwd_ret", "mean"), ("_w", "fwd_w", "mean"), ("_med", "fwd_ret", "median"), ("_strict", "fwd_ret_12m_strict", "mean")]:
        RES["subperiod_spread"][k]["pre" + sfx] = spread_sub(df, f"z_{k}", PRE, oc, st); RES["subperiod_spread"][k]["post" + sfx] = spread_sub(df, f"z_{k}", POST, oc, st)
v = RES["subperiod_spread"]["AI_COMPOS"]; print("COMPOS pre/post raw", v["pre"]["mean"], v["post"]["mean"], "| wins", v["pre_w"]["mean"], v["post_w"]["mean"], "| median", v["pre_med"]["mean"], v["post_med"]["mean"], "| post cohorts wins", v["post_w"]["by_cohort"], flush=True)

# ---------------- pooled Q1/Q5 regressions with firm / two-way / wild inference [E] ----------------
def pooled(d, yrs, dv="fwd_ret", z="z_AI_COMPOS", boot=True):
    d = d[d.fiscal_year.isin(list(yrs))].dropna(subset=[z, dv]).copy(); d["q"] = d.groupby("fiscal_year")[z].transform(lambda s: pd.qcut(s.rank(method="first"), 5, labels=False) + 1)
    d = d[d.q.isin([1, 5])].copy(); d["hi"] = (d.q == 5).astype(float); yd = pd.get_dummies(d.fiscal_year.astype(int), prefix="y", drop_first=True).astype(float)
    X = np.c_[np.ones(len(d)), d["hi"].values, yd.values]; y = d[dv].values; r = sm.OLS(y, X).fit(); gf = pd.factorize(d.ticker)[0]; gy = pd.factorize(d.fiscal_year)[0]
    se_f = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gf).cov_params())); G = gy.max() + 1
    out = {"coef_pct": round(float(r.params[1]) * 100, 2), "t_firm": round(float(r.params[1] / se_f[1]), 2), "ci95_firm_pct": [round(float(r.params[1] - 1.96 * se_f[1]) * 100, 2), round(float(r.params[1] + 1.96 * se_f[1]) * 100, 2)], "n": int(r.nobs), "G_years": int(G)}
    if G >= 2:
        se_tw = np.sqrt(np.diag(I.cgm_cov(r, gf, gy))); se_y = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gy).cov_params())); out["t_twoway"] = round(float(r.params[1] / se_tw[1]), 2); out["t_year"] = round(float(r.params[1] / se_y[1]), 2)
        if boot:
            for wname in ["rademacher", "webb"]: out[f"p_wild_{wname}"] = I.wild_cluster_boot(y, X, gy, 1, B=9999, weights=wname)["p"]
    return out, d
RES["pooled"] = {}
for sfx, dv in [("", "fwd_ret"), ("_w", "fwd_w"), ("_med", None)]:
    if dv is None: continue
    RES["pooled"]["post" + sfx], dpost = pooled(df, POST, dv); RES["pooled"]["pre" + sfx], dpre = pooled(df, PRE, dv); RES["pooled"]["post_excl_2024" + sfx], _ = pooled(df, range(2022, 2024), dv); RES["pooled"]["post_2023_2024" + sfx], _ = pooled(df, range(2023, 2025), dv)
    dall = pd.concat([dpre.assign(post=0), dpost.assign(post=1)]); dall["hp"] = dall.hi * dall.post; yd = pd.get_dummies(dall.fiscal_year.astype(int), prefix="y", drop_first=True).astype(float)
    X = np.c_[np.ones(len(dall)), dall["hi"].values, dall["hp"].values, yd.values]; y = dall[dv].values; r = sm.OLS(y, X).fit(); gf = pd.factorize(dall.ticker)[0]; gy = pd.factorize(dall.fiscal_year)[0]
    se_f = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gf).cov_params())); se_tw = np.sqrt(np.diag(I.cgm_cov(r, gf, gy)))
    RES["pooled"]["diff" + sfx] = {"coef_pct": round(float(r.params[2]) * 100, 2), "t_firm": round(float(r.params[2] / se_f[2]), 2), "t_twoway": round(float(r.params[2] / se_tw[2]), 2), "ci95_firm_pct": [round(float(r.params[2] - 1.96 * se_f[2]) * 100, 2), round(float(r.params[2] + 1.96 * se_f[2]) * 100, 2)], "n": int(r.nobs),
                                  "p_wild_rademacher": I.wild_cluster_boot(y, X, gy, 2, B=9999, weights="rademacher")["p"], "p_wild_webb": I.wild_cluster_boot(y, X, gy, 2, B=9999, weights="webb")["p"], "mde_pct_firm": round(I.mde(se_f[2]) * 100, 2)}
print("pooled post", RES["pooled"]["post"], "\npooled post wins", RES["pooled"]["post_w"], "\ndiff", RES["pooled"]["diff"], flush=True)

# ---------------- size groups (both partitions) and index membership ----------------
RES["groups"] = {}
def group_block(d, gname):
    blk = {"n": int(len(d)), "firms": int(d.ticker.nunique())}
    for sfx, oc, st in [("", "fwd_ret", "mean"), ("_w", "fwd_w", "mean"), ("_med", "fwd_ret", "median")]:
        blk["sort_EW" + sfx] = sort_table(d, "z_AI_COMPOS", None, oc, st); blk["pre" + sfx] = spread_sub(d, "z_AI_COMPOS", PRE, oc, st); blk["post" + sfx] = spread_sub(d, "z_AI_COMPOS", POST, oc, st)
        cs = blk["sort_EW" + sfx]["by_cohort"]; post = {k: v for k, v in cs.items() if k >= 2022}
        blk["post_loo" + sfx] = {k: round(float(np.mean([v for kk, v in post.items() if kk != k])), 4) for k in post}
        blk["pooled_post" + sfx], _ = pooled(d, POST, oc, boot=False) if st == "mean" else (None, None); blk["pooled_post_excl_2024" + sfx], _ = pooled(d, range(2022, 2024), oc, boot=False) if st == "mean" else (None, None)
    blk["fm_compos"] = fama_macbeth(d, ["z_AI_COMPOS", "z_size", "z_mom", "z_vol"]); blk["fm_compos_w"] = fama_macbeth(d, ["z_AI_COMPOS", "z_size", "z_mom", "z_vol"], oc="fwd_w")
    ts = ls_monthly(d, "z_AI_COMPOS"); ts.to_parquet(f"{V}/ls_ts_v2_AI_COMPOS_{gname}.parquet", index=False); blk["ls"] = {"full": factor_alpha(ts), "pre": factor_alpha(ts[ts.date < POST_START]), "post": factor_alpha(ts[ts.date >= POST_START])}
    return blk
for gname, mask in [("large_top1000", df.size_group == "large"), ("small_rest", df.size_group == "small"), ("large_pct42", df.size_group_pct42 == "large"), ("small_pct42", df.size_group_pct42 == "small"), ("large_top1000_v1", df.size_group_v1 == "large"), ("small_rest_v1", df.size_group_v1 == "small"), ("sp500", df.in_sp500 == True), ("non_sp500", df.in_sp500 == False)]:
    RES["groups"][gname] = group_block(df[mask], gname); g = RES["groups"][gname]
    print(gname, "n", g["n"], "full", round(g["sort_EW"]["Q5_minus_Q1"], 4), "| pre/post raw", g["pre"]["mean"], g["post"]["mean"], "| wins", g["pre_w"]["mean"], g["post_w"]["mean"], "| LOO wins", g["post_loo_w"], "| ff6 post", g["ls"]["post"].get("ff6_alpha_ann_pct"), flush=True)
df["z_AI_COMPOS_sizeneu"] = df.groupby(["fiscal_year", "size_group"])["z_AI_COMPOS"].transform(lambda g: zscore(g) if g.std(ddof=0) > 0 else g * 0)
RES["size_neutral"] = {"full": sort_table(df, "z_AI_COMPOS_sizeneu"), "full_w": sort_table(df, "z_AI_COMPOS_sizeneu", None, "fwd_w"), "pre": spread_sub(df, "z_AI_COMPOS_sizeneu", PRE), "post": spread_sub(df, "z_AI_COMPOS_sizeneu", POST), "pre_w": spread_sub(df, "z_AI_COMPOS_sizeneu", PRE, "fwd_w"), "post_w": spread_sub(df, "z_AI_COMPOS_sizeneu", POST, "fwd_w")}

# ---------------- contrasts: hi x post x large, hi x post x member [E] ----------------
def contrast(d, grp, dv="fwd_w"):
    d = d.dropna(subset=["z_AI_COMPOS", dv, grp]).copy(); d["q"] = d.groupby("fiscal_year")["z_AI_COMPOS"].transform(lambda s: pd.qcut(s.rank(method="first"), 5, labels=False) + 1); d = d[d.q.isin([1, 5])].copy()
    d["hi"] = (d.q == 5).astype(float); d["g"] = d[grp].astype(float); d["hp"] = d.hi * d.post; d["hg"] = d.hi * d.g; d["pg"] = d.post * d.g; d["hpg"] = d.hi * d.post * d.g
    yd = pd.get_dummies(d.fiscal_year.astype(int), prefix="y", drop_first=True).astype(float); X = np.c_[np.ones(len(d)), d[["hi", "hp", "g", "hg", "pg", "hpg"]].values, yd.values]; y = d[dv].values
    r = sm.OLS(y, X).fit(); gf = pd.factorize(d.ticker)[0]; gy = pd.factorize(d.fiscal_year)[0]; se_f = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gf).cov_params())); se_tw = np.sqrt(np.diag(I.cgm_cov(r, gf, gy)))
    out = {}
    for j, nm in [(2, "hi_x_post"), (4, "hi_x_group"), (6, "hi_x_post_x_group")]:
        out[nm] = {"coef_pct": round(float(r.params[j]) * 100, 2), "t_firm": round(float(r.params[j] / se_f[j]), 2), "t_twoway": round(float(r.params[j] / se_tw[j]), 2), "ci95_firm_pct": [round(float(r.params[j] - 1.96 * se_f[j]) * 100, 2), round(float(r.params[j] + 1.96 * se_f[j]) * 100, 2)], "mde_pct_firm": round(I.mde(se_f[j]) * 100, 2)}
    out["hi_x_post_x_group"]["p_wild_rademacher"] = I.wild_cluster_boot(y, X, gy, 6, B=9999, weights="rademacher")["p"]; out["hi_x_post_x_group"]["p_wild_webb"] = I.wild_cluster_boot(y, X, gy, 6, B=9999, weights="webb")["p"]; out["n"] = int(r.nobs); return out
df["is_large"] = (df.size_group == "large").astype(float).where(df.size_group.notna()); df["is_member"] = df.in_sp500.astype(float)
RES["contrasts"] = {"large_vs_small_w": contrast(df, "is_large"), "large_vs_small_raw": contrast(df, "is_large", "fwd_ret"), "member_vs_non_w": contrast(df, "is_member"), "member_vs_non_raw": contrast(df, "is_member", "fwd_ret")}
# cohort-block bootstrap of the difference in (post - pre) spreads between groups
def group_diff_boot(ga, gb, sfx="_w", B=9999):
    a = RES["groups"][ga]["sort_EW" + sfx]["by_cohort"]; b = RES["groups"][gb]["sort_EW" + sfx]["by_cohort"]; yrs = sorted(set(a) & set(b)); pre = [y for y in yrs if y < 2022]; post = [y for y in yrs if y >= 2022]
    def stat(pre_i, post_i): return (np.mean([a[y] for y in post_i]) - np.mean([a[y] for y in pre_i])) - (np.mean([b[y] for y in post_i]) - np.mean([b[y] for y in pre_i]))
    obs = stat(pre, post); rng = np.random.default_rng(I.RNG_SEED); bs = [stat(rng.choice(pre, len(pre)), rng.choice(post, len(post))) for _ in range(B)]
    return {"diff_in_diff_pct": round(obs * 100, 2), "ci95_boot_pct": [round(float(np.quantile(bs, .025)) * 100, 2), round(float(np.quantile(bs, .975)) * 100, 2)], "n_pre": len(pre), "n_post": len(post)}
RES["contrasts"]["large_minus_small_postpre_boot_w"] = group_diff_boot("large_top1000", "small_rest"); RES["contrasts"]["member_minus_non_postpre_boot_w"] = group_diff_boot("sp500", "non_sp500"); RES["contrasts"]["large_minus_small_postpre_boot_raw"] = group_diff_boot("large_top1000", "small_rest", "")
print("contrast large", RES["contrasts"]["large_vs_small_w"]["hi_x_post_x_group"], "| member", RES["contrasts"]["member_vs_non_w"]["hi_x_post_x_group"], flush=True)

# ---------------- SPID binary contrasts, CI, MDE [I] ----------------
def binary_contrast(d, flag, yrs, oc="fwd_ret", min_treated=15):
    diffs = {}; ns = {}
    for yr, g in d[d.fiscal_year.isin(list(yrs))].groupby("fiscal_year"):
        a = g[g[flag] == 1][oc].dropna(); b = g[g[flag] == 0][oc].dropna()
        if len(a) >= min_treated and len(b) >= min_treated: diffs[int(yr)] = float(a.mean() - b.mean()); ns[int(yr)] = int(len(a))
    ls = list(diffs.values()); mu, t = I.nw_t(ls); se = (np.std(ls, ddof=1) / np.sqrt(len(ls))) if len(ls) > 1 else np.nan
    return {"mean_pct": round(mu * 100, 2) if ls else None, "t": round(t, 2) if t == t else None, "t_plain": round(I.plain_t(ls), 2) if len(ls) > 1 else None, "ci95_pct": [round(x * 100, 2) for x in I.ci95(ls)] if len(ls) > 1 else [None, None], "n_cohorts": len(ls), "n_treated_by_cohort": ns, "mde_pct": round(I.mde(se) * 100, 2) if se == se else None, "by_cohort_pct": {k: round(v * 100, 2) for k, v in diffs.items()}}
def firm_level_contrast(d, flag, yrs, oc="fwd_w"):
    d = d[d.fiscal_year.isin(list(yrs))].dropna(subset=[oc, flag, "z_size", "z_mom"]).copy(); yd = pd.get_dummies(d.fiscal_year.astype(int), prefix="y", drop_first=True).astype(float)
    X = np.c_[np.ones(len(d)), d[flag].values, d[["z_size", "z_mom"]].values, yd.values]; y = d[oc].values; r = sm.OLS(y, X).fit(); gf = pd.factorize(d.ticker)[0]; se = np.sqrt(np.diag(r.get_robustcov_results(cov_type="cluster", groups=gf).cov_params()))
    return {"coef_pct": round(float(r.params[1]) * 100, 2), "t_firm": round(float(r.params[1] / se[1]), 2), "ci95_pct": [round(float(r.params[1] - 1.96 * se[1]) * 100, 2), round(float(r.params[1] + 1.96 * se[1]) * 100, 2)], "mde_pct": round(I.mde(se[1]) * 100, 2), "n": int(r.nobs), "n_treated": int(d[flag].sum())}
RES["spid"] = {"share_firm_years_any_ai": round(float(df.spid_ai_any.mean()), 4), "share_firm_years_any_tech": round(float(df.spid_tech_any.mean()), 4), "share_post_any_ai": round(float(df[df.post == 1].spid_ai_any.mean()), 4), "any_ai_by_year": df.groupby("fiscal_year").spid_ai_any.mean().round(4).to_dict(), "n_any_ai_by_year": df.groupby("fiscal_year").spid_ai_any.sum().astype(int).to_dict(),
               "spid_tech_quantiles": {q: float(df.spid_mentions_tech.quantile(q)) for q in [.5, .75, .9, .95]}}
for flag in ["spid_ai_any", "spid_tech_any"]:
    RES["spid"][flag] = {}
    for per, yrs in [("full", range(2010, 2025)), ("pre", PRE), ("post", POST)]:
        RES["spid"][flag][per] = {"cohort_raw": binary_contrast(df, flag, yrs), "cohort_w": binary_contrast(df, flag, yrs, "fwd_w"), "firm_w": firm_level_contrast(df, flag, yrs), "firm_raw": firm_level_contrast(df, flag, yrs, "fwd_ret")}
# terciles among non-zero tech-project firms (>= 15 treated per cohort)
nz = df[df.spid_mentions_tech > 0].copy(); nz["spid_terc"] = nz.groupby("fiscal_year")["spid_mentions_tech"].transform(lambda s: pd.qcut(s.rank(method="first"), 3, labels=False) + 1 if len(s) >= 45 else np.nan)
tsp = {}
for yr, g in nz.dropna(subset=["spid_terc"]).groupby("fiscal_year"):
    a = g[g.spid_terc == 3].fwd_w; b = g[g.spid_terc == 1].fwd_w
    if len(a) >= 15 and len(b) >= 15: tsp[int(yr)] = float(a.mean() - b.mean())
mu, t = I.nw_t(list(tsp.values())); RES["spid"]["tercile_T3_minus_T1_nonzero_w"] = {"mean_pct": round(mu * 100, 2), "t": round(t, 2) if t == t else None, "t_plain": round(I.plain_t(list(tsp.values())), 2), "ci95_pct": [round(x * 100, 2) for x in I.ci95(list(tsp.values()))], "n_cohorts": len(tsp), "by_cohort_pct": {k: round(v * 100, 2) for k, v in tsp.items()}}
# legacy quintile spread (SPID_TECH) kept only for the before/after, with its MDE
sp_leg = RES["subperiod_spread"]["SPID_TECH"]; ls_leg = list(RES["sorts"]["SPID_TECH_EW"]["by_cohort"].values()); RES["spid"]["legacy_quintile_full_mde_pct"] = round(I.mde(np.std(ls_leg, ddof=1) / np.sqrt(len(ls_leg))) * 100, 2)
print("SPID any-AI post cohort", RES["spid"]["spid_ai_any"]["post"]["cohort_raw"], "\n firm", RES["spid"]["spid_ai_any"]["post"]["firm_w"], flush=True)

# ---------------- side-by-side samples: Russell 3000 vs S&P 500 members ----------------
RES["samples"] = {}
for sname, d in [("r3k", df), ("sp500", df[df.in_sp500 == True].copy())]:
    for k, col in SIGNALS.items(): d[f"zz_{k}"] = d.groupby("fiscal_year")[col].transform(zscore)
    d["zz_AI_COMPOS"] = d[["zz_AI_ADOPT", "zz_AI_INFRA"]].mean(axis=1); d["zz_SPID_TECH"] = d.groupby("fiscal_year")["spid_tech_log"].transform(zscore); d["zz_AI_PC1"] = d.groupby("fiscal_year")["z_AI_PC1"].transform(zscore); d["zz_AI_formw"] = d.groupby("fiscal_year")["z_AI_formw"].transform(zscore)
    for c in ["size", "mom", "vol"]: d[f"zz_{c}"] = d.groupby("fiscal_year")[{"size": "logcap_w", "mom": "mom_w", "vol": "vol_w"}[c]].transform(zscore)
    for c, w in [("fwd_ret", "fwd_w_s"), ("fwd_ret_mktadj", "fwd_madj_w_s")]: d[w] = d.groupby("fiscal_year")[c].transform(wins_within)
    blk = {"n": int(len(d)), "firms": int(d.ticker.nunique()), "sorts": {}, "ls": {}, "fm": {}, "break": {}, "sub": {}}
    for k in ["AI_ADOPT", "AI_INFRA", "AI_RISK", "AI_COMPOS", "AI_PC1", "AI_formw", "SPID_TECH"]:
        for sfx, oc, st in [("", "fwd_ret", "mean"), ("_w", "fwd_w_s", "mean"), ("_med", "fwd_ret", "median")]:
            blk["sorts"][f"{k}_EW{sfx}"] = sort_table(d, f"zz_{k}", None, oc, st)
            if sfx != "_med": blk["sorts"][f"{k}_VW{sfx}"] = sort_table(d, f"zz_{k}", "mktcap", oc, st, diag_label=f"{sname}_{k}_VW{sfx}")
            blk["sub"].setdefault(k, {})["pre" + sfx] = spread_sub(d, f"zz_{k}", PRE, oc, st); blk["sub"][k]["post" + sfx] = spread_sub(d, f"zz_{k}", POST, oc, st)
        ts = ls_monthly(d, f"zz_{k}"); ts.to_parquet(f"{V}/ls_ts_v2_{k}_{sname}.parquet", index=False); blk["ls"][k] = {"full": factor_alpha(ts), "pre": factor_alpha(ts[ts.date < POST_START]), "post": factor_alpha(ts[ts.date >= POST_START])}
        if k not in ("AI_RISK",): blk["break"][k] = break_reg(d.rename(columns={f"zz_{k}": f"zzz_{k}", "zz_size": "z_size_s", "zz_mom": "z_mom_s"}).assign(z_size=lambda x: x.z_size_s, z_mom=lambda x: x.z_mom_s), f"zzz_{k}", oc="fwd_w_s", boot=(k == "AI_COMPOS"))
    for sfx, oc in [("", "fwd_ret"), ("_w", "fwd_w_s")]:
        blk["fm"]["compos" + sfx] = fama_macbeth(d, ["zz_AI_COMPOS", "zz_size", "zz_mom", "zz_vol"], oc=oc); blk["fm"]["all" + sfx] = fama_macbeth(d, ["zz_AI_ADOPT", "zz_AI_INFRA", "zz_AI_RISK", "zz_size", "zz_mom", "zz_vol"], oc=oc); blk["fm"]["spid" + sfx] = fama_macbeth(d, ["zz_SPID_TECH", "zz_size", "zz_mom", "zz_vol"], oc=oc); blk["fm"]["pc1" + sfx] = fama_macbeth(d, ["zz_AI_PC1", "zz_size", "zz_mom", "zz_vol"], oc=oc); blk["fm"]["formw" + sfx] = fama_macbeth(d, ["zz_AI_formw", "zz_size", "zz_mom", "zz_vol"], oc=oc)
    blk["cohorts_pct"] = {k_: round(v * 100, 1) for k_, v in blk["sorts"]["AI_COMPOS_EW"]["by_cohort"].items()}; blk["cohorts_w_pct"] = {k_: round(v * 100, 1) for k_, v in blk["sorts"]["AI_COMPOS_EW_w"]["by_cohort"].items()}
    blk["pooled_post"], _ = pooled(d, POST, "fwd_ret", "zz_AI_COMPOS"); blk["pooled_post_w"], _ = pooled(d, POST, "fwd_w_s", "zz_AI_COMPOS"); blk["pooled_post_excl_2024_w"], _ = pooled(d, range(2022, 2024), "fwd_w_s", "zz_AI_COMPOS", boot=False)
    RES["samples"][sname] = blk
    print("SAMPLE", sname, "n", blk["n"], "compos EW raw", round(blk["sorts"]["AI_COMPOS_EW"]["Q5_minus_Q1"], 4), "wins", round(blk["sorts"]["AI_COMPOS_EW_w"]["Q5_minus_Q1"], 4), "(NW t", round(blk["sorts"]["AI_COMPOS_EW_w"]["t_stat"], 2), ") pre/post wins", blk["sub"]["AI_COMPOS"]["pre_w"]["mean"], blk["sub"]["AI_COMPOS"]["post_w"]["mean"], "| pooled post wins", blk["pooled_post_w"], flush=True)
pd.DataFrame(VWD).to_csv(f"{T}/vw_weights_diag.csv", index=False)
RES["sr_trials_ann"] = [round(float(x), 4) for x in SR_TRIALS]; RES["runtime_s"] = round(time.time() - T0)
json.dump(RES, open(f"{T}/results.json", "w"), indent=2, default=str); df.to_parquet(f"{V}/analysis_panel_v2.parquet", index=False); print("saved results/r3k_v2/results.json in", RES["runtime_s"], "s", flush=True)