"""Lesson 12: censoring, prefix-only fate prediction, randomized effect."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from lifelines import KaplanMeierFitter, CoxPHFitter
from lifelines.utils import concordance_index
from sklearn.calibration import calibration_curve
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import brier_score_loss, log_loss, roc_auc_score
from sklearn.preprocessing import StandardScaler
from common import make_dataset


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    n_sessions = len(data["duration"])
    sessions = np.arange(n_sessions)
    first = np.array([np.flatnonzero(data["session"] == s)[0] for s in sessions])
    fit = data["train"][first]
    prefix = data["time"] <= .4
    # Fit EVERY feature transform using only training sessions' prefix frames.
    scaler = StandardScaler().fit(data["X"][prefix & data["train"]])
    pca = PCA(n_components=2).fit(scaler.transform(data["X"][prefix & data["train"]]))
    prefix_z = pca.transform(scaler.transform(data["X"][prefix]))
    prefix_session = data["session"][prefix]
    features = np.vstack([prefix_z[prefix_session == s].mean(0) for s in sessions])
    model = LogisticRegression(C=1, max_iter=1000).fit(features[fit], data["fate"][fit])
    prob = model.predict_proba(features[~fit])[:, 1]
    y = data["fate"][~fit]
    constant = np.full(len(y), data["fate"][fit].mean())
    table = pd.DataFrame({"duration": data["duration"], "event": data["event"],
                          "baseline": data["baseline"], "treatment": data["treatment"]})
    cox = CoxPHFitter(penalizer=.05).fit(table[fit], "duration", "event")
    risk = cox.predict_partial_hazard(table[~fit]).to_numpy()
    cindex = concordance_index(data["duration"][~fit], -risk, data["event"][~fit])
    treated = data["outcome"][data["treatment"] == 1]
    control = data["outcome"][data["treatment"] == 0]
    effect = float(treated.mean() - control.mean())
    se = float(np.sqrt(treated.var(ddof=1)/len(treated) + control.var(ddof=1)/len(control)))
    fig, axes = plt.subplots(1, 3, figsize=(12.5, 4))
    km = KaplanMeierFitter()
    km.fit(data["duration"], event_observed=data["event"], label="All sessions")
    km.plot_survival_function(ax=axes[0], ci_show=True)
    axes[0].set(title="Separate post-window follow-up", xlabel="Follow-up time", ylabel="S(t)")
    observed, predicted = calibration_curve(y, prob, n_bins=4, strategy="quantile")
    axes[1].plot(predicted, observed, "o-", label="Small held-out sample")
    axes[1].plot([0, 1], [0, 1], "k:")
    axes[1].set(xlabel="Predicted fate probability", ylabel="Observed fraction", title="Calibration diagnostic")
    axes[2].errorbar([effect], [0], xerr=1.96*se, fmt="o", capsize=5)
    axes[2].axvline(.8, linestyle=":", color="#dc5a32", label="Simulator ATE = 0.8")
    axes[2].axvline(0, color="gray", linewidth=.6)
    axes[2].set(yticks=[], xlabel="Mean outcome difference", title="Randomized effect\nApproximate 95% CI")
    axes[2].legend(fontsize=7)
    fig.tight_layout()
    fig.savefig(output_dir / "lesson12.png", dpi=160)
    plt.close(fig)
    return {"test_sessions": int((~fit).sum()), "fate_auc": float(roc_auc_score(y, prob)),
            "fate_brier": float(brier_score_loss(y, prob)),
            "baseline_brier": float(brier_score_loss(y, constant)),
            "fate_log_loss": float(log_loss(y, prob, labels=[0, 1])),
            "cox_test_concordance": float(cindex), "event_fraction": float(data["event"].mean()),
            "randomized_effect": effect, "effect_ci_low": effect-1.96*se,
            "effect_ci_high": effect+1.96*se, "simulator_true_ate": .8}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    args = parser.parse_args()
    print(json.dumps(run(args.output), indent=2, allow_nan=False))