"""Survival analysis of project lifecycles.

Among multi-filing projects (tracked across >=2 filings), we measure the span
from first to last appearance. Event = the project reaches 'completed' status;
projects last observed as planned/in-progress are right-censored. Kaplan-Meier
curves by project family and a Cox proportional-hazards model.
"""
from __future__ import annotations

import warnings
import numpy as np
import pandas as pd
import statsmodels.api as sm
from statsmodels.duration.survfunc import SurvfuncRight
from statsmodels.duration.hazard_regression import PHReg

from common import DATA, set_style, savefig, write_table, save_results
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")
set_style()

p = pd.read_parquet(DATA / "projects.parquet")
d = p[(p["n_filings"] >= 2) & (p["duration_days"] > 0)].copy()
d["months"] = d["duration_days"] / 30.44
d["event"] = (d["status"] == "completed").astype(int)
d["log_amount_f"] = np.log(d["total_amount_usd"].fillna(d["total_amount_usd"].median()).clip(lower=1))
d["is_ai_i"] = d["is_ai"].astype(int)

# ---- Kaplan-Meier by project family --------------------------------------
FAMILIES = {
    "M&A integration": ["ma_integration"],
    "Capex / construction": ["capex_program", "plant_construction", "infrastructure"],
    "R&D / product": ["rd_program", "product_launch", "drug_pipeline"],
    "Digital / AI": ["ai_initiative", "digital_transformation", "cloud_migration",
                     "data_center", "automation", "technology_deployment"],
}
fig, ax = plt.subplots(figsize=(7.4, 4.4))
median_surv = {}
for (lbl, types), color in zip(FAMILIES.items(), ["#1b3a5b", "#c44e52", "#2e8b57", "#8172b3"]):
    sub = d[d["project_type"].isin(types)]
    if len(sub) < 20:
        continue
    sf = SurvfuncRight(sub["months"], sub["event"])
    ax.step(sf.surv_times, sf.surv_prob, where="post", label=f"{lbl} (n={len(sub)})", color=color, lw=1.8)
    # median survival
    below = sf.surv_times[sf.surv_prob <= 0.5]
    median_surv[lbl] = float(below[0]) if len(below) else float("nan")
ax.set_xlabel("Months since first appearance in filings")
ax.set_ylabel("Share not yet completed")
ax.set_title("Kaplan-Meier project-survival curves by family")
ax.legend(fontsize=8)
ax.set_xlim(0, d["months"].quantile(0.98))
savefig(fig, "fig_survival_km.pdf")

# ---- Cox proportional hazards --------------------------------------------
base = "ma_integration"
X = pd.get_dummies(d["project_type"].astype(str), prefix="pt", drop_first=False, dtype=float)
X = X.drop(columns=[f"pt_{base}"], errors="ignore")          # baseline category
X = X.loc[:, X.std() > 0]                                     # drop all-zero/constant dummies
X["is_ai"] = d["is_ai_i"].values
X["log_amount"] = d["log_amount_f"].values
X["n_tech"] = d["n_tech"].values
X = X.loc[:, X.std() > 0]
cox = PHReg(d["months"].values, X.values.astype(float), status=d["event"].values).fit()
params = pd.Series(cox.params, index=X.columns)
bse = pd.Series(cox.bse, index=X.columns)
pvals = pd.Series(cox.pvalues, index=X.columns)

SHOW = ["is_ai", "log_amount", "n_tech"]
lt = [r"\begin{tabular}{lcc}", r"\toprule",
      r"Covariate & Hazard ratio & $p$-value \\", r"\midrule"]
for k in SHOW:
    hr = np.exp(params[k]); pv = pvals[k]
    star = "***" if pv < .01 else "**" if pv < .05 else "*" if pv < .1 else ""
    lt.append(f"{k.replace('_',' ')} & {hr:.3f}{star} & {pv:.3f} \\\\")
lt += [r"\midrule", r"Project-type controls & Yes & \\",
       f"Projects (events) & {int(d['event'].sum()):,} of {len(d):,} & \\\\",
       r"\bottomrule", r"\end{tabular}"]
write_table("tab_cox", "\n".join(lt))

save_results("survival", {
    "n": int(len(d)), "events": int(d["event"].sum()),
    "median_survival_months": median_surv,
    "cox_is_ai_hr": float(np.exp(params["is_ai"])), "cox_is_ai_p": float(pvals["is_ai"]),
    "cox_logamount_hr": float(np.exp(params["log_amount"])), "cox_logamount_p": float(pvals["log_amount"]),
})
print(f"survival done | n={len(d):,}, events={int(d['event'].sum()):,}")