{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a674c7f4",
   "metadata": {},
   "source": [
    "# Imports & load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "34b64194",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shezin/.local/lib/python3.8/site-packages/pandas/core/computation/expressions.py:20: UserWarning: Pandas requires version '2.7.3' or newer of 'numexpr' (version '2.7.1' currently installed).\n",
      "  from pandas.core.computation.check import NUMEXPR_INSTALLED\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "df = pd.read_pickle(\"../data/cleaned/NSDUH_2023_clean.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2033d363",
   "metadata": {},
   "source": [
    "# Drop globally excluded features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ff2a5197",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Removing these features (person must already have a depression to asnwer these questions)\n",
    "\n",
    "FEATURES_TO_REMOVE = [\n",
    "    \"ADRX12MO\",\n",
    "    \"IRMHTRXMED\",\n",
    "    \"ADWRENRG\",\n",
    "    \"ADWRSLNO\",\n",
    "    \"ADPSHMGT\",\n",
    "]\n",
    "\n",
    "df = df.drop(columns=[c for c in FEATURES_TO_REMOVE if c in df.columns])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b52f5430",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ADRX12MO: Dropped\n",
      "IRMHTRXMED: Dropped\n",
      "ADWRENRG: Dropped\n",
      "ADWRSLNO: Dropped\n",
      "ADPSHMGT: Dropped\n"
     ]
    }
   ],
   "source": [
    "# Checking to make sure all removed\n",
    "for col in FEATURES_TO_REMOVE:\n",
    "    print(f\"{col}: {'Still exists' if col in df.columns else 'Dropped'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08144a0e",
   "metadata": {},
   "source": [
    "# Switching raw variable to recoded variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ed878820",
   "metadata": {},
   "outputs": [],
   "source": [
    "RAW_TO_RECODED = {\n",
    "    \"AGE3\": \"CATAG3\",\n",
    "    \"CAMHRCVR\": \"RCVYMHPRB\",\n",
    "    \"MJSKNYR\": \"MJCMOTHYR\",\n",
    "    \"DIFGETCRK\": \"DIFOBTCRK\",\n",
    "    \"MHTUNENFCV\": \"MHNTENFCV\",\n",
    "    \"MHTUNINSCV\": \"MHNTINSCV\",\n",
    "}\n",
    "\n",
    "for raw, recoded in RAW_TO_RECODED.items():\n",
    "    if raw in df.columns:\n",
    "        df = df.drop(columns=[raw])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f59e9f18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Removing these original non recoded features\n",
    "\n",
    "FEATURES_TO_REMOVE_NOT_RECODED = [\n",
    "    \"AGE3\",\n",
    "    \"CAMHRCVR\",\n",
    "    \"MJSKNYR\",\n",
    "    \"DIFGETCRK\",\n",
    "    \"MHTUNENFCV\",\n",
    "    \"MHTUNINSCV\",\n",
    "]\n",
    "\n",
    "df = df.drop(columns=[c for c in FEATURES_TO_REMOVE_NOT_RECODED if c in df.columns])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "879fae8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AGE3: Dropped\n",
      "CAMHRCVR: Dropped\n",
      "MJSKNYR: Dropped\n",
      "DIFGETCRK: Dropped\n",
      "MHTUNENFCV: Dropped\n",
      "MHTUNINSCV: Dropped\n"
     ]
    }
   ],
   "source": [
    "# Checking to make sure all removed\n",
    "for col in FEATURES_TO_REMOVE_NOT_RECODED:\n",
    "    print(f\"{col}: {'Still exists' if col in df.columns else 'Dropped'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eb85251",
   "metadata": {},
   "source": [
    "# Class abstraction / collapsing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b95358e",
   "metadata": {},
   "source": [
    "IRMARIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3b9ed644",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'IRMARIT Len : 2 IMPUTATION REVISED MARITAL STATUS\\nFreq Pct\\n1 = Married ......................................................................................................................... 18316 32.30\\n2 = Widowed....................................................................................................................... 1383 2.44\\n3 = Divorced or Separated .................................................................................................. 4670 8.24\\n4 = Never Been Married...................................................................................................... 26755 47.18\\n99 = LEGITIMATE SKIP Respondent is <= 14 years old.................................................. 5581 9.84'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"IRMARIT Len : 2 IMPUTATION REVISED MARITAL STATUS\n",
    "Freq Pct\n",
    "1 = Married ......................................................................................................................... 18316 32.30\n",
    "2 = Widowed....................................................................................................................... 1383 2.44\n",
    "3 = Divorced or Separated .................................................................................................. 4670 8.24\n",
    "4 = Never Been Married...................................................................................................... 26755 47.18\n",
    "99 = LEGITIMATE SKIP Respondent is <= 14 years old.................................................. 5581 9.84\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c8a58cfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# IRMARIT Abstracted\n",
    "IRMARIT_MAP = {\n",
    "    1: \"Married\",\n",
    "    2: \"Widowed-Divorced-or-Separated\",\n",
    "    3: \"Widowed-Divorced-or-Separated\",\n",
    "    4: \"Never Married\",\n",
    "}\n",
    "\n",
    "IRMARIT_ORDER = [\n",
    "    \"Married\",\n",
    "    \"Widowed-Divorced-or-Separated\",\n",
    "    \"Never Married\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1d1931b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IRMARIT_ABS missing: 0\n",
      "IRMARIT_ABS\n",
      "Married                          18309\n",
      "Widowed-Divorced-or-Separated     6040\n",
      "Never Married                    20784\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"IRMARIT_ABS\"] = df[\"IRMARIT\"].map(IRMARIT_MAP)\n",
    "\n",
    "df[\"IRMARIT_ABS\"] = pd.Categorical(\n",
    "    df[\"IRMARIT_ABS\"],\n",
    "    categories=IRMARIT_ORDER,\n",
    "    ordered=True\n",
    ")\n",
    "\n",
    "# Sanity checks\n",
    "print(\"IRMARIT_ABS missing:\", df[\"IRMARIT_ABS\"].isna().sum())\n",
    "print(df[\"IRMARIT_ABS\"].value_counts().sort_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8179e2cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drop the original variable\n",
    "df = df.drop(columns=[\"IRMARIT\"])\n",
    "\"IRMARIT\" in df.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a0f3f10",
   "metadata": {},
   "source": [
    "IRPINC3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0f2c8cc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'IRPINC3 Len : 2 RECODE -RESP TOT INCOME - IMPUTATION REVISED\\nFreq Pct\\n1 = Less than $10,000 (Including Loss).............................................................................. 21742 38.34\\n2 = $10,000 - $19,999......................................................................................................... 7740 13.65\\n3 = $20,000 - $29,999......................................................................................................... 5110 9.01\\n4 = $30,000 - $39,999......................................................................................................... 4689 8.27\\n5 = $40,000 - $49,999......................................................................................................... 3845 6.78\\n6 = $50,000 - $74,999......................................................................................................... 5592 9.86\\n7 = $75,000 or more............................................................................................................ 7987 14.09'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"IRPINC3 Len : 2 RECODE -RESP TOT INCOME - IMPUTATION REVISED\n",
    "Freq Pct\n",
    "1 = Less than $10,000 (Including Loss).............................................................................. 21742 38.34\n",
    "2 = $10,000 - $19,999......................................................................................................... 7740 13.65\n",
    "3 = $20,000 - $29,999......................................................................................................... 5110 9.01\n",
    "4 = $30,000 - $39,999......................................................................................................... 4689 8.27\n",
    "5 = $40,000 - $49,999......................................................................................................... 3845 6.78\n",
    "6 = $50,000 - $74,999......................................................................................................... 5592 9.86\n",
    "7 = $75,000 or more............................................................................................................ 7987 14.09\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2b4b39d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# IRPINC3 Abstracted\n",
    "\n",
    "IRPINC3_MAP = {\n",
    "    1: \"a_<10k\",\n",
    "    2: \"b_10k-to-40k\",\n",
    "    3: \"b_10k-to-40k\",\n",
    "    4: \"b_10k-to-40k\",\n",
    "    5: \"c_50k-to-75k\",\n",
    "    6: \"c_50k-to-75k\",\n",
    "    7: \"d_75k+\",\n",
    "}\n",
    "\n",
    "IRPINC3_ORDER = [\"a_<10k\", \"b_10k-to-40k\", \"c_50k-to-75k\", \"d_75k+\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d97d4932",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IRPINC3_ABS missing: 0\n",
      "IRPINC3_ABS\n",
      "a_<10k          10720\n",
      "b_10k-to-40k    17116\n",
      "c_50k-to-75k     9384\n",
      "d_75k+           7913\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"IRPINC3_ABS\"] = df[\"IRPINC3\"].map(IRPINC3_MAP)\n",
    "\n",
    "df[\"IRPINC3_ABS\"] = pd.Categorical(\n",
    "    df[\"IRPINC3_ABS\"],\n",
    "    categories=IRPINC3_ORDER,\n",
    "    ordered=True\n",
    ")\n",
    "\n",
    "# Sanity checks\n",
    "print(\"IRPINC3_ABS missing:\", df[\"IRPINC3_ABS\"].isna().sum())\n",
    "print(df[\"IRPINC3_ABS\"].value_counts().sort_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1947094e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drop the original variable\n",
    "df = df.drop(columns=[\"IRPINC3\"])\n",
    "\"IRPINC3\" in df.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45f3cf77",
   "metadata": {},
   "source": [
    "WRKDRGHLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3e917827",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"WRKDRGHLP Len : 2 ANY ASSISTANCE PROGRAM OFFERED THROUGH WRK\\nFreq Pct\\n1 = Yes................................................................................................................................ 14585 25.72\\n2 = No................................................................................................................................. 13733 24.22\\n85 = BAD DATA Logically assigned ................................................................................. 8 0.01\\n94 = DON'T KNOW........................................................................................................... 1569 2.77\\n97 = REFUSED .................................................................................................................. 481 0.85\\n98 = BLANK (NO ANSWER) ........................................................................................... 1435 2.53\\n99 = LEGITIMATE SKIP................................................................................................... 24894 43.90\""
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"WRKDRGHLP Len : 2 ANY ASSISTANCE PROGRAM OFFERED THROUGH WRK\n",
    "Freq Pct\n",
    "1 = Yes................................................................................................................................ 14585 25.72\n",
    "2 = No................................................................................................................................. 13733 24.22\n",
    "85 = BAD DATA Logically assigned ................................................................................. 8 0.01\n",
    "94 = DON'T KNOW........................................................................................................... 1569 2.77\n",
    "97 = REFUSED .................................................................................................................. 481 0.85\n",
    "98 = BLANK (NO ANSWER) ........................................................................................... 1435 2.53\n",
    "99 = LEGITIMATE SKIP................................................................................................... 24894 43.90\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3ca88d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "WRKDRGHLP_MAP = {\n",
    "    1: \"Having access to assistance program at work\",\n",
    "    2: \"Not having access to assistance program at work\",\n",
    "}\n",
    "\n",
    "WRKDRGHLP_ORDER = [\n",
    "    \"Having access to assistance program at work\",\n",
    "    \"Not having access to assistance program at work\",\n",
    "    \"No need or do not know about assistance program at work\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "eadf1ea9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WRKDRGHLP_ABS missing: 0\n",
      "WRKDRGHLP_ABS\n",
      "Having access to assistance program at work                14216\n",
      "Not having access to assistance program at work            12347\n",
      "No need or do not know about assistance program at work    18570\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"WRKDRGHLP_ABS\"] = df[\"WRKDRGHLP\"].map(WRKDRGHLP_MAP)\n",
    "\n",
    "# IMPORTANT: collapse skip logic / NA into third class\n",
    "df[\"WRKDRGHLP_ABS\"] = df[\"WRKDRGHLP_ABS\"].fillna(\n",
    "    \"No need or do not know about assistance program at work\"\n",
    ")\n",
    "\n",
    "# Set categorical\n",
    "df[\"WRKDRGHLP_ABS\"] = pd.Categorical(\n",
    "    df[\"WRKDRGHLP_ABS\"],\n",
    "    categories=WRKDRGHLP_ORDER,\n",
    "    ordered=True\n",
    ")\n",
    "\n",
    "# Sanity checks\n",
    "print(\"WRKDRGHLP_ABS missing:\", df[\"WRKDRGHLP_ABS\"].isna().sum())\n",
    "print(df[\"WRKDRGHLP_ABS\"].value_counts().sort_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "07473b8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WRKDRGHLP missing: 18570\n",
      "WRKDRGHLP\n",
      "1.0    14216\n",
      "2.0    12347\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(\"WRKDRGHLP missing:\", df[\"WRKDRGHLP\"].isna().sum())\n",
    "print(df[\"WRKDRGHLP\"].value_counts().sort_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dd053fd5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drop the original variable\n",
    "df = df.drop(columns=[\"WRKDRGHLP\"])\n",
    "\"WRKDRGHLP\" in df.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5f506e6",
   "metadata": {},
   "source": [
    "KSSLR6MAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a93f9793",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'KSSLR6MAX Len : 2 RC-WORST K6 TOTAL SCORE IN PAST YEAR\\nFreq Pct\\nRANGE = 0 - 24 ................................................................................................................. 45133 79.59\\n. = Aged 12-17 .................................................................................................................... 11572 20.41'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"KSSLR6MAX Len : 2 RC-WORST K6 TOTAL SCORE IN PAST YEAR\n",
    "Freq Pct\n",
    "RANGE = 0 - 24 ................................................................................................................. 45133 79.59\n",
    ". = Aged 12-17 .................................................................................................................... 11572 20.41\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "30550d5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# KSSLR6MAX Abstracted\n",
    "\n",
    "KSSLR6_BINS = [-1, 6, 12, 18, 24]\n",
    "KSSLR6_LABELS = [\"0-6\", \"07-12\", \"13-18\", \"19-24\"]\n",
    "\n",
    "df[\"KSSLR6MAX_ABS\"] = pd.cut(\n",
    "    df[\"KSSLR6MAX\"],\n",
    "    bins=KSSLR6_BINS,\n",
    "    labels=KSSLR6_LABELS\n",
    ")\n",
    "\n",
    "\n",
    "# Making it ordered categorical\n",
    "df[\"KSSLR6MAX_ABS\"] = pd.Categorical(\n",
    "    df[\"KSSLR6MAX_ABS\"],\n",
    "    categories=KSSLR6_LABELS,\n",
    "    ordered=True\n",
    ")\n",
    "\n",
    "df[\"KSSLR6MAX_ABS_ENC\"] = df[\"KSSLR6MAX_ABS\"].cat.codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f5cdb69c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Missing KSSLR6MAX_ABS: 0\n",
      "KSSLR6MAX_ABS\n",
      "0-6      27458\n",
      "07-12     8847\n",
      "13-18     5643\n",
      "19-24     3185\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(\"Missing KSSLR6MAX_ABS:\", df[\"KSSLR6MAX_ABS\"].isna().sum())\n",
    "print(\n",
    "    df[\"KSSLR6MAX_ABS\"]\n",
    "    .value_counts()\n",
    "    .sort_index()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7ccf4d9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drop the original variable\n",
    "df = df.drop(columns=[\"KSSLR6MAX\"])\n",
    "\"KSSLR6MAX\" in df.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9dac5c",
   "metadata": {},
   "source": [
    "RCVYMHPRB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0cc78f6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'RCVYMHPRB Len : 1 RC-PERCEIVED RECOVERY FROM MENTAL HEALTH ISSUE\\nFreq Pct\\n. = Aged 12-17/Unknown (Otherwise)................................................................................ 13018 22.96\\n0 = No/No Issue (CAMHRCVR=2 OR CAMHPROB2=0)................................................ 34775 61.33\\n1 = Yes (CAMHRCVR=1) ................................................................................................. 8912 15.72'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"RCVYMHPRB Len : 1 RC-PERCEIVED RECOVERY FROM MENTAL HEALTH ISSUE\n",
    "Freq Pct\n",
    ". = Aged 12-17/Unknown (Otherwise)................................................................................ 13018 22.96\n",
    "0 = No/No Issue (CAMHRCVR=2 OR CAMHPROB2=0)................................................ 34775 61.33\n",
    "1 = Yes (CAMHRCVR=1) ................................................................................................. 8912 15.72\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "67381a3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RCVYMHPRB\n",
       "0.0    34775\n",
       "1.0     8912\n",
       "NaN     1446\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"RCVYMHPRB\"].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "604a58e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop rows with missing RCVYMHPRB\n",
    "df = df.dropna(subset=[\"RCVYMHPRB\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d6822eb7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RCVYMHPRB\n",
       "0.0    34775\n",
       "1.0     8912\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"RCVYMHPRB\"].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bffdabc",
   "metadata": {},
   "source": [
    "IRIMPGOUT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c72fff3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"IRIMPGOUT Len : 2 DIFFICULTY GOING OUT ONE MO IN PST 12 MOS - IMP REV\\nFreq Pct\\n1 = No difficulty ................................................................................................................. 20700 36.50\\n2 = Mild difficulty............................................................................................................... 7015 12.37\\n3 = Moderate difficulty ....................................................................................................... 4505 7.94\\n4 = Severe difficulty............................................................................................................ 2059 3.63\\n5 = You didn't leave the house on your own ....................................................................... 1188 2.10\\n99 = LEGITIMATE SKIP................................................................................................... 21238 37.45\""
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"IRIMPGOUT Len : 2 DIFFICULTY GOING OUT ONE MO IN PST 12 MOS - IMP REV\n",
    "Freq Pct\n",
    "1 = No difficulty ................................................................................................................. 20700 36.50\n",
    "2 = Mild difficulty............................................................................................................... 7015 12.37\n",
    "3 = Moderate difficulty ....................................................................................................... 4505 7.94\n",
    "4 = Severe difficulty............................................................................................................ 2059 3.63\n",
    "5 = You didn't leave the house on your own ....................................................................... 1188 2.10\n",
    "99 = LEGITIMATE SKIP................................................................................................... 21238 37.45\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "c97dcace",
   "metadata": {},
   "outputs": [],
   "source": [
    "# IRIMPGOUT Abstracted\n",
    "\n",
    "IRIMPGOUT_MAP = {\n",
    "    1: \"b_No difficulty\",\n",
    "    2: \"c_Mild difficulty\",\n",
    "    3: \"d_Moderate difficulty\",\n",
    "    4: \"e_Severe difficulty\",\n",
    "    5: \"a_Not-Going-Out-Alone-Or-No-Difficulty\",\n",
    "    99: \"a_Not-Going-Out-Alone-Or-No-Difficulty\",\n",
    "}\n",
    "\n",
    "IRIMPGOUT_ORDER = [\n",
    "    \"a_Not-Going-Out-Alone-Or-No-Difficulty\",\n",
    "    \"b_No difficulty\",\n",
    "    \"c_Mild difficulty\",\n",
    "    \"d_Moderate difficulty\",\n",
    "    \"e_Severe difficulty\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "3b342a09",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IRIMPGOUT_ABS missing: 0\n",
      "IRIMPGOUT_ABS\n",
      "a_Not-Going-Out-Alone-Or-No-Difficulty    10583\n",
      "b_No difficulty                           20070\n",
      "c_Mild difficulty                          6726\n",
      "d_Moderate difficulty                      4315\n",
      "e_Severe difficulty                        1993\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"IRIMPGOUT_ABS\"] = df[\"IRIMPGOUT\"].map(IRIMPGOUT_MAP)\n",
    "\n",
    "# Collapse skip logic / NA into the combined class\n",
    "df[\"IRIMPGOUT_ABS\"] = df[\"IRIMPGOUT_ABS\"].fillna(\n",
    "    \"a_Not-Going-Out-Alone-Or-No-Difficulty\"\n",
    ")\n",
    "\n",
    "# Set categorical\n",
    "df[\"IRIMPGOUT_ABS\"] = pd.Categorical(\n",
    "    df[\"IRIMPGOUT_ABS\"],\n",
    "    categories=IRIMPGOUT_ORDER,\n",
    "    ordered=True\n",
    ")\n",
    "\n",
    "# Sanity checks\n",
    "print(\"IRIMPGOUT_ABS missing:\", df[\"IRIMPGOUT_ABS\"].isna().sum())\n",
    "print(df[\"IRIMPGOUT_ABS\"].value_counts().sort_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "65a8ace7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drop the original variable\n",
    "df = df.drop(columns=[\"IRIMPGOUT\"])\n",
    "\"IRIMPGOUT\" in df.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c2641ff",
   "metadata": {},
   "source": [
    "WHODASDASC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "deabd330",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'WHODASDASC Len : 2 RC-ALTERNATIVE WHODAS SCORE\\nFreq Pct\\nRANGE = 0 - 8................................................................................................................... 45133 79.59\\n. = Aged 12-17 .................................................................................................................... 11572 20.41'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"WHODASDASC Len : 2 RC-ALTERNATIVE WHODAS SCORE\n",
    "Freq Pct\n",
    "RANGE = 0 - 8................................................................................................................... 45133 79.59\n",
    ". = Aged 12-17 .................................................................................................................... 11572 20.41\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "4efe9f59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# WHODASDASC Abstracted\n",
    "\n",
    "WHODASDASC_MAP = {\n",
    "    0: \"0\",\n",
    "    1: \"1-2\",\n",
    "    2: \"1-2\",\n",
    "    3: \"3-4\",\n",
    "    4: \"3-4\",\n",
    "    5: \"5-6\",\n",
    "    6: \"5-6\",\n",
    "    7: \"7-8\",\n",
    "    8: \"7-8\",\n",
    "}\n",
    "\n",
    "WHODASDASC_ORDER = [\n",
    "    \"0\",\n",
    "    \"1-2\",\n",
    "    \"3-4\",\n",
    "    \"5-6\",\n",
    "    \"7-8\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "160d2e68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WHODASDASC_ABS missing: 0\n",
      "WHODASDASC_ABS\n",
      "0      27951\n",
      "1-2     5588\n",
      "3-4     3708\n",
      "5-6     3039\n",
      "7-8     3401\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"WHODASDASC_ABS\"] = df[\"WHODASDASC\"].map(WHODASDASC_MAP)\n",
    "\n",
    "df[\"WHODASDASC_ABS\"] = pd.Categorical(\n",
    "    df[\"WHODASDASC_ABS\"],\n",
    "    categories=WHODASDASC_ORDER,\n",
    "    ordered=True\n",
    ")\n",
    "\n",
    "# Sanity checks\n",
    "print(\"WHODASDASC_ABS missing:\", df[\"WHODASDASC_ABS\"].isna().sum())\n",
    "print(df[\"WHODASDASC_ABS\"].value_counts().sort_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "0a516936",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    }
   ],
   "source": [
    "# Drop the original variable\n",
    "df = df.drop(columns=[\"WHODASDASC\"])\n",
    "\n",
    "print(\"WHODASDASC\" in df.columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10c76c96",
   "metadata": {},
   "source": [
    "COCLALCUSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "161a0755",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'COCLALCUSE Len : 1 RC-COLLAPSED HOW COVID-19 AFFECTED AMOUNT ALCOHOL DRANK\\nFreq Pct\\n. = Unknown/No PY Alcohol Use (Otherwise)................................................................... 26038 45.92\\n1 = Drink much less or little less (COALCUSE=1,2) ......................................................... 8670 15.29\\n2 = Drink about the same (COALCUSE=3)........................................................................ 18142 31.99\\n3 = Drink a little more or much more (COALCUSE=4,5) .................................................. 3855 6.80'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"COCLALCUSE Len : 1 RC-COLLAPSED HOW COVID-19 AFFECTED AMOUNT ALCOHOL DRANK\n",
    "Freq Pct\n",
    ". = Unknown/No PY Alcohol Use (Otherwise)................................................................... 26038 45.92\n",
    "1 = Drink much less or little less (COALCUSE=1,2) ......................................................... 8670 15.29\n",
    "2 = Drink about the same (COALCUSE=3)........................................................................ 18142 31.99\n",
    "3 = Drink a little more or much more (COALCUSE=4,5) .................................................. 3855 6.80\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "d995c14e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Map original values to semantic labels\n",
    "COCLALCUSE_MAP = {\n",
    "    1: \"Drink much less or little less\",\n",
    "    2: \"Drink about the same\",\n",
    "    3: \"Drink a little more or much more\",\n",
    "}\n",
    "\n",
    "df[\"COCLALCUSE_ABS\"] = df[\"COCLALCUSE\"].map(COCLALCUSE_MAP)\n",
    "\n",
    "# Explicitly assign missing / skip logic\n",
    "df[\"COCLALCUSE_ABS\"] = df[\"COCLALCUSE_ABS\"].fillna(\n",
    "    \"Unknown/No PY Alcohol Use\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "71e33771",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COCLALCUSE Abstracted\n",
    "\n",
    "COCLALCUSE_LABEL_MAP = {\n",
    "    \"Unknown/No PY Alcohol Use\": \"a_Unknown/No PY Alcohol Use\",\n",
    "    \"Drink much less or little less\": \"b_Drink much less or little less\",\n",
    "    \"Drink about the same\": \"c_Drink about the same\",\n",
    "    \"Drink a little more or much more\": \"d_Drink a little more or much more\",\n",
    "}\n",
    "\n",
    "df[\"COCLALCUSE_ABS\"] = df[\"COCLALCUSE_ABS\"].map(COCLALCUSE_LABEL_MAP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "bec3ef19",
   "metadata": {},
   "outputs": [],
   "source": [
    "COCLALCUSE_ORDER = [\n",
    "    \"a_Unknown/No PY Alcohol Use\",\n",
    "    \"b_Drink much less or little less\",\n",
    "    \"c_Drink about the same\",\n",
    "    \"d_Drink a little more or much more\",\n",
    "]\n",
    "\n",
    "df[\"COCLALCUSE_ABS\"] = pd.Categorical(\n",
    "    df[\"COCLALCUSE_ABS\"],\n",
    "    categories=COCLALCUSE_ORDER,\n",
    "    ordered=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "17407ea4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "COCLALCUSE_ABS missing: 0\n",
      "COCLALCUSE_ABS\n",
      "a_Unknown/No PY Alcohol Use           14755\n",
      "b_Drink much less or little less       8023\n",
      "c_Drink about the same                17447\n",
      "d_Drink a little more or much more     3462\n",
      "Name: count, dtype: int64\n",
      "category\n"
     ]
    }
   ],
   "source": [
    "print(\"COCLALCUSE_ABS missing:\", df[\"COCLALCUSE_ABS\"].isna().sum())\n",
    "print(df[\"COCLALCUSE_ABS\"].value_counts().sort_index())\n",
    "print(df[\"COCLALCUSE_ABS\"].dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "e0097f08",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.drop(columns=[\"COCLALCUSE\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b8910be",
   "metadata": {},
   "source": [
    "# Save prepared dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "afa49dd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepared dataset saved to: ../data/prepared/NSDUH_2023_prepared.pkl\n",
      "Final shape: (43687, 2306)\n"
     ]
    }
   ],
   "source": [
    "OUTPUT_PATH = \"../data/prepared/NSDUH_2023_prepared.pkl\"\n",
    "\n",
    "df.to_pickle(OUTPUT_PATH)\n",
    "\n",
    "print(f\"Prepared dataset saved to: {OUTPUT_PATH}\")\n",
    "print(f\"Final shape: {df.shape}\")"
   ]
  }
 ],
 "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
}