{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c467f8df",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "030f5163",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import time\n",
    "from sklearn.model_selection import train_test_split, GridSearchCV, StratifiedKFold\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import (\n",
    "    roc_auc_score, recall_score, precision_score, accuracy_score, \n",
    "    confusion_matrix, fbeta_score, precision_recall_curve, roc_curve\n",
    ")\n",
    "from sklearn.dummy import DummyClassifier\n",
    "from sklearn.utils import resample\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import cross_validate\n",
    "from sklearn.metrics import make_scorer\n",
    "import os\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.preprocessing import OneHotEncoder\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5ba3077",
   "metadata": {},
   "source": [
    "Load tiers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "babee804",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved: C:\\Users\\agila\\Documents\\STM\\Predicting-Past-Year-Major-Depressive-Episode\\tiers\\tier_1_basic_features.csv\n",
      "Saved: C:\\Users\\agila\\Documents\\STM\\Predicting-Past-Year-Major-Depressive-Episode\\tiers\\tier_2_clinical_features.csv\n",
      "Saved: C:\\Users\\agila\\Documents\\STM\\Predicting-Past-Year-Major-Depressive-Episode\\tiers\\tier_3_personalized_features.csv\n"
     ]
    }
   ],
   "source": [
    "from tiering_preparation import TIERS\n",
    "\n",
    "tiers = {\n",
    "    \"Tier_1_Basic\": TIERS[\"Tier_1_Basic\"],\n",
    "    \"Tier_2_Clinical\": TIERS[\"Tier_2_Clinical\"],\n",
    "    \"Tier_3_Personalized\": TIERS[\"Tier_3_Personalized\"],\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fa8e944",
   "metadata": {},
   "source": [
    "Load dataset, drop weight from X, but SAVE w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6d728bf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_pickle(\"../data/prepared/NSDUH_2023_prepared.pkl\")  \n",
    "\n",
    "TARGET = \"IRAMDEYR\"\n",
    "WEIGHT = \"ANALWT2_C\"\n",
    "\n",
    "y = df[TARGET]\n",
    "w = df[WEIGHT]\n",
    "X = df.drop(columns=[TARGET, WEIGHT])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a51fcfb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "categorical_features = X.select_dtypes(\n",
    "    include=[\"object\", \"category\"]\n",
    ").columns.tolist()\n",
    "\n",
    "numeric_features = X.select_dtypes(\n",
    "    include=[\"int64\", \"float64\"]\n",
    ").columns.tolist()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a84949a",
   "metadata": {},
   "source": [
    "Split train/test (consistent with feature selection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "739bca76",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test, w_train, w_test = train_test_split(\n",
    "    X, y, w,\n",
    "    test_size=0.2,\n",
    "    stratify=y,\n",
    "    random_state=42\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60d4bb0e",
   "metadata": {},
   "source": [
    "Define metric helper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "75f72c77",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(y_true, preds, probs):\n",
    "    tn, fp, fn, tp = confusion_matrix(y_true, preds).ravel()\n",
    "\n",
    "    specificity = tn / (tn + fp)\n",
    "    sensitivity = recall_score(y_true, preds)  # same as recall\n",
    "    ppv = precision_score(y_true, preds)\n",
    "    # npv = tn / (tn + fn)\n",
    "    npv = tn / (tn + fn) if (tn + fn) > 0 else 0\n",
    "\n",
    "    return {\n",
    "        \"Accuracy\": accuracy_score(y_true, preds),\n",
    "        \"Sensitivity\": sensitivity,\n",
    "        \"Specificity\": specificity,\n",
    "        \"PPV\": ppv,\n",
    "        \"NPV\": npv,\n",
    "        \"AUC\": roc_auc_score(y_true, probs),\n",
    "        \"F1\": fbeta_score(y_true, preds, beta=1),\n",
    "        \"F2\": fbeta_score(y_true, preds, beta=2),\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abf73828",
   "metadata": {},
   "source": [
    "CV metric helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fcf9ee4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def specificity_cv(y_true, y_pred):\n",
    "    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()\n",
    "    return tn / (tn + fp) if (tn + fp) > 0 else 0\n",
    "\n",
    "def npv_cv(y_true, y_pred):\n",
    "    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()\n",
    "    return tn / (tn + fn) if (tn + fn) > 0 else 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b51469ec",
   "metadata": {},
   "source": [
    "CV scoring dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f7eb338a",
   "metadata": {},
   "outputs": [],
   "source": [
    "cv_scoring = {\n",
    "    \"accuracy\": make_scorer(accuracy_score),\n",
    "    \"sensitivity\": make_scorer(recall_score),\n",
    "    \"specificity\": make_scorer(specificity_cv),\n",
    "    \"ppv\": make_scorer(precision_score, zero_division=0),\n",
    "    \"npv\": make_scorer(npv_cv),\n",
    "    \"f1\": make_scorer(fbeta_score, beta=1),\n",
    "    \"f2\": make_scorer(fbeta_score, beta=2),\n",
    "    \"auc\": \"roc_auc\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "391c4136",
   "metadata": {},
   "source": [
    "Logistic Regression Pipeline w/ Weighted Fitting + GridSearchCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "966b088c",
   "metadata": {},
   "outputs": [],
   "source": [
    "param_grid = {\n",
    "    \"clf__C\": [0.01, 0.1, 1, 10],\n",
    "    \"clf__class_weight\": [None, \"balanced\"],\n",
    "}\n",
    "\n",
    "cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n",
    "\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "\n",
    "def build_model(numeric_features, categorical_features):\n",
    "\n",
    "    numeric_transformer = Pipeline([\n",
    "        (\"imputer\", SimpleImputer(strategy=\"median\")),\n",
    "        (\"scaler\", StandardScaler())\n",
    "    ])\n",
    "\n",
    "    categorical_transformer = Pipeline([\n",
    "        (\"imputer\", SimpleImputer(strategy=\"most_frequent\")),\n",
    "        (\"onehot\", OneHotEncoder(\n",
    "            handle_unknown=\"ignore\",\n",
    "            sparse_output=False\n",
    "        ))\n",
    "    ])\n",
    "\n",
    "    preprocessor = ColumnTransformer(\n",
    "        transformers=[\n",
    "            (\"num\", numeric_transformer, numeric_features),\n",
    "            (\"cat\", categorical_transformer, categorical_features),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    pipe = Pipeline([\n",
    "        (\"preprocessor\", preprocessor),\n",
    "        (\"clf\", LogisticRegression(\n",
    "            max_iter=3000,\n",
    "            penalty=\"l2\",\n",
    "            solver=\"lbfgs\"\n",
    "        ))\n",
    "    ])\n",
    "\n",
    "    model = GridSearchCV(\n",
    "        estimator=pipe,\n",
    "        param_grid=param_grid,\n",
    "        scoring=\"recall\",\n",
    "        cv=cv,\n",
    "        n_jobs=-1\n",
    "    )\n",
    "\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff87b9f7",
   "metadata": {},
   "source": [
    "Bootstrapped Confidence Intervals (ROC + PR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0517b4e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bootstrap_auc_ci(model, X, y, n_boot=1000):\n",
    "    aucs = []\n",
    "\n",
    "    for _ in range(n_boot):\n",
    "        X_bs, y_bs = resample(X, y, replace=True)\n",
    "        probs = model.predict_proba(X_bs)[:, 1]\n",
    "        aucs.append(roc_auc_score(y_bs, probs))\n",
    "\n",
    "    return np.percentile(aucs, [2.5, 97.5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d84eeafd",
   "metadata": {},
   "source": [
    "Baseline Model (Dummy Classifier)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6a8839f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\agila\\Documents\\STM\\Predicting-Past-Year-Major-Depressive-Episode\\venv\\Lib\\site-packages\\sklearn\\metrics\\_classification.py:1731: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 due to no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", result.shape[0])\n"
     ]
    }
   ],
   "source": [
    "dummy = DummyClassifier(strategy=\"most_frequent\")\n",
    "dummy.fit(X_train, y_train)\n",
    "dummy_preds = dummy.predict(X_test)\n",
    "dummy_probs = np.zeros_like(dummy_preds)\n",
    "\n",
    "dummy_metrics = evaluate_model(y_test, dummy_preds, dummy_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53186065",
   "metadata": {},
   "source": [
    "Threshold tuning helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b8d77cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_at_threshold(y_true, probs, threshold):\n",
    "    preds = (probs >= threshold).astype(int)\n",
    "    return evaluate_model(y_true, preds, probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7515b795",
   "metadata": {},
   "source": [
    "Train + Evaluate Logistic Regression on Each tier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1cfd0e26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----- Training Logistic Regression on Tier_1_Basic -----\n",
      "Saved fold table to: ../results/logreg_5fold_threshold0.50_by_tier.csv\n",
      "\n",
      "----- Training Logistic Regression on Tier_2_Clinical -----\n",
      "Saved fold table to: ../results/logreg_5fold_threshold0.50_by_tier.csv\n",
      "\n",
      "----- Training Logistic Regression on Tier_3_Personalized -----\n",
      "Saved fold table to: ../results/logreg_5fold_threshold0.50_by_tier.csv\n"
     ]
    }
   ],
   "source": [
    "results = []\n",
    "\n",
    "for tier_name, features in tiers.items():\n",
    "\n",
    "    print(f\"\\n----- Training Logistic Regression on {tier_name} -----\")\n",
    "\n",
    "    Xtr = X_train[features]\n",
    "    Xte = X_test[features]\n",
    "\n",
    "    # Identify feature types for THIS tier only\n",
    "    categorical_features = Xtr.select_dtypes(\n",
    "        include=[\"object\", \"category\"]\n",
    "    ).columns.tolist()\n",
    "\n",
    "    numeric_features = Xtr.select_dtypes(\n",
    "        include=[\"int64\", \"float64\"]\n",
    "    ).columns.tolist()\n",
    "\n",
    "    model = build_model(\n",
    "        numeric_features=numeric_features,\n",
    "        categorical_features=categorical_features\n",
    "    )\n",
    "\n",
    "\n",
    "    # Fit WITH survey weights\n",
    "    model.fit(Xtr, y_train, clf__sample_weight=w_train)\n",
    "    cv_mean_recall = model.best_score_\n",
    "    cv_std_recall = model.cv_results_[\"std_test_score\"][model.best_index_]\n",
    "\n",
    "    cv_scores = {k: [] for k in cv_scoring.keys()}\n",
    "\n",
    "    for train_idx, val_idx in cv.split(Xtr, y_train):\n",
    "        X_cv_tr, X_cv_val = Xtr.iloc[train_idx], Xtr.iloc[val_idx]\n",
    "        y_cv_tr, y_cv_val = y_train.iloc[train_idx], y_train.iloc[val_idx]\n",
    "        w_cv_tr = w_train.iloc[train_idx]\n",
    "\n",
    "        est = model.best_estimator_\n",
    "        est.fit(X_cv_tr, y_cv_tr, clf__sample_weight=w_cv_tr)\n",
    "\n",
    "        preds = est.predict(X_cv_val)\n",
    "        probs = est.predict_proba(X_cv_val)[:, 1]\n",
    "\n",
    "        # compute metrics fold-by-fold\n",
    "        fold_metrics = {\n",
    "            \"accuracy\": accuracy_score(y_cv_val, preds),\n",
    "            \"sensitivity\": recall_score(y_cv_val, preds),\n",
    "            \"specificity\": specificity_cv(y_cv_val, preds),\n",
    "            \"ppv\": precision_score(y_cv_val, preds, zero_division=0),\n",
    "            \"npv\": npv_cv(y_cv_val, preds),\n",
    "            \"f1\": fbeta_score(y_cv_val, preds, beta=1),\n",
    "            \"f2\": fbeta_score(y_cv_val, preds, beta=2),\n",
    "            \"auc\": roc_auc_score(y_cv_val, probs),\n",
    "        }\n",
    "\n",
    "        for k, v in fold_metrics.items():\n",
    "            cv_scores[k].append(v)\n",
    "\n",
    "    # summarize CV metrics\n",
    "    cv_summary = {}\n",
    "    for metric, values in cv_scores.items():\n",
    "        cv_summary[f\"CV_{metric}_mean\"] = np.mean(values)\n",
    "        cv_summary[f\"CV_{metric}_std\"] = np.std(values)\n",
    "\n",
    "\n",
    "    metrics_order = [\n",
    "    \"accuracy\", \"sensitivity\", \"specificity\",\n",
    "    \"ppv\", \"npv\", \"auc\", \"f1\", \"f2\"\n",
    "    ]\n",
    "\n",
    "    fold_df = pd.DataFrame({\n",
    "        f\"Fold {i+1}\": [cv_scores[m][i] for m in metrics_order]\n",
    "        for i in range(len(next(iter(cv_scores.values()))))\n",
    "    }, index=metrics_order)\n",
    "\n",
    "    fold_df[\"Mean\"] = fold_df.mean(axis=1)\n",
    "    fold_df[\"Std Dev\"] = fold_df.std(axis=1)\n",
    "\n",
    "    # Add metadata\n",
    "    fold_df.insert(0, \"Tier\", tier_name)\n",
    "    fold_df.insert(1, \"Model\", \"LogisticRegression\")\n",
    "\n",
    "    # Reset index so Metric is a column\n",
    "    fold_df = fold_df.reset_index().rename(columns={\"index\": \"Metric\"})\n",
    "\n",
    "    # (optional but recommended) save per-tier fold table here\n",
    "    output_csv = \"../results/logreg_5fold_threshold0.50_by_tier.csv\"\n",
    "\n",
    "    if os.path.exists(output_csv):\n",
    "        fold_df.to_csv(output_csv, mode=\"a\", header=False, index=False)\n",
    "    else:\n",
    "        fold_df.to_csv(output_csv, index=False)\n",
    "\n",
    "    print(f\"Saved fold table to: {output_csv}\")\n",
    "\n",
    "\n",
    "    # Probabilities (computed once)\n",
    "    probs = model.predict_proba(Xte)[:, 1]\n",
    "\n",
    "    # ---- Threshold sweep ----\n",
    "    thresholds = np.arange(0.1, 0.51, 0.05)\n",
    "    rows = []\n",
    "\n",
    "    for t in thresholds:\n",
    "        m = evaluate_at_threshold(y_test, probs, threshold=t)\n",
    "        m[\"Threshold\"] = t\n",
    "        rows.append(m)\n",
    "\n",
    "    threshold_df = pd.DataFrame(rows)\n",
    "\n",
    "    '''# Select screening threshold (Recall ≥ 0.90)\n",
    "    screening_candidates = threshold_df[\n",
    "    (threshold_df[\"Sensitivity\"] >= 0.90) &\n",
    "    (threshold_df[\"NPV\"] >= 0.95)\n",
    "    ]\n",
    "\n",
    "    screening_row = screening_candidates.sort_values(\n",
    "    by=[\"Specificity\", \"F2\"],\n",
    "    ascending=[False, False]\n",
    "    ).iloc[0]\n",
    "\n",
    "    # Sanity check print\n",
    "    print(\n",
    "        f\"Selected screening threshold for {tier_name}: \"\n",
    "        f\"{screening_row['Threshold']:.2f} \"\n",
    "        f\"(Sens={screening_row['Sensitivity']:.2f}, \"\n",
    "        f\"Spec={screening_row['Specificity']:.2f}, \"\n",
    "        f\"NPV={screening_row['NPV']:.2f})\"\n",
    "    )'''\n",
    "\n",
    "    # ---- Bootstrapped CI (AUC is threshold-independent) ----\n",
    "    ci_low, ci_high = bootstrap_auc_ci(model, Xte, y_test)\n",
    "\n",
    "    # ---- Default threshold metrics (0.5) ----\n",
    "    default_metrics = evaluate_at_threshold(y_test, probs, threshold=0.5)\n",
    "    default_metrics.update({\n",
    "        \"Tier\": tier_name,\n",
    "        \"Threshold\": 0.5,\n",
    "        \"Best_Params\": model.best_params_,\n",
    "        **cv_summary,\n",
    "        \"AUC_CI_Low\": ci_low,\n",
    "        \"AUC_CI_High\": ci_high,\n",
    "        \"Mode\": \"Default\"\n",
    "    })\n",
    "\n",
    "    '''# ---- Screening threshold metrics ----\n",
    "    screening_metrics = screening_row.to_dict()\n",
    "    screening_metrics.update({\n",
    "        \"Tier\": tier_name,\n",
    "        \"Best_Params\": model.best_params_,\n",
    "        **cv_summary,                 # ← ALL CV mean/std metrics\n",
    "        \"AUC_CI_Low\": ci_low,\n",
    "        \"AUC_CI_High\": ci_high,\n",
    "        \"Mode\": \"Screening\"\n",
    "    })'''\n",
    "\n",
    "    results.append(default_metrics)\n",
    "    # results.append(screening_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83b9aad4",
   "metadata": {},
   "source": [
    "Convert to DataFrame + Export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "46eefea0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                  Tier     Mode  Threshold  Sensitivity       NPV       AUC\n",
      "0         Tier_1_Basic  Default        0.5     0.764765  0.943414  0.701264\n",
      "1      Tier_2_Clinical  Default        0.5     0.893894  0.983286  0.922726\n",
      "2  Tier_3_Personalized  Default        0.5     0.905906  0.985301  0.930458\n"
     ]
    }
   ],
   "source": [
    "# Convert to DataFrame + Export\n",
    "results_df = pd.DataFrame(results)\n",
    "#results_df.to_csv(\"../results/logreg_tier_performance.csv\", index=False)\n",
    "\n",
    "# 🔍 Sanity check view (THIS LINE)\n",
    "print(results_df[[\"Tier\", \"Mode\", \"Threshold\", \"Sensitivity\", \"NPV\", \"AUC\"]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}