{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cfc14fde",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "225d950b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import time\n",
    "\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.model_selection import train_test_split\n",
    "from sklearn.metrics import roc_auc_score, recall_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34645c0b",
   "metadata": {},
   "source": [
    "Load the cleaned dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "60e81251",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X: (45133, 2315)  y: (45133,)\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_pickle(\"../data/cleaned/NSDUH_2023_clean.pkl\")\n",
    "\n",
    "TARGET = \"IRAMDEYR\"\n",
    "WEIGHT = \"ANALWT2_C\"   # we drop it immediately\n",
    "\n",
    "y = df[TARGET]\n",
    "X = df.drop(columns=[TARGET, WEIGHT])\n",
    "\n",
    "print(\"X:\", X.shape, \" y:\", y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28d5815c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train: (36106, 2315)  Test: (9027, 2315)\n"
     ]
    }
   ],
   "source": [
    "# Train Test Split\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, y,\n",
    "    test_size=0.2,\n",
    "    stratify=y,\n",
    "    random_state=42\n",
    ")\n",
    "\n",
    "print(\"Train:\", X_train.shape, \" Test:\", X_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8e732ba",
   "metadata": {},
   "source": [
    "ElasticNet model (AUC version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "54fe3002",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Running ElasticNet (AUC)\n",
      "Completed in 15.58 minutes\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\agila\\Documents\\STM\\Predicting-Past-Year-Major-Depressive-Episode\\venv\\Lib\\site-packages\\sklearn\\linear_model\\_sag.py:348: ConvergenceWarning: The max_iter was reached which means the coef_ did not converge\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "elasticnet_auc = Pipeline([\n",
    "    ('imputer', SimpleImputer(strategy='median')),\n",
    "    ('scaler', StandardScaler()),\n",
    "    ('clf', LogisticRegression(\n",
    "        penalty='elasticnet',\n",
    "        solver='saga',\n",
    "        C=1.0,\n",
    "        l1_ratio=0.5,\n",
    "        max_iter=500,\n",
    "        random_state=42\n",
    "    ))\n",
    "])\n",
    "\n",
    "print(\"\\nRunning ElasticNet (AUC)\")\n",
    "start = time.time()\n",
    "\n",
    "elasticnet_auc.fit(X_train, y_train)\n",
    "\n",
    "print(f\"Completed in {(time.time() - start)/60:.2f} minutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da9a4ffa",
   "metadata": {},
   "source": [
    "Evaluate on test set (AUC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bbb02027",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AUC Score: 0.9990896286703606\n"
     ]
    }
   ],
   "source": [
    "probs = elasticnet_auc.predict_proba(X_test)[:, 1]\n",
    "auc_score = roc_auc_score(y_test, probs)\n",
    "\n",
    "print(\"AUC Score:\", auc_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adfc0b6c",
   "metadata": {},
   "source": [
    "Extract Top 100 features (AUC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b275de3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved elasticnet_auc_top100.csv\n"
     ]
    }
   ],
   "source": [
    "coefs = np.abs(elasticnet_auc.named_steps['clf'].coef_[0])\n",
    "importances = pd.Series(coefs, index=X.columns)\n",
    "\n",
    "top100_auc = importances.sort_values(ascending=False).head(100).index.tolist()\n",
    "\n",
    "pd.Series(top100_auc).to_csv(\"../results/elasticnet_auc_top100.csv\", index=False)\n",
    "\n",
    "print(\"Saved elasticnet_auc_top100.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d347c4de",
   "metadata": {},
   "source": [
    "ElasticNet model (Recall version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c0e63067",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Running ElasticNet (Recall)\n",
      "Completed in 16.57 minutes\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\agila\\Documents\\STM\\Predicting-Past-Year-Major-Depressive-Episode\\venv\\Lib\\site-packages\\sklearn\\linear_model\\_sag.py:348: ConvergenceWarning: The max_iter was reached which means the coef_ did not converge\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "elasticnet_recall = Pipeline([\n",
    "    ('imputer', SimpleImputer(strategy='median')),\n",
    "    ('scaler', StandardScaler()),\n",
    "    ('clf', LogisticRegression(\n",
    "        penalty='elasticnet',\n",
    "        solver='saga',\n",
    "        C=0.5,\n",
    "        l1_ratio=0.5,\n",
    "        max_iter=500,\n",
    "        random_state=42\n",
    "    ))\n",
    "])\n",
    "\n",
    "print(\"\\nRunning ElasticNet (Recall)\")\n",
    "start = time.time()\n",
    "\n",
    "elasticnet_recall.fit(X_train, y_train)\n",
    "\n",
    "print(f\"Completed in {(time.time() - start)/60:.2f} minutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20519e69",
   "metadata": {},
   "source": [
    "Evaluate Recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2384f767",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall Score: 0.9537126325940212\n"
     ]
    }
   ],
   "source": [
    "pred = elasticnet_recall.predict(X_test)\n",
    "rec = recall_score(y_test, pred)\n",
    "\n",
    "print(\"Recall Score:\", rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a04f6a6",
   "metadata": {},
   "source": [
    "Extract Top 100 Recall features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21d80dd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved elasticnet_recall_top100.csv\n"
     ]
    }
   ],
   "source": [
    "coefs_recall = np.abs(elasticnet_recall.named_steps['clf'].coef_[0])\n",
    "importances_recall = pd.Series(coefs_recall, index=X.columns)\n",
    "\n",
    "top100_recall = importances_recall.sort_values(ascending=False).head(100).index.tolist()\n",
    "\n",
    "pd.Series(top100_recall).to_csv(\"../results/elasticnet_recall_top100.csv\", index=False)\n",
    "\n",
    "print(\"Saved elasticnet_recall_top100.csv\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.13.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}