{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d4854640",
   "metadata": {},
   "source": [
    "# Sample Code to Run SC-Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ded9e92",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0c40ad1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "\n",
    "seed = 0\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea4180a0",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9c9b2829",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   pid  event  time fac_grade fac_stage fac_ageGroup\n",
      "0    0      0    18      mode       T1c          80+\n",
      "1    1      0    23      mode      T1ab        75-79\n",
      "2    2      0    37      poor       T1c        75-79\n",
      "3    3      0    27      mode        T2        70-74\n",
      "4    4      0    42      mode       T1c        70-74\n"
     ]
    }
   ],
   "source": [
    "# Prerequisite: pip install SurvSet\n",
    "from SurvSet.data import SurvLoader\n",
    "\n",
    "dataset_name = \"prostateSurvival\"\n",
    "loader = SurvLoader()\n",
    "df, _ = loader.load_dataset(ds_name=dataset_name).values()\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4fd6e35",
   "metadata": {},
   "source": [
    "## Encode Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "985299bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   pid  event  time  fac_grade_poor  fac_stage_T1c  fac_stage_T2  \\\n",
      "0    0      0    18             0.0            1.0           0.0   \n",
      "1    1      0    23             0.0            0.0           0.0   \n",
      "2    2      0    37             1.0            1.0           0.0   \n",
      "3    3      0    27             0.0            0.0           1.0   \n",
      "4    4      0    42             0.0            1.0           0.0   \n",
      "\n",
      "   fac_ageGroup_70-74  fac_ageGroup_75-79  fac_ageGroup_80+  \n",
      "0                 0.0                 0.0               1.0  \n",
      "1                 0.0                 1.0               0.0  \n",
      "2                 0.0                 1.0               0.0  \n",
      "3                 1.0                 0.0               0.0  \n",
      "4                 1.0                 0.0               0.0  \n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "\n",
    "# rescale numerical features into [0, 1]\n",
    "cols_cont = [col for col in df.columns if col.startswith(\"num_\")]\n",
    "if len(cols_cont) > 0:\n",
    "    df[cols_cont] = df[cols_cont].fillna(df[cols_cont].median())\n",
    "    mm = MinMaxScaler()\n",
    "    df[cols_cont] = mm.fit_transform(df[cols_cont])\n",
    "\n",
    "# encode categorical columns\n",
    "cols_cat = []\n",
    "for col in df.columns:\n",
    "    if col.startswith(\"fac_\"):\n",
    "        cols_cat.append(col)\n",
    "df = pd.get_dummies(df, columns=cols_cat, drop_first=True, dtype=\"float32\")\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2832a09b",
   "metadata": {},
   "source": [
    "## Split Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "20fb7da2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(11435, 9) (2859, 9)\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "min_y = 0.0\n",
    "max_y = df[\"time\"].max()\n",
    "\n",
    "df_train, df_test = train_test_split(df, random_state=42, test_size=0.2)\n",
    "print(df_train.shape, df_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a38325cd",
   "metadata": {},
   "source": [
    "## Train with Neural Network (SC-Net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "01d02620",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cenreg.pytorch.copula_torch\n",
    "from cenreg.pytorch.datamodule import SurvDataModule\n",
    "from cenreg.pytorch.loss_cont import CopulaNegativeLogLikelihood\n",
    "from cenreg.pytorch.mlp import SMM_MultiHead\n",
    "\n",
    "# drop the columns that are not needed\n",
    "x_train = df_train.drop(columns=[\"pid\", \"time\", \"event\"]).values\n",
    "x_test = df_test.drop(columns=[\"pid\", \"time\", \"event\"]).values\n",
    "\n",
    "# preparation for PyTorch training\n",
    "sdm = SurvDataModule(128)\n",
    "train_dataloader = sdm.train_dataloader(x_train, df_train[\"time\"].values, df_train[\"event\"].values)\n",
    "model = SMM_MultiHead(x_train.shape[1], 1, 2, 16)\n",
    "\n",
    "# create a copula object\n",
    "copula_torch = cenreg.pytorch.copula_torch.create(\"independence\")\n",
    "# prepare the loss function\n",
    "loss_fn = CopulaNegativeLogLikelihood(copula_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "581e2dc4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train Loss: 0.31274175867438314\n",
      "Epoch 2, Train Loss: 0.1699182650488284\n",
      "Epoch 3, Train Loss: 0.14413190956951843\n",
      "Epoch 4, Train Loss: 0.17235662221080728\n",
      "Epoch 5, Train Loss: 0.1399189831275079\n",
      "Epoch 6, Train Loss: 0.15406331359926198\n",
      "Epoch 7, Train Loss: 0.15370798817732267\n",
      "Epoch 8, Train Loss: 0.1371321897332867\n",
      "Epoch 9, Train Loss: 0.13695469473799068\n",
      "Epoch 10, Train Loss: 0.14057688855876524\n"
     ]
    }
   ],
   "source": [
    "import cenreg.pytorch.utils\n",
    "\n",
    "optimizer_mlp = torch.optim.Adam(model.parameters(), lr=1.0)\n",
    "for epoch in range(10):\n",
    "    loss_sum = 0.0\n",
    "    for i, data in enumerate(train_dataloader):\n",
    "        x, t, e = data\n",
    "\n",
    "        optimizer_mlp.zero_grad()\n",
    "        t_grad = t.view(-1, 1).clone().detach().requires_grad_(True)\n",
    "        t_grad = cenreg.pytorch.utils.normalize_y(t_grad, min_y, max_y)\n",
    "        F_pred = model(x, t_grad)\n",
    "        loss = loss_fn.loss(F_pred, t_grad, e).mean()\n",
    "        loss_sum += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer_mlp.step()\n",
    "    print(f\"Epoch {epoch+1}, Train Loss: {loss_sum / len(train_dataloader)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15bda88f",
   "metadata": {},
   "source": [
    "## Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a93782c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F_pred tensor([[0.0000, 0.1252, 0.1932, 0.2826, 0.3730, 0.4326, 0.4713, 0.5055, 0.5388,\n",
      "         0.5716, 0.6039, 0.6352, 0.6655, 0.6944, 0.7219, 0.7478, 0.7721, 0.7946,\n",
      "         0.8155, 0.8347, 0.8523, 0.8682, 0.8827, 0.8958, 0.9076, 0.9182, 0.9276,\n",
      "         0.9361, 0.9436, 0.9503, 0.9562, 0.9614, 1.0000],\n",
      "        [0.0000, 0.0052, 0.0074, 0.0107, 0.0152, 0.0198, 0.0214, 0.0232, 0.0251,\n",
      "         0.0271, 0.0293, 0.0317, 0.0343, 0.0370, 0.0400, 0.0432, 0.0466, 0.0504,\n",
      "         0.0543, 0.0586, 0.0632, 0.0682, 0.0734, 0.0791, 0.0852, 0.0916, 0.0985,\n",
      "         0.1059, 0.1138, 0.1221, 0.1310, 0.1405, 1.0000]]) torch.Size([2859, 2, 33])\n"
     ]
    }
   ],
   "source": [
    "import cenreg.utils\n",
    "from cenreg.distribution.linear import LinearCDF\n",
    "\n",
    "num_bins = 32\n",
    "bins = torch.tensor(cenreg.utils.create_bins(max_y, 0.0, num_bins), dtype=torch.float32)\n",
    "\n",
    "test_dataloader = sdm.test_dataloader(x_test, df_test[\"time\"].values, df_test[\"event\"].values)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for i, data in enumerate(test_dataloader):\n",
    "        x, _, _ = data\n",
    "        list_pred = []\n",
    "        for t in bins:\n",
    "            t = cenreg.pytorch.utils.normalize_y(t, min_y, max_y)\n",
    "            t_expand = t.expand(x.shape[0], 1)\n",
    "            pred = model(x, t_expand)\n",
    "            list_pred.append(pred)\n",
    "        F_pred = torch.stack(list_pred, dim=2)\n",
    "        F_pred[:,:,0] = 0.0\n",
    "        F_pred[:,:,-1] = 1.0\n",
    "        break  # assuming that the first batch contains all the test data\n",
    "print(\"F_pred\", F_pred[0,:,:], F_pred.shape)\n",
    "\n",
    "bins_np = bins.detach().cpu().numpy()\n",
    "F_pred_np = F_pred.detach().cpu().numpy()\n",
    "list_t_dist = []\n",
    "for k in range(2):\n",
    "    list_t_dist.append(LinearCDF(bins_np, F_pred_np[:,k,:], apply_cumsum=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d61d1a0",
   "metadata": {},
   "source": [
    "## Convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "007e3853",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cenreg.model.copula_np\n",
    "import cenreg.model.F2cjd_np\n",
    "\n",
    "num_risks = 2\n",
    "list_Fpred = []\n",
    "for k, d in enumerate(list_t_dist):\n",
    "    F_pred = d.cdf(bins_np[1:-1], add_edges=True)\n",
    "    if F_pred.ndim == 1:\n",
    "        F_pred = F_pred.reshape(1, -1)\n",
    "    list_Fpred.append(F_pred)\n",
    "F_pred = np.stack(list_Fpred, axis=1)\n",
    "\n",
    "copula_np = cenreg.model.copula_np.IndependenceCopula()\n",
    "list_K = list(range(num_risks))\n",
    "w = [0.0] + [1 / i for i in range(1, num_risks + 1)]\n",
    "jd_pred = np.zeros((F_pred.shape[0], num_risks, bins_np.shape[0] - 1))\n",
    "for k in list_K:\n",
    "    for t in range(F_pred.shape[2] - 1):\n",
    "        temp = F_pred[:, :, t : t + 2]\n",
    "        ones = np.ones((temp.shape[0], temp.shape[1], 1))\n",
    "        temp = np.concatenate([temp, ones], axis=2)\n",
    "        jd_pred[:, k, t] = cenreg.model.F2cjd_np.convert(\n",
    "            temp, copula_np, w, num_risks, k, list_K\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e5257c9",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "5183f4cd",
   "metadata": {},
   "source": [
    "## Evaluate Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b050886e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CJD-Brier 0.9645153392245996\n",
      "CJD-Logarithmic 3.5418869960575394\n",
      "CJD-KS 0.34061252873258013\n",
      "NLL-SC 4.885996452186602\n",
      "Cen-log 0.3756013711177897\n",
      "D-calibration 0.0005720768376036918\n",
      "KM-calibration 0.07679972\n"
     ]
    }
   ],
   "source": [
    "import cenreg.metric.cdf\n",
    "import cenreg.metric.cjd\n",
    "import cenreg.metric.quantile\n",
    "import cenreg.model.copula_np\n",
    "\n",
    "cjd_pred = jd_pred.reshape(-1, num_bins * 2)\n",
    "observed_times = df_test['time'].values\n",
    "events = df_test[\"event\"].astype(bool).values\n",
    "\n",
    "# Compute Brier score on CJD representation\n",
    "cjd_brier = cenreg.metric.cjd.brier(observed_times, events, 2, cjd_pred, bins_np)\n",
    "print(\"CJD-Brier\", cjd_brier.mean())\n",
    "\n",
    "# Compute Logarithmic score on CDF representation\n",
    "cjd_logarithmic = cenreg.metric.cjd.negative_loglikelihood(observed_times, events, cjd_pred, bins_np)\n",
    "print(\"CJD-Logarithmic\", cjd_logarithmic.mean())\n",
    "\n",
    "# Compute KS calibration error on CDF representation\n",
    "cjd_ks = cenreg.metric.cjd.kolmogorov_smirnov_calibration_error(observed_times, events, cjd_pred, bins_np)\n",
    "print(\"CJD-KS\", cjd_ks)\n",
    "\n",
    "# Compute NLL-SC metric\n",
    "copula_np = cenreg.model.copula_np.IndependenceCopula()\n",
    "survival_copula_np = cenreg.model.copula_np.SurvivalCopula(copula_np)\n",
    "nll_sc = cenreg.metric.cdf.nll_sc(list_t_dist, observed_times, events, survival_copula_np)\n",
    "print(\"NLL-SC\", nll_sc.mean())\n",
    "\n",
    "# Compute cen-log metric\n",
    "nll = cenreg.metric.cdf.negative_loglikelihood(list_t_dist[1], observed_times, events)\n",
    "print(\"Cen-log\", nll.mean())\n",
    "\n",
    "# Compute D-calibration\n",
    "quantiles_cal = np.linspace(0.0, 1.0, 11)\n",
    "dcal = cenreg.metric.quantile.d_calibration(list_t_dist[1], observed_times, events, quantiles_cal)\n",
    "print(\"D-calibration\", dcal)\n",
    "\n",
    "# Compute KM-calibration\n",
    "kmcal = cenreg.metric.cdf.km_calibration(list_t_dist[1], observed_times, events, bins_np)\n",
    "print(\"KM-calibration\", kmcal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "543f50f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20112 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "x = list_t_dist[1].boundaries\n",
    "y = 1.0 - list_t_dist[1].cdf(x)\n",
    "\n",
    "plt.style.use(\"tableau-colorblind10\")\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.title(dataset_name)\n",
    "plt.xlabel(\"Time\")\n",
    "plt.ylabel(\"Survival probability\")\n",
    "plt.ylim(0.0, 1.05)\n",
    "plt.plot(x[:-1], y[50,:-1], marker='o')  # Plot the 50th survival function\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cenreg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
