# ============================
# Linear SES×Age charts (95% CIs)
# Interpersonal + Opportunity-taking
# ============================

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

# --------- INPUT ---------
# CHANGE THIS if your CSV is in a different location/name
path = "Public Copy of socioeconomic background and interpersonal risks - Socioeconomic Factors in Career Risk Appetite.csv"

# --------- Helpers ---------
def col_match(df, prefix):
"""Find the first column whose name starts with `prefix`."""
for c in df.columns:
if str(c).startswith(prefix):
return c
raise KeyError(f"Column starting with '{prefix}' not found.")

def parse_age(age):
"""Convert age band labels to numeric midpoints."""
s = str(age).strip().lower()
if "-" in s:
parts = s.replace(" ", "").split("-")
try:
low, high = int(parts[0]), int(parts[1])
return (low + high) / 2
except:
pass
if "under" in s:
try:
n = int("".join(ch for ch in s if ch.isdigit()))
return max(18, n - 1)
except:
return 19
if ("over" in s) or ("plus" in s) or ("+" in s) or ("above" in s):
digits = "".join(ch for ch in s if ch.isdigit())
return int(digits) + 2.5 if digits else 72.5
try:
return float(age)
except:
return np.nan

def ols_hc3(y, X_df):
"""
OLS with intercept + HC3 robust covariance.
Returns beta (including intercept), covariance, dof, tcrit, and R^2.
"""
X = np.column_stack([np.ones(len(X_df)), X_df.values])
yv = y.values.astype(float)

XtX = X.T @ X
XtX_inv = np.linalg.inv(XtX)
beta = XtX_inv @ (X.T @ yv)

yhat = X @ beta
resid = yv - yhat

H = X @ XtX_inv @ X.T
h = np.clip(np.diag(H), 0, 0.999999)
w = (resid / (1 - h))**2
S = X.T * w @ X
cov_hc3 = XtX_inv @ S @ XtX_inv

dof = max(len(yv) - X.shape[1], 1)
tcrit = stats.t.ppf(0.975, df=dof)

sst = float(((yv - yv.mean())**2).sum())
sse = float((resid**2).sum())
r2 = 1 - sse/sst if sst > 0 else np.nan
return beta, cov_hc3, dof, float(tcrit), float(r2)

# --------- Load & prepare data ---------
df_raw = pd.read_csv(path)

token_col = "Token" if "Token" in df_raw.columns else col_match(df_raw, "Token")
socio_col = col_match(df_raw, "Thinking about your childhood generally")
inter_col = col_match(df_raw, "How safe do you feel taking interpersonal risks")
opp_col = col_match(df_raw, "When high-profile, higher-risk opportunities arise")
age_col = col_match(df_raw, "What's your age?")

# Deduplicate respondents by token
agg = (
df_raw[[token_col, socio_col, inter_col, opp_col, age_col]]
.groupby(token_col, dropna=False)
.agg({socio_col:"mean", inter_col:"mean", opp_col:"mean", age_col:"first"})
.reset_index()
.rename(columns={
socio_col:"SES",
inter_col:"PS_interpersonal",
opp_col:"PS_opportunity"
})
)
agg["Age_numeric"] = agg[age_col].apply(parse_age)

# --------- Linear-age model: outcome ~ SES + Age + SES×Age ---------
def fit_linear_meta(df, outcome_col):
dm = df[["SES","Age_numeric", outcome_col]].dropna().copy()
dm["Age_c"] = dm["Age_numeric"] - dm["Age_numeric"].mean()
dm["SES_c"] = dm["SES"] - dm["SES"].mean()
dm["SESxAge"] = dm["SES_c"] * dm["Age_c"]

