{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83b4897",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import narwhals as nw\n",
    "import polars as pl\n",
    "import polars.selectors as cs\n",
    "\n",
    "from survey_kit.utilities.random import RandomData\n",
    "from survey_kit.utilities.dataframe import summary\n",
    "\n",
    "from survey_kit.imputation.variable import Variable\n",
    "from survey_kit.imputation.parameters import Parameters\n",
    "from survey_kit.imputation.srmi import SRMI\n",
    "from survey_kit.imputation.selection import Selection\n",
    "\n",
    "from survey_kit import logger, config\n",
    "from survey_kit.utilities.dataframe import summary, columns_from_list\n",
    "from survey_kit.utilities.formula_builder import FormulaBuilder\n",
    "\n",
    "\n",
    "path = Path(config.code_root)\n",
    "sys.path.append(os.path.normpath(path.parent.parent / \"tests\"))\n",
    "from scratch import path_scratch\n",
    "\n",
    "\n",
    "config.data_root = path_scratch(temp_file_suffix=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a822db8f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Draw some random data\n",
    "\n",
    "n_rows = 10_000\n",
    "impute_share = 0.25\n",
    "\n",
    "\n",
    "df = (\n",
    "    RandomData(n_rows=n_rows, seed=32565437)\n",
    "    .index(\"index\")\n",
    "    .integer(\"year\", 2016, 2020)\n",
    "    .integer(\"month\", 1, 12)\n",
    "    .integer(\"var2\", 0, 10)\n",
    "    .integer(\"var3\", 0, 50)\n",
    "    .float(\"var4\", 0, 1)\n",
    "    .integer(\"var5\", 0, 1)\n",
    "    .float(\"unrelated_1\", 0, 1)\n",
    "    .float(\"unrelated_2\", 0, 1)\n",
    "    .float(\"unrelated_3\", 0, 1)\n",
    "    .float(\"unrelated_4\", 0, 1)\n",
    "    .float(\"unrelated_5\", 0, 1)\n",
    "    .np_distribution(\"epsilon_reg1\", \"normal\", scale=5)\n",
    "    .np_distribution(\"epsilon_reg2\", \"normal\", scale=5)\n",
    "    .float(\"missing_reg1\", 0, 1)\n",
    "    .float(\"missing_reg2\", 0, 1)\n",
    "    .to_df()\n",
    ")\n",
    "\n",
    "\n",
    "#   Convenience references to them for creating dependent variables\n",
    "c_var2 = pl.col(\"var2\")\n",
    "c_var3 = pl.col(\"var3\")\n",
    "c_var4 = pl.col(\"var4\")\n",
    "c_var5 = pl.col(\"var5\")\n",
    "\n",
    "c_e_reg1 = pl.col(\"epsilon_reg1\")\n",
    "c_e_reg2 = pl.col(\"epsilon_reg2\")\n",
    "\n",
    "\n",
    "#   Convenience references to them for creating dependent variables\n",
    "c_var2 = pl.col(\"var2\")\n",
    "c_var3 = pl.col(\"var3\")\n",
    "c_var4 = pl.col(\"var4\")\n",
    "c_var5 = pl.col(\"var5\")\n",
    "\n",
    "\n",
    "logger.info(\"var_reg1 is binary and conditional on other variables\")\n",
    "c_reg1 = ((c_var2 * 2 - c_var3 * 3 * c_var5 + c_e_reg1) > 0).alias(\"var_reg1\")\n",
    "\n",
    "logger.info(\"var_reg2 is != 0 only if var_reg1 == True\")\n",
    "c_reg2 = (\n",
    "    pl.when(pl.col(\"var_reg1\"))\n",
    "    .then((c_var2 * 1.5 - c_var3 * 1 * c_var4 + c_e_reg2))\n",
    "    .otherwise(pl.lit(0))\n",
    "    .alias(\"var_reg2\")\n",
    ")\n",
    "#   Create a bunch of variables that are functions of the variables created above\n",
    "df = (\n",
    "    df.with_columns(c_reg1)\n",
    "    .with_columns(c_reg2)\n",
    "    .drop(columns_from_list(df=df, columns=\"epsilon*\"))\n",
    "    .with_row_index(name=\"_row_index_\")\n",
    ")\n",
    "\n",
    "df_original = df\n",
    "\n",
    "#   Set variables to missing according to the uniform random variables missing_\n",
    "clear_missing = []\n",
    "for prefixi in [\"reg\"]:\n",
    "    for i in range(1, 3):\n",
    "        vari = f\"var_{prefixi}{i}\"\n",
    "        missingi = f\"missing_{prefixi}{i}\"\n",
    "\n",
    "        clear_missing.append(\n",
    "            pl.when(pl.col(missingi) < impute_share)\n",
    "            .then(pl.lit(None))\n",
    "            .otherwise(pl.col(vari))\n",
    "            .alias(vari)\n",
    "        )\n",
    "df = df.with_columns(clear_missing).drop(cs.starts_with(\"missing_\"))\n",
    "\n",
    "#   Make a fully collinear var for testing\n",
    "df = df.with_columns(pl.col(\"unrelated_1\").alias(\"repeat_1\"))\n",
    "\n",
    "\n",
    "summary(df)\n",
    "\n",
    "\n",
    "#   Actually do the imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03426c86",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "logger.info(\n",
    "    \"Define the regression model (intentionally include some extraneous variables\"\n",
    ")\n",
    "\n",
    "f_model = FormulaBuilder(df=df)\n",
    "f_model.formula_with_varnames_in_brackets(\n",
    "    \"~1+{var_*}+var2+var4+var4*var3*C(var5)+{unrelated_*}+{repeat_*}\"\n",
    ")\n",
    "logger.info(f_model.formula)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc1d330",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Set up the variable to be imputed\n",
    "vars_impute = []\n",
    "\n",
    "logger.info(\"Impute the boolean variable (var_reg1)\")\n",
    "logger.info(\"   to the default setup for predicted mean matching\")\n",
    "logger.info(\"   using logit regression\")\n",
    "v_reg1 = Variable(\n",
    "    impute_var=\"var_reg1\",\n",
    "    modeltype=Variable.ModelType.pmm,\n",
    "    model=f_model.formula,\n",
    "    parameters=Parameters.Regression(model=Parameters.RegressionModel.Logit),\n",
    ")\n",
    "logger.info(\"Add the variable to the list to be imputed\")\n",
    "vars_impute.append(v_reg1)\n",
    "\n",
    "logger.info(\"Impute the continuous variable (var_reg2) \")\n",
    "logger.info(\"   conditional on var_reg1, using narwhals (nw.col('var_reg1'))\")\n",
    "logger.info(\"   by setting the model type\")\n",
    "logger.info(\"   and the formula\")\n",
    "logger.info(\"   as well as a post-processing edit to set var_reg2=0 when var_reg1==0\")\n",
    "v_reg2 = Variable(\n",
    "    impute_var=\"var_reg2\",\n",
    "    Where=nw.col(\"var_reg1\"),\n",
    "    modeltype=Variable.ModelType.pmm,\n",
    "    model=f_model.formula,\n",
    "    #   Default parameters\n",
    "    parameters=Parameters.Regression(),\n",
    "    postFunctions=(\n",
    "        nw.when(nw.col(\"var_reg1\"))\n",
    "        .then(nw.col(\"var_reg2\"))\n",
    "        .otherwise(nw.lit(0))\n",
    "        .alias(\"var_reg2\")\n",
    "    ),\n",
    ")\n",
    "\n",
    "vars_impute.append(v_reg2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8cf64ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Set up the imputation\")\n",
    "logger.info(\"Add LASSO selection before each imputation\")\n",
    "srmi = SRMI(\n",
    "    df=df,\n",
    "    variables=vars_impute,\n",
    "    n_implicates=2,\n",
    "    n_iterations=2,\n",
    "    parallel=False,\n",
    "    selection=Selection(method=Selection.Method.LASSO),\n",
    "    modeltype=Variable.ModelType.pmm,\n",
    "    model=f_model.formula,\n",
    "    bayesian_bootstrap=True,\n",
    "    path_model=f\"{config.path_temp_files}/py_srmi_test_regression\",\n",
    "    force_start=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89209e27",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "logger.info(\"Run it\")\n",
    "srmi.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1835b29a",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Get the results\")\n",
    "_ = df_list = srmi.df_implicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b3150ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"\\n\\nLook at the original\")\n",
    "_ = summary(df_original)\n",
    "\n",
    "logger.info(\"\\n\\nLook at the imputes\")\n",
    "_ = df_list.pipe(summary)\n",
    "\n",
    "logger.info(\"\\n\\nLook at the imputes | var_reg1 == 0\")\n",
    "_ = df_list.filter(~nw.col(\"var_reg1\")).pipe(summary)\n",
    "\n",
    "logger.info(\"\\n\\nLook at the imputes | var_reg1 == 1\")\n",
    "_ = df_list.filter(nw.col(\"var_reg1\")).pipe(summary)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
