"""
Cross-sectional versus within-runner age gradients in marathon performance.
Boston Marathon finisher records, 2001-2014.

Thornton OR & Li W.

Data source: llimllib/bostonmarathon public data dumps (GitHub).

Design note on identification:
  Within a runner, age = year - birth cohort exactly, so runner fixed effects,
  a linear age term and year fixed effects are perfectly collinear (the classic
  age-period-cohort problem). Year effects are therefore omitted from the
  fixed-effects models and their potential influence is probed by excluding
  anomalous race years in sensitivity analyses.
"""
import glob, json
import numpy as np
import pandas as pd
from scipy import stats

# ---------------------------------------------------------------- LOAD
frames = []
for f in sorted(glob.glob("boston/*.csv")):
    year = int(f.split("/")[-1][:4])
    d = pd.read_csv(f, dtype=str, low_memory=False)
    cols = [c for c in ["name", "gender", "age", "official", "country", "state", "city", "bib"] if c in d.columns]
    d = d[cols].copy()
    d["year"] = year
    frames.append(d)

df = pd.concat(frames, ignore_index=True)
df["age"] = pd.to_numeric(df["age"], errors="coerce")
df["finish"] = pd.to_numeric(df["official"], errors="coerce")   # decimal minutes
df = df.dropna(subset=["name", "gender", "age", "finish"])
df["gender"] = df["gender"].str.strip().str.upper()
df = df[df.gender.isin(["M", "F"])]

N_RAW = len(df)

# --- exclusions -------------------------------------------------------------
# Wheelchair / handcycle divisions ride rather than run; identified by bib prefix.
if "bib" in df:
    wheel = df["bib"].astype(str).str.match(r"^W", na=False)
else:
    wheel = pd.Series(False, index=df.index)
N_WHEEL = int(wheel.sum())
df = df[~wheel]

# 2013 was stopped by the bombing at 4:09:43 elapsed. Its "finishers" are a
# speed-truncated subset of starters, so the year is excluded from all
# performance models. It is retained only for the attrition/return analysis.
df_perf = df[df.year != 2013].copy()
N_2013 = int((df.year == 2013).sum())

# Implausible records
plaus = df_perf.finish.between(100, 480) & df_perf.age.between(18, 85)
N_IMPLAUS = int((~plaus).sum())
df_perf = df_perf[plaus].copy()

# ---------------------------------------------------------------- LINKAGE
def add_key(d):
    d = d.copy()
    d["key"] = (d["name"].str.strip().str.lower().str.replace(r"\s+", " ", regex=True)
                + "|" + d["gender"] + "|" + d["country"].fillna("").str.strip().str.upper())
    d["cohort"] = d["year"] - d["age"]
    return d

df_perf = add_key(df_perf)

# A key is accepted as one person only if the implied birth cohort is stable
# (tolerance of 1 year absorbs birthday-relative-to-race-day timing) and the
# person appears in at least three distinct race years.
grp = df_perf.groupby("key")
cohort_ok = grp["cohort"].agg(lambda s: (s.max() - s.min()) <= 1)
n_years = grp["year"].nunique()
valid = cohort_ok & (n_years >= 3)
panel_keys = set(valid[valid].index)

panel = df_perf[df_perf.key.isin(panel_keys)].copy()
panel = panel.sort_values(["key", "year"]).drop_duplicates(["key", "year"])

N_PANEL_RUNNERS = panel.key.nunique()
N_PANEL_OBS = len(panel)

# ---------------------------------------------------------------- MODELS
def ols(y, X, names):
    """OLS with heteroskedasticity-consistent (HC1) standard errors."""
    XtX_inv = np.linalg.pinv(X.T @ X)
    b = XtX_inv @ X.T @ y
    resid = y - X @ b
    n, k = X.shape
    S = (X * resid[:, None]).T @ (X * resid[:, None])
    cov = XtX_inv @ S @ XtX_inv * (n / max(n - k, 1))
    se = np.sqrt(np.diag(cov))
    t = b / se
    p = 2 * (1 - stats.norm.cdf(np.abs(t)))
    ss_res = float(resid @ resid)
    ss_tot = float(((y - y.mean()) ** 2).sum())
    return dict(names=names, b=b, se=se, t=t, p=p, n=int(n),
                r2=1 - ss_res / ss_tot,
                terms={nm: dict(b=float(bb), se=float(ss), t=float(tt), p=float(pp))
                       for nm, bb, ss, tt, pp in zip(names, b, se, t, p)})


AGE_C = 45.0   # centring constant, near the sample median


