{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "891d82df",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c8004118",
   "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.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": "83ff60e4",
   "metadata": {},
   "source": [
    "Load Cleaned Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "db836b46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X shape: (45133, 2315)\n",
      "y mean: 0.11488268007887799\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_pickle(\"../data/cleaned/NSDUH_2023_clean.pkl\")\n",
    "\n",
    "TARGET = \"IRAMDEYR\"\n",
    "WEIGHT = \"ANALWT2_C\"   # remove for feature selection\n",
    "\n",
    "y = df[TARGET]\n",
    "X = df.drop(columns=[TARGET, WEIGHT])  # DROP WEIGHT COLUMN\n",
    "\n",
    "print(\"X shape:\", X.shape)\n",
    "print(\"y mean:\", y.mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50c57b36",
   "metadata": {},
   "source": [
    "Train/Test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8b90f133",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train: (36106, 2315)  Test: (9027, 2315)\n"
     ]
    }
   ],
   "source": [
    "# numeric coercion\n",
    "X = X.apply(pd.to_numeric, errors=\"coerce\").fillna(0)\n",
    "\n",
    "# Split (same rule as all other feature-selection notebooks)\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": "5639efdd",
   "metadata": {},
   "source": [
    "Standardize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d856bfd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = StandardScaler()\n",
    "X_train_scaled = scaler.fit_transform(X_train)\n",
    "X_test_scaled = scaler.transform(X_test)\n",
    "\n",
    "feature_names = X.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05b4dfa3",
   "metadata": {},
   "source": [
    "Define a LASSO Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9a98fddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_lasso(C_value, label):\n",
    "    print(f\"\\nRunning LASSO ({label}) with C={C_value}\")\n",
    "    start = time.time()\n",
    "\n",
    "    model = LogisticRegression(\n",
    "        penalty=\"l1\",\n",
    "        solver=\"saga\",\n",
    "        C=C_value,\n",
    "        max_iter=3000,\n",
    "        n_jobs=-1\n",
    "    )\n",
    "\n",
    "    model.fit(X_train_scaled, y_train)\n",
    "\n",
    "    # Time\n",
    "    minutes = (time.time() - start) / 60\n",
    "    print(f\"Completed in {minutes:.2f} minutes\")\n",
    "\n",
    "    # evaluate on test\n",
    "    if label.lower().startswith(\"auc\"):\n",
    "        probs = model.predict_proba(X_test_scaled)[:, 1]\n",
    "        score = roc_auc_score(y_test, probs)\n",
    "        print(\"AUC:\", score)\n",
    "\n",
    "    else:\n",
    "        preds = model.predict(X_test_scaled)\n",
    "        score = recall_score(y_test, preds)\n",
    "        print(\"Recall:\", score)\n",
    "\n",
    "    # coefficient importance\n",
    "    coef = np.abs(model.coef_[0])\n",
    "    coef_series = pd.Series(coef, index=feature_names).sort_values(ascending=False)\n",
    "\n",
    "    top100 = coef_series.head(100).index.tolist()\n",
    "\n",
    "    return top100, coef_series"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e2b65b6",
   "metadata": {},
   "source": [
    "Run LASSO (AUC version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "05f983ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Running LASSO (AUC-Optimized) with C=1.0\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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed in 84.03 minutes\n",
      "AUC: 0.9988266432365432\n",
      "\n",
      "Saved lasso_auc_top100.csv\n"
     ]
    }
   ],
   "source": [
    "lasso_auc_top100, coef_auc = run_lasso(\n",
    "    C_value=1.0,    # stronger regularization → AUC-focused\n",
    "    label=\"AUC-Optimized\"\n",
    ")\n",
    "\n",
    "pd.Series(lasso_auc_top100).to_csv(\n",
    "    \"../results/lasso_auc_top100.csv\",\n",
    "    index=False\n",
    ")\n",
    "\n",
    "print(\"\\nSaved lasso_auc_top100.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a3d2459",
   "metadata": {},
   "source": [
    "Run LASSO (Recall version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d504c266",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Running LASSO (Recall-Optimized) with C=0.5\n",
      "Completed in 77.27 minutes\n",
      "Recall: 0.9816779170684667\n",
      "\n",
      "Saved lasso_recall_top100.csv\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": [
    "lasso_recall_top100, coef_recall = run_lasso(\n",
    "    C_value=0.5,     # weaker regularization → increases recall bias\n",
    "    label=\"Recall-Optimized\"\n",
    ")\n",
    "\n",
    "pd.Series(lasso_recall_top100).to_csv(\n",
    "    \"../results/lasso_recall_top100.csv\",\n",
    "    index=False\n",
    ")\n",
    "\n",
    "print(\"\\nSaved lasso_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
}