"""
Empirical rigor suite for SAE-1.4 against synthetic HOPE logs.
Sections:
 A. Build SAE-1.4 predictors from raw logs
 B. Model comparison (baselines vs SAE) with walk-forward out-of-sample
 C. Belief-regime stratification (the key methodological insight)
 D. Cross-lag / Granger-style causality
 E. Fairness / non-moralization subgroup check
"""
import numpy as np, pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, r2_score
from scipy import stats
from sae14 import purity, blockers, shakti_effective, belief_total, awareness_raw

df = pd.read_csv("/home/claude/hope_synth.csv").sort_values(["user","day"]).reset_index(drop=True)

# ---------- A. Build SAE-1.4 terms ----------
Pcols = ["P1","P2","P3","P4","P5","P6"]
df["Purity"]   = df[Pcols].clip(0,1).mean(axis=1)
df["Blockers"] = df["blk"].clip(0,1)
df["ShaktiEff"] = np.maximum(0, 1 - df["phys"] - df["mental"]) * df["sh_strength"] * df["sh_stability"]
df["BeliefTotal"] = np.minimum(df["B1"], df["B2"]).clip(-1,1)
df["SAE_raw"] = df["BeliefTotal"] * (1.0 - df["Blockers"] + df["Purity"] + df["ShaktiEff"]) + df["grace_mag"]
df["Overdrive"] = df["sh_strength"] * (1 - df["Purity"])
y = df["ObservedAwareness"].values

def fit_eval(Xtr, ytr, Xte, yte):
    m = LinearRegression().fit(Xtr, ytr)
    pred = m.predict(Xte)
    return mean_absolute_error(yte, pred), r2_score(yte, pred)

# ---------- B. Walk-forward out-of-sample (train days<80, test days>=80) ----------
train = df[df.day < 80]; test = df[df.day >= 80]
def cols(c): return train[c].values, test[c].values

print("="*70)
print("B. MODEL COMPARISON  (walk-forward: train day<80, test day>=80)")
print("="*70)
results = {}
# baseline 0: constant
const_pred = np.full(len(test), train["ObservedAwareness"].mean())
results["B0 constant"] = (mean_absolute_error(test.ObservedAwareness, const_pred),
                          r2_score(test.ObservedAwareness, const_pred))
# baseline: purity only
results["B1 purity-only"] = fit_eval(train[["Purity"]], train.ObservedAwareness, test[["Purity"]], test.ObservedAwareness)
# baseline: blockers only
results["B2 blockers-only"] = fit_eval(train[["Blockers"]], train.ObservedAwareness, test[["Blockers"]], test.ObservedAwareness)
# baseline: wellness (mood/sleep proxy = mental+phys)
results["B3 wellness(phys+mental)"] = fit_eval(train[["phys","mental"]], train.ObservedAwareness, test[["phys","mental"]], test.ObservedAwareness)
# SAE-1.4 single composite predictor
results["SAE-1.4 raw composite"] = fit_eval(train[["SAE_raw"]], train.ObservedAwareness, test[["SAE_raw"]], test.ObservedAwareness)
# SAE-1.4 components as separate regressors
saecomp = ["BeliefTotal","Purity","Blockers","ShaktiEff","grace_mag"]
results["SAE-1.4 components"] = fit_eval(train[saecomp], train.ObservedAwareness, test[saecomp], test.ObservedAwareness)
# SAE-1.4 with belief interaction terms (regime-aware)
df["bt_x_pur"]=df.BeliefTotal*df.Purity; df["bt_x_blk"]=df.BeliefTotal*df.Blockers; df["bt_x_sh"]=df.BeliefTotal*df.ShaktiEff
train=df[df.day<80]; test=df[df.day>=80]
saeint=["BeliefTotal","bt_x_pur","bt_x_blk","bt_x_sh","grace_mag"]
results["SAE-1.4 + interactions"] = fit_eval(train[saeint], train.ObservedAwareness, test[saeint], test.ObservedAwareness)

print(f"{'model':32s} {'MAE':>8s} {'R2(oos)':>9s}")
for k,(mae,r2) in results.items():
    print(f"{k:32s} {mae:8.3f} {r2:9.3f}")

# ---------- C. Belief-regime stratification ----------
print("\n"+"="*70)
print("C. BELIEF-REGIME STRATIFICATION  (the key insight)")
print("="*70)
for label, sub in [("POOLED (all)", df), ("POSITIVE belief", df[df.BeliefTotal>=0]), ("NEGATIVE belief", df[df.BeliefTotal<0])]:
    rp = stats.pearsonr(sub.Purity, sub.ObservedAwareness)[0]
    rb = stats.pearsonr(sub.Blockers, sub.ObservedAwareness)[0]
    print(f"{label:18s}  corr(Purity,Aw)={rp:+.3f}   corr(Blockers,Aw)={rb:+.3f}   n={len(sub)}")
print("Note: pooling positive+negative regimes can flip/wash signs — must stratify.")

# ---------- D. Cross-lag causality ----------
print("\n"+"="*70)
print("D. CROSS-LAG / DIRECTIONALITY  (does Blockers(t) lead Awareness(t+1)?)")
print("="*70)
d = df.copy()
d["Aw_next"] = d.groupby("user")["ObservedAwareness"].shift(-1)
d["Blk_next"] = d.groupby("user")["Blockers"].shift(-1)
d["Pur_next"] = d.groupby("user")["Purity"].shift(-1)
dd = d.dropna()
# Blk(t)->Aw(t+1) controlling Aw(t)
import statsmodels.api as sm
def lagreg(yname, xs):
    X = sm.add_constant(dd[xs]); m = sm.OLS(dd[yname], X).fit()
    return {x: (m.params[x], m.pvalues[x]) for x in xs}
fwd = lagreg("Aw_next", ["Blockers","ObservedAwareness"])
rev = lagreg("Blk_next", ["ObservedAwareness","Blockers"])
print(f"Blockers(t) -> Awareness(t+1): beta={fwd['Blockers'][0]:+.3f} (p={fwd['Blockers'][1]:.1e})")
print(f"Awareness(t) -> Blockers(t+1): beta={rev['ObservedAwareness'][0]:+.3f} (p={rev['ObservedAwareness'][1]:.1e})")
purlag = lagreg("Blk_next", ["Purity","Blockers"])
print(f"Purity(t) -> Blockers(t+1):    beta={purlag['Purity'][0]:+.3f} (p={purlag['Purity'][1]:.1e})  (expect negative: purity reduces future blockers)")

# ---------- E. Fairness / non-moralization ----------
print("\n"+"="*70)
print("E. FAIRNESS CHECK  (does high constraint unfairly imply low awareness?)")
print("="*70)
hi = df[df.constraint_user > df.constraint_user.median()]
lo = df[df.constraint_user <= df.constraint_user.median()]
# match on Purity and Belief bins, compare ObservedAwareness
def matched_mean(g):
    return g.groupby([pd.cut(g.Purity,5), pd.cut(g.BeliefTotal,5)], observed=True)["ObservedAwareness"].mean().mean()
print(f"High-constraint users, matched mean Awareness: {matched_mean(hi):.2f}")
print(f"Low-constraint users,  matched mean Awareness: {matched_mean(lo):.2f}")
diff = matched_mean(lo)-matched_mean(hi)
print(f"Difference (lo-hi) after matching on Purity+Belief: {diff:+.2f}")
print("Ground truth was built so inner awareness does NOT depend on physical constraint.")
print("If diff is small, SAE measurement layer is fair; if large, constraints leak into awareness.")
print("\nDONE.")