def cross_sectional(d, label):
    """Between-person age gradient: one row per runner-year, year dummies included."""
    a = d.age.values - AGE_C
    female = (d.gender.values == "F").astype(float)
    years = sorted(d.year.unique())[1:]           # first year is the reference
    ydum = np.column_stack([(d.year.values == y).astype(float) for y in years])
    X = np.column_stack([np.ones(len(d)), a, a ** 2, female, ydum])
    names = ["intercept", "age", "age^2", "female"] + [f"year_{y}" for y in years]
    r = ols(d.finish.values, X, names)
    r["label"] = label
    r["model"] = "Cross-sectional OLS (HC1), year fixed effects"
    return r


def within_runner(d, label):
    """Within-person age gradient: runner fixed effects via demeaning.

    Year fixed effects are NOT included: within a runner, age is an exact linear
    function of year, so the two cannot be separately identified alongside
    runner fixed effects.
    """
    d = d.copy()
    a = d.age.values - AGE_C
    X = np.column_stack([a, a ** 2])
    y = d.finish.values
    codes = pd.factorize(d.key.values)[0]
    n_g = codes.max() + 1

    def demean(M):
        M = np.atleast_2d(M.T).T
        out = np.empty_like(M, dtype=float)
        for j in range(M.shape[1]):
            s = np.bincount(codes, weights=M[:, j], minlength=n_g)
            c = np.bincount(codes, minlength=n_g)
            out[:, j] = M[:, j] - (s / c)[codes]
        return out

    Xd, yd = demean(X), demean(y.reshape(-1, 1))[:, 0]
    r = ols(yd, Xd, ["age", "age^2"])
    # correct dof for the absorbed fixed effects
    dof_scale = (len(d) - 2) / max(len(d) - n_g - 2, 1)
    r["se"] = r["se"] * np.sqrt(dof_scale)
    r["t"] = r["b"] / r["se"]
    r["p"] = 2 * (1 - stats.norm.cdf(np.abs(r["t"])))
    r["terms"] = {nm: dict(b=float(b), se=float(s), t=float(t), p=float(p))
                  for nm, b, s, t, p in zip(r["names"], r["b"], r["se"], r["t"], r["p"])}
    r.update(label=label, n_groups=int(n_g),
             model="Within-runner (runner fixed effects), HC1, dof-corrected")
    return r


def slope_at(r, age):
    """Marginal minutes per year of age at a given age, with delta-method SE."""
    a = age - AGE_C
    i_age = r["names"].index("age")
    i_a2 = r["names"].index("age^2")
    b = r["b"][i_age] + 2 * a * r["b"][i_a2]
    se = np.sqrt(r["se"][i_age] ** 2 + (2 * a) ** 2 * r["se"][i_a2] ** 2)
    return float(b), float(se)


# --- year-effect adjustment -------------------------------------------------
# Race-day conditions differ sharply between years (2012 was run in record heat).
# Within a runner these period shocks cannot be separated from age. They CAN be
# identified in the full cross-section, where age varies within every year.
# Step 1: estimate year effects on all finishers, holding age and sex constant.
# Step 2: subtract them, and fit the within-runner model to adjusted times.
_cs = cross_sectional(df_perf, "year-effect estimation")
_year_eff = {int(nm.split("_")[1]): float(_cs["terms"][nm]["b"])
             for nm in _cs["names"] if nm.startswith("year_")}
_ref_year = min(df_perf.year.unique())
_year_eff[_ref_year] = 0.0
YEAR_EFFECTS = _year_eff

for _d in (df_perf, panel):
    _d["finish_adj"] = _d["finish"] - _d["year"].map(YEAR_EFFECTS).astype(float)


def within_runner_adj(d, label):
    """Within-runner model fitted to year-adjusted finish times."""
    d2 = d.copy()
    d2["finish"] = d2["finish_adj"]
    r = within_runner(d2, label)
    r["model"] = ("Within-runner (runner fixed effects) on year-adjusted times; "
                  "year effects estimated from the full cross-section")
    return r


results = {}
results["cross_all"] = cross_sectional(df_perf, "Cross-sectional, all finishers")
results["cross_panel"] = cross_sectional(panel, "Cross-sectional, restricted to panel runners")
results["within"] = within_runner(panel, "Within-runner, panel")
results["within_m"] = within_runner(panel[panel.gender == "M"], "Within-runner, men")
results["within_f"] = within_runner(panel[panel.gender == "F"], "Within-runner, women")

# Sensitivity: drop 2012, an extreme-heat race day
results["within_adj"] = within_runner_adj(panel, "Within-runner, year-adjusted (primary)")
results["within_adj_m"] = within_runner_adj(panel[panel.gender == "M"], "Within-runner, year-adjusted, men")
results["within_adj_f"] = within_runner_adj(panel[panel.gender == "F"], "Within-runner, year-adjusted, women")
results["within_adj_no2012"] = within_runner_adj(panel[panel.year != 2012],
                                                 "Within-runner, year-adjusted, excluding 2012")
