{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9b5a883f",
   "metadata": {},
   "source": [
    "Importing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b6114b26",
   "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.feature_selection import RFE\n",
    "from sklearn.metrics import roc_auc_score, recall_score\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be202f8b",
   "metadata": {},
   "source": [
    "Loading data and define X / y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8bf27a52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X shape: (45133, 2315)\n",
      "y distribution:\n",
      " IRAMDEYR\n",
      "0.0    39948\n",
      "1.0     5185\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Load cleaned dataset\n",
    "df = pd.read_pickle(\"../data/cleaned/NSDUH_2023_clean.pkl\")\n",
    "\n",
    "TARGET = \"IRAMDEYR\"\n",
    "\n",
    "# If ANALWT2_C is present, drop it from features (RFE doesn't use weights)\n",
    "cols_to_drop = [TARGET]\n",
    "if \"ANALWT2_C\" in df.columns:\n",
    "    cols_to_drop.append(\"ANALWT2_C\")\n",
    "\n",
    "X = df.drop(columns=cols_to_drop)\n",
    "y = df[TARGET]\n",
    "\n",
    "print(\"X shape:\", X.shape)\n",
    "print(\"y distribution:\\n\", y.value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "307d95d3",
   "metadata": {},
   "source": [
    "Ensure numeric & handle NaNs (features only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "46374125",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure everything in X is numeric; convert errors to NaN then fill with 0\n",
    "X = X.apply(pd.to_numeric, errors=\"coerce\").fillna(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3197f90",
   "metadata": {},
   "source": [
    "Train/Test split (NO LEAKAGE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "49c78288",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train: (36106, 2315) Test: (9027, 2315)\n"
     ]
    }
   ],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X,\n",
    "    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": "b6a3d89e",
   "metadata": {},
   "source": [
    "Standardize (fit on train only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1658a923",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = StandardScaler()\n",
    "\n",
    "X_train_scaled = scaler.fit_transform(X_train)  # fit only on train\n",
    "X_test_scaled  = scaler.transform(X_test)       # transform test with same scaler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24babe3f",
   "metadata": {},
   "source": [
    "RFE (AUC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0bb86926",
   "metadata": {},
   "outputs": [],
   "source": [
    "auc_model = LogisticRegression(\n",
    "    max_iter=2000,\n",
    "    n_jobs=-1,\n",
    "    # no class_weight here → more balanced between precision/recall\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3d2e5e00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting RFE (AUC)\n",
      "RFE (AUC) completed in 1.69 minutes\n"
     ]
    }
   ],
   "source": [
    "n_features_to_select = 100\n",
    "step_size = 0.1  # drop 10% per iteration\n",
    "\n",
    "start_time = time.time()\n",
    "print(\"Starting RFE (AUC)\")\n",
    "\n",
    "rfe_auc = RFE(\n",
    "    estimator=auc_model,\n",
    "    n_features_to_select=n_features_to_select,\n",
    "    step=step_size\n",
    ")\n",
    "\n",
    "rfe_auc.fit(X_train_scaled, y_train)\n",
    "\n",
    "end_time = time.time()\n",
    "print(f\"RFE (AUC) completed in {(end_time - start_time)/60:.2f} minutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfff4d93",
   "metadata": {},
   "source": [
    "Evaluate AUC and save top-100 list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6a1c5f16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AUC Score using RFE_AUC features: 0.9989\n",
      "\n",
      "RFE_AUC Top 100 Features:\n",
      "AGE3\n",
      "IRMARIT\n",
      "CATAG7\n",
      "DRVINDETAG\n",
      "WRKDHRSWK2\n",
      "WRK35WKUS\n",
      "WRKDPSTYR\n",
      "WRKDRGHLP\n",
      "WRKOKRAND\n",
      "IRHH65_2\n",
      "IRPINC3\n",
      "BKDRVINF\n",
      "UADOTHM\n",
      "IITRQANYREC\n",
      "CIGYR\n",
      "TOBMON\n",
      "MJSKNYR\n",
      "MJSKNMON\n",
      "CNSANYYR\n",
      "ILLEMMON\n",
      "PSILCYEVER\n",
      "RXHYDMANY\n",
      "FUCGR18\n",
      "FUNICVAP21\n",
      "FULSD18\n",
      "PYUD5MRJ\n",
      "EDUD5PNRUNM\n",
      "UD5ILALANY\n",
      "SVYRDUDANY\n",
      "SUTOUTHER\n",
      "IRSUTINRHAB\n",
      "SUTOUTCNSPY\n",
      "SUTNEEDPY\n",
      "SUNTNOHLP\n",
      "RSKMRJMON\n",
      "DIFGETCRK\n",
      "GRSKMRJWK\n",
      "DSTNRV30\n",
      "IMPWORK\n",
      "COSUITHNK\n",
      "KSSLR6YRED\n",
      "WHODASDAED\n",
      "IIDSTHOP30\n",
      "IIDSTEFF30\n",
      "IRIMPGOUT\n",
      "IIIMPPEOP\n",
      "IIIMPPEOPM\n",
      "IRIMPSOC\n",
      "IRIMPHHLD\n",
      "IIIMPHHLD\n",
      "IIIMPHHLDM\n",
      "IRIMPRESP\n",
      "IRIMPRESPM\n",
      "IRIMPWORK\n",
      "IIIMPWORK\n",
      "IRSUICTHNK\n",
      "IRCOSUITHNK\n",
      "IICOSUIPLNYR\n",
      "KSSLR6MAX\n",
      "SPDPSTYR\n",
      "AKSSLR6WRST\n",
      "WHODASTOTSC\n",
      "WHODASDASC\n",
      "SMIPPPY\n",
      "SMIPY\n",
      "AMIPY\n",
      "SMMIPY\n",
      "MMIPY\n",
      "LMIPY\n",
      "LMMIPY\n",
      "MICATPY\n",
      "AMISUD5ANY\n",
      "LMMISUD5ANY\n",
      "SMIRSUD5ANY\n",
      "AMIRSUD5ANY\n",
      "AMISUD5ANYO\n",
      "ADSUITPAYR\n",
      "ADWRELES\n",
      "ADWREMOR\n",
      "ADWRSLEP\n",
      "ADWRSMOR\n",
      "ADWRENRG\n",
      "ADWRSLOW\n",
      "ADWRSLNO\n",
      "ADWRTHOT\n",
      "ADPSHMGT\n",
      "ADRX12MO\n",
      "ASDSHOM2\n",
      "ASDSWRK2\n",
      "ASDSREL2\n",
      "ASDSSOC2\n",
      "IRMHTRXMED\n",
      "MHTRTPY\n",
      "MHNTENFCV\n",
      "COPDAGE\n",
      "CASUPROB\n",
      "CAMHRCVR\n",
      "IMFNDLEVER\n",
      "IIIMFREC\n",
      "COCLALCUSE\n"
     ]
    }
   ],
   "source": [
    "# Predict probabilities on test using only selected features\n",
    "X_test_auc = X_test_scaled[:, rfe_auc.support_]\n",
    "pred_auc   = rfe_auc.estimator_.predict_proba(X_test_auc)[:, 1]\n",
    "\n",
    "auc_score = roc_auc_score(y_test, pred_auc)\n",
    "print(f\"AUC Score using RFE_AUC features: {auc_score:.4f}\")\n",
    "\n",
    "# Get feature names for the selected mask\n",
    "feature_names = X.columns\n",
    "rfe_auc_features = feature_names[rfe_auc.support_].tolist()\n",
    "\n",
    "# Save top-100 list\n",
    "pd.Series(rfe_auc_features).to_csv(\n",
    "    \"../results/rfe_auc_top100.csv\",\n",
    "    index=False\n",
    ")\n",
    "\n",
    "print(\"\\nRFE_AUC Top 100 Features:\")\n",
    "for f in rfe_auc_features:\n",
    "    print(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6969c48",
   "metadata": {},
   "source": [
    "RFE (Recall)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "2768cafd",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_model = LogisticRegression(\n",
    "    max_iter=2000,\n",
    "    n_jobs=-1,\n",
    "    class_weight=\"balanced\"  # tilt toward minority class → better recall\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "82f27aee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting RFE (Recall)\n",
      "RFE (Recall) completed in 1.43 minutes\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "print(\"Starting RFE (Recall)\")\n",
    "\n",
    "rfe_rec = RFE(\n",
    "    estimator=rec_model,\n",
    "    n_features_to_select=n_features_to_select,\n",
    "    step=step_size\n",
    ")\n",
    "\n",
    "rfe_rec.fit(X_train_scaled, y_train)\n",
    "\n",
    "end_time = time.time()\n",
    "print(f\"RFE (Recall) completed in {(end_time - start_time)/60:.2f} minutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f81d4dc5",
   "metadata": {},
   "source": [
    "Evaluate Recall and save top-100 list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a8c2c74f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall using RFE_Recall features: 0.9894\n",
      "\n",
      "RFE_Recall Top 100 Features:\n",
      "AGE3\n",
      "IRSEX\n",
      "CATAG7\n",
      "DRVINDETAG\n",
      "SEXAGE\n",
      "WRKDHRSWK2\n",
      "WRK35WKUS\n",
      "WRKDPSTYR\n",
      "WRKDRGEDU\n",
      "WRKDRGHLP\n",
      "IRKI17_2\n",
      "IIKI17_2\n",
      "CAIDCHIP\n",
      "IRFAMSSI\n",
      "IRPINC3\n",
      "BKMVTHFT\n",
      "IIBZONMYR\n",
      "IISTMNMINIT\n",
      "CIGYR\n",
      "TOBMON\n",
      "TOBVNICFLAG\n",
      "MJSMKYR\n",
      "MJSKNYR\n",
      "MJSKNMON\n",
      "CNSANYYR\n",
      "MJONLYFLAG\n",
      "PSILCYEVER\n",
      "RXHYDMANY\n",
      "FUNICVAP21\n",
      "FULSD18\n",
      "PYUD5MRJ\n",
      "UD5ILALANY\n",
      "SUTOUTHER\n",
      "SUTOUTCNSPY\n",
      "SUNTNOHLP\n",
      "GRSKMRJWK\n",
      "SNRLDCSN\n",
      "SNRLFRND\n",
      "DSTNRV30\n",
      "DSTHOP30\n",
      "IMPREMEM\n",
      "IMPGOUT\n",
      "IMPWORK\n",
      "COSUITHNK\n",
      "KSSLR6YRED\n",
      "WHODASDAED\n",
      "IIDSTEFF30\n",
      "IRDSTNGD30\n",
      "IRIMPGOUT\n",
      "IRIMPSOC\n",
      "IRIMPHHLD\n",
      "IRIMPRESP\n",
      "IRIMPRESPM\n",
      "IRIMPWORK\n",
      "IRSUICTHNK\n",
      "IRCOSUITHNK\n",
      "KSSLR6MAX\n",
      "AKSSLR6WRST\n",
      "WHODASTOTSC\n",
      "WHODASDASC\n",
      "SMIPPPY\n",
      "SMIPY\n",
      "AMIPY\n",
      "SMMIPY\n",
      "MMIPY\n",
      "LMIPY\n",
      "LMMIPY\n",
      "MICATPY\n",
      "AMIRSUD5ANY\n",
      "AMISUD5ANYO\n",
      "ADSUITPAYR\n",
      "ADWRELES\n",
      "ADWRSLEP\n",
      "ADWRSMOR\n",
      "ADWRENRG\n",
      "ADWRSLOW\n",
      "ADWRSLNO\n",
      "ADWRTHOT\n",
      "ADPSHMGT\n",
      "ADRX12MO\n",
      "ASDSHOM2\n",
      "ASDSWRK2\n",
      "ASDSREL2\n",
      "ASDSSOC2\n",
      "IRMHTRXMED\n",
      "MHTRTPY\n",
      "MHTNSEEKPY\n",
      "MHTSKTHPY\n",
      "DIABETEAG\n",
      "COPDAGE\n",
      "ASTHMAAGE\n",
      "ASTHMANOW\n",
      "CASUPROB\n",
      "CAMHRCVR\n",
      "SYNSTMEVR\n",
      "IMFEVER\n",
      "IISYNSTMREC\n",
      "CASUPROB2\n",
      "CAMHPROB2\n",
      "COCLALCUSE\n"
     ]
    }
   ],
   "source": [
    "# Predict class labels on test using only selected features\n",
    "X_test_rec = X_test_scaled[:, rfe_rec.support_]\n",
    "pred_rec   = rfe_rec.estimator_.predict(X_test_rec)\n",
    "\n",
    "rec_score = recall_score(y_test, pred_rec)\n",
    "print(f\"Recall using RFE_Recall features: {rec_score:.4f}\")\n",
    "\n",
    "# Get feature names for the selected mask\n",
    "rfe_rec_features = feature_names[rfe_rec.support_].tolist()\n",
    "\n",
    "# Save top-100 list\n",
    "pd.Series(rfe_rec_features).to_csv(\n",
    "    \"../results/rfe_recall_top100.csv\",\n",
    "    index=False\n",
    ")\n",
    "\n",
    "print(\"\\nRFE_Recall Top 100 Features:\")\n",
    "for f in rfe_rec_features:\n",
    "    print(f)"
   ]
  }
 ],
 "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
}