X = dm[["SES_c","Age_c","SESxAge"]]
beta, cov, dof, tcrit, r2 = ols_hc3(dm[outcome_col], X)
meta = {
"N": len(dm),
"R2": r2,
"age_mean": dm["Age_numeric"].mean(),
"ses_mean": dm["SES"].mean(),
"p20": float(np.nanpercentile(dm["SES"], 20)),
"p50": float(np.nanpercentile(dm["SES"], 50)),
"p80": float(np.nanpercentile(dm["SES"], 80)),
"beta": beta,
"cov": cov,
"tcrit": tcrit
}
return meta

def preds_linear(meta, s_val, ages_grid):
A = ages_grid - meta["age_mean"]
S = s_val - meta["ses_mean"]
# Design: [1, SES_c, Age_c, SESxAge]
Xg = np.column_stack([np.ones_like(A), S*np.ones_like(A), A, S*A])
yhat = Xg @ meta["beta"]
var = np.einsum('ij,jk,ik->i', Xg, meta["cov"], Xg)
se = np.sqrt(np.maximum(var, 0))
lo = yhat - meta["tcrit"]*se
hi = yhat + meta["tcrit"]*se
return yhat, lo, hi

def compute_ylim(meta, ages_grid):
"""Tight y-limits across the 3 SES lines (2nd–98th pct) + 10% padding."""
ses_levels = [meta["p20"], meta["p50"], meta["p80"]]
y_all = []
for s in ses_levels:
yhat, lo, hi = preds_linear(meta, s, ages_grid)
y_all += [yhat, lo, hi]
y_all = np.concatenate(y_all)
low = np.nanpercentile(y_all, 2)
high = np.nanpercentile(y_all, 98)
pad = (high - low) * 0.10
return low - pad, high + pad

def plot_three_lines_linear(meta, title, ylabel, out_path, ages_grid=None, truncate_axis=True):
if ages_grid is None:
ages_grid = np.linspace(agg["Age_numeric"].min(), agg["Age_numeric"].max(), 160)

ses_levels = [meta["p20"], meta["p50"], meta["p80"]]
labels = [f"Low (20th≈{ses_levels[0]:.2f})",
f"Mid (50th≈{ses_levels[1]:.2f})",
f"High (80th≈{ses_levels[2]:.2f})"]

# Optional truncation to make slopes visually steeper
if truncate_axis:
ylo, yhi = compute_ylim(meta, ages_grid)
else:
ylo, yhi = 0.0, 6.0

for s, lab in zip(ses_levels, labels):
yhat, lo, hi = preds_linear(meta, s, ages_grid)
plt.plot(ages_grid, yhat, label=lab)
plt.fill_between(ages_grid, lo, hi, alpha=0.15)

plt.ylim(ylo, yhi)
plt.xlabel("Age (years)")
plt.ylabel(ylabel)
subtitle = f"(Linear age; SES×Age; HC3 95% CIs; N={meta['N']}; R²={meta['R2']:.3f}; {'y-axis truncated' if truncate_axis else 'full 0–6 y-axis'})"
plt.title(f"{title}\n{subtitle}")
plt.legend()
plt.tight_layout()
plt.savefig(out_path, dpi=200)
plt.show()
return out_path, (ylo, yhi)

# --------- Fit models & plot ---------
meta_inter = fit_linear_meta(agg, "PS_interpersonal")
meta_opp = fit_linear_meta(agg, "PS_opportunity")

ages_grid = np.linspace(agg["Age_numeric"].min(), agg["Age_numeric"].max(), 160)

plot_three_lines_linear(
meta_inter,
"Interpersonal Risk-Taking Safety Across Career By Childhood SES",
"Predicted Interpersonal PS (0–6)",
"linear_interpersonal_truncated.png",
ages_grid=ages_grid,
truncate_axis=True,
)

plot_three_lines_linear(
meta_opp,
"Opportunity-Taking Risk Safety Across Career By Childhood SES",
"Predicted Opportunity-taking PS (0–6)",
"linear_opportunity_truncated.png",
ages_grid=ages_grid,
truncate_axis=True,
)