results["within_no2012"] = within_runner(panel[panel.year != 2012],
                                         "Within-runner, excluding 2012 heat race")
# Sensitivity: require >=5 appearances
keys5 = panel.groupby("key").year.nunique()
results["within_k5"] = within_runner(panel[panel.key.isin(keys5[keys5 >= 5].index)],
                                     "Within-runner, runners with >=5 appearances")

# ---------------------------------------------------------------- ATTRITION
# Does returning depend on how the previous race went? If slower runners stop
# returning, the cross-sectional age gradient is flattened by survivorship.
panel_sorted = panel.sort_values(["key", "year"])
panel_sorted["next_year"] = panel_sorted.groupby("key")["year"].shift(-1)
panel_sorted["returned"] = panel_sorted["next_year"].notna()

# performance percentile within each race year and sex
panel_sorted["pctile"] = (panel_sorted.groupby(["year", "gender"])["finish"]
                          .rank(pct=True))
att = panel_sorted[panel_sorted.year < 2014].dropna(subset=["pctile"])
a = att.age.values - AGE_C
Xa = np.column_stack([np.ones(len(att)), a, att.pctile.values,
                      (att.gender.values == "F").astype(float)])
ya = att.returned.values.astype(float)
attrition = ols(ya, Xa, ["intercept", "age", "finish_percentile", "female"])
attrition.update(label="Linear probability model of returning to a later Boston Marathon",
                 model="OLS (HC1); outcome = runner appears in any later year")

# observed return rate by performance quartile
q = pd.qcut(att.pctile, 4, labels=["Q1 fastest", "Q2", "Q3", "Q4 slowest"])
ret_by_q = att.groupby(q, observed=True).returned.mean().round(4).to_dict()

# ---------------------------------------------------------------- OUTPUT
summary = dict(
    counts=dict(raw_records=int(N_RAW), wheelchair_removed=int(N_WHEEL),
                year2013_removed=int(N_2013), implausible_removed=int(N_IMPLAUS),
                analysis_records=int(len(df_perf)),
                panel_runners=int(N_PANEL_RUNNERS), panel_obs=int(N_PANEL_OBS),
                years=sorted(int(y) for y in df_perf.year.unique())),
    slopes={},
    models={k: {kk: vv for kk, vv in v.items()
                if kk in ("label", "model", "n", "n_groups", "r2", "terms")}
            for k, v in results.items()},
    attrition={kk: vv for kk, vv in attrition.items()
               if kk in ("label", "model", "n", "r2", "terms")},
    return_rate_by_quartile=ret_by_q,
    year_effects={str(k): round(v, 3) for k, v in sorted(YEAR_EFFECTS.items())},
)

for key in ("cross_all", "cross_panel", "within", "within_adj", "within_adj_m",
            "within_adj_f", "within_adj_no2012", "within_m", "within_f",
            "within_no2012", "within_k5"):
    summary["slopes"][key] = {f"age_{ag}": dict(zip(("slope_min_per_yr", "se"),
                                                    slope_at(results[key], ag)))
                              for ag in (35, 45, 55, 65)}

with open("/home/claude/boston_results.json", "w") as fh:
    json.dump(summary, fh, indent=2)

# also save the panel for figures
panel.to_csv("/home/claude/boston_panel.csv", index=False)
df_perf.to_csv("/home/claude/boston_analysis_sample.csv", index=False)

# ---------------------------------------------------------------- PRINT
c = summary["counts"]
print("=" * 76)
print(f"Records: {c['raw_records']:,} raw -> {c['analysis_records']:,} analysed")
print(f"  removed: {c['wheelchair_removed']:,} wheelchair, "
      f"{c['year2013_removed']:,} from 2013, {c['implausible_removed']:,} implausible")
print(f"Panel: {c['panel_runners']:,} runners, {c['panel_obs']:,} observations")
print("=" * 76)
for key in ("cross_all", "cross_panel", "within", "within_adj", "within_adj_m",
            "within_adj_f", "within_adj_no2012", "within_m", "within_f",
            "within_no2012", "within_k5"):
    r = results[key]
    print(f"\n{r['label']}   (n = {r['n']:,}"
          + (f", runners = {r['n_groups']:,}" if "n_groups" in r else "") + ")")
    for ag in (35, 45, 55, 65):
        b, se = slope_at(r, ag)
        print(f"    age {ag}: {b:+.3f} min/yr  (95% CI {b-1.96*se:+.3f} to {b+1.96*se:+.3f})")
print("\n" + "=" * 76)
print(attrition["label"], f"(n = {attrition['n']:,})")
for nm in ("age", "finish_percentile", "female"):
    t = attrition["terms"][nm]
    print(f"    {nm:<20} b = {t['b']:+.4f}  SE {t['se']:.4f}  p = {t['p']:.3g}")
print("  observed return rate by finish-time quartile:", ret_by_q)