import sys
import os
from pathlib import Path

import narwhals as nw
import polars as pl
import polars.selectors as cs

from survey_kit.utilities.random import RandomData
from survey_kit.utilities.dataframe import summary

from survey_kit.imputation.variable import Variable
from survey_kit.imputation.parameters import Parameters
from survey_kit.imputation.srmi import SRMI
from survey_kit.orchestration.config import Config

from survey_kit import logger, config
from survey_kit.utilities.dataframe import summary, columns_from_list

# Draw some random data

n_rows = 10_000
impute_share = 0.25

df = (
    RandomData(n_rows=n_rows, seed=32565437)
    .index("index")
    .integer("year", 2016, 2020)
    .integer("month", 1, 12)
    .integer("var2", 0, 10)
    .integer("var3", 0, 50)
    .float("var4", 0, 1)
    .integer("var5", 0, 1)
    .np_distribution("epsilon_hd1", "normal", scale=5)
    .np_distribution("epsilon_hd2", "normal", scale=5)
    .float("missing_hd1", 0, 1)
    .float("missing_hd2", 0, 1)
    .to_df()
)


#   Convenience references to them for creating dependent variables
c_var2 = pl.col("var2")
c_var3 = pl.col("var3")
c_var4 = pl.col("var4")
c_var5 = pl.col("var5")

c_e_hd1 = pl.col("epsilon_hd1")
c_e_hd2 = pl.col("epsilon_hd2")


logger.info("var_hd1 is binary and conditional on other variables")
c_hd1 = ((c_var2 * 2 - c_var3 * 3 * c_var5 + c_e_hd1)  > 0).alias("var_hd1")

logger.info("var_hd2 is != 0 only if var_hd1 == True")
c_hd2 = (
    pl.when(pl.col("var_hd1"))
      .then(((c_var2 * 1.5 - c_var3 * 1 * c_var4 + c_e_hd2)))
      .otherwise(pl.lit(0))
      .alias("var_hd2")
)
#   Create a bunch of variables that are functions of the variables created above
df = (
    df.with_columns(c_hd1)
    .with_columns(c_hd2)
    .drop(columns_from_list(df=df, columns="epsilon*"))
    .with_row_index(name="_row_index_")
)
df_original = df

#   Set variables to missing according to the uniform random variables missing_
clear_missing = []
for prefixi in ["hd"]:
    for i in range(1, 3):
        vari = f"var_{prefixi}{i}"
        missingi = f"missing_{prefixi}{i}"

        clear_missing.append(
            pl.when(pl.col(missingi) < impute_share)
            .then(pl.lit(None))
            .otherwise(pl.col(vari))
            .alias(vari)
        )
df = df.with_columns(clear_missing).drop(cs.starts_with("missing_"))

summary(df)


#   Actually do the imputation

#       The list of variables to impute (eventually)
vars_impute = []


#   1) Impute some variables to impute using stat match/hot deck
modeltype = Variable.ModelType.StatMatch
modeltype_binary = Variable.ModelType.HotDeck

#       Hot deck a continuous variable
#           Each model has a set of possible parameters
#           that determine what happens in the model
parameters_hd1 = Parameters.HotDeck(
    #   model_list - a list of variables to match
    #       donors and recipients
    model_list=["var2", "var3", "var5"],
    #   Donate anything other than the variable
    #       (i.e. donate together)
    #       In this case, it's redundant and does nothing...
    donate_list=["var_hd1"],
)

var_hd1 is binary and conditional on other variables

var_hd2 is != 0 only if var_hd1 == True

┌─────────────┬────────┬─────────────┬────────────┬─────────────┬────────────┬───────────┐
│    Variable ┆      n ┆ n (missing) ┆       mean ┆         std ┆        min ┆       max │
╞═════════════╪════════╪═════════════╪════════════╪═════════════╪════════════╪═══════════╡
│ _row_index_ ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│       index ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│        year ┆ 10,000 ┆           0 ┆ 2,017.9851 ┆    1.415937 ┆    2,016.0 ┆   2,020.0 │
│       month ┆ 10,000 ┆           0 ┆     6.5137 ┆    3.432141 ┆        1.0 ┆      12.0 │
│        var2 ┆ 10,000 ┆           0 ┆     4.9782 ┆    3.154508 ┆        0.0 ┆      10.0 │
│        var3 ┆ 10,000 ┆           0 ┆    25.1084 ┆   14.752302 ┆        0.0 ┆      50.0 │
│        var4 ┆ 10,000 ┆           0 ┆   0.505666 ┆    0.287861 ┆   0.000027 ┆  0.999997 │
│        var5 ┆ 10,000 ┆           0 ┆     0.4999 ┆    0.500025 ┆        0.0 ┆       1.0 │
│     var_hd1 ┆ 10,000 ┆       2,486 ┆   0.530077 ┆    0.499128 ┆        0.0 ┆       1.0 │
│     var_hd2 ┆ 10,000 ┆       2,515 ┆  -2.409894 ┆   10.378084 ┆ -51.051306 ┆ 26.725014 │
└─────────────┴────────┴─────────────┴────────────┴─────────────┴────────────┴───────────┘

# Set up the variable to be imputed

logger.info("Impute the boolean variable (var_hd1)")
logger.info("   by setting the model type (a stat match)")
logger.info("   and the list of match variables")
v_hd1 = Variable(
    impute_var="var_hd1",
    modeltype=Variable.ModelType.StatMatch,
    parameters=Parameters.HotDeck(
        model_list=["var2", "var3", "var5"]
    )
)

logger.info("Add the variable to the list to be imputed")
vars_impute.append(v_hd1)


logger.info("Impute the continuous variable (var_hd2) ")
logger.info("   conditional on var_hd1, using narwhals (nw.col('var_hd1'))")
logger.info("   by setting the model type (a hot deck)")
logger.info("   and the list of match variables")
logger.info("   as well as a post-processing edit to set var_hd2=0 when var_hd1==0")

v_hd2 = Variable(
    impute_var="var_hd2",
    Where=nw.col("var_hd1"),
    By=["year", "month"],
    modeltype=Variable.ModelType.HotDeck,
    parameters=Parameters.HotDeck(
        model_list=["var2", "var3", "var5"]
    ),
    postFunctions=(
        nw.when(nw.col("var_hd1"))
          .then(nw.col("var_hd2"))
          .otherwise(nw.lit(0))
          .alias("var_hd2")
    )
)
vars_impute.append(v_hd2)

Impute the boolean variable (var_hd1)

   by setting the model type (a stat match)

   and the list of match variables

Add the variable to the list to be imputed

Impute the continuous variable (var_hd2)

   conditional on var_hd1, using narwhals (nw.col('var_hd1'))

   by setting the model type (a hot deck)

   and the list of match variables

   as well as a post-processing edit to set var_hd2=0 when var_hd1==0

logger.info("Set up the imputation")
srmi = SRMI(
    df=df,
    variables=vars_impute,
    n_implicates=2,
    n_iterations=1,
    parallel=False,
    bayesian_bootstrap=True,
    parallel_testing=False,
    path_model=f"{config.path_temp_files}/py_srmi_test_hd",
    force_start=True,
)

Set up the imputation

logger.info("Run it")
srmi.run()

Run it

Variable selection before SRMI run, if necessary

     var_hd1: Method.No

     var_hd2: Method.No

Hyperparameter tuning before SRMI run, if necessary

Removing existing directory C:\Users\jonro\OneDrive\Documents\Coding\survey_kit\.scratch\temp_files/py_srmi_test_hd.srmi/1.srmi.implicate

Removing existing directory C:\Users\jonro\OneDrive\Documents\Coding\survey_kit\.scratch\temp_files/py_srmi_test_hd.srmi/2.srmi.implicate

     Imputation using statistical matching

[['var2', 'var3', 'var5'], ['var2', 'var3'], ['var2']]

     Matching on: ['var2', 'var3', 'var5']

logger.info("Get the results")
_ = df_list = srmi.df_implicates

Get the results

logger.info("\n\nLook at the original")
_ = summary(df_original)

logger.info("\n\nLook at the imputes")
_ = df_list.pipe(summary)

logger.info("\n\nLook at the imputes | var_hd1 == 0")
_ = df_list.filter(~nw.col("var_hd1")).pipe(summary)

logger.info("\n\nLook at the imputes | var_hd1 == 1")
_ = df_list.filter(nw.col("var_hd1")).pipe(summary)


Look at the original

┌─────────────┬────────┬─────────────┬────────────┬─────────────┬────────────┬──────────┐
│    Variable ┆      n ┆ n (missing) ┆       mean ┆         std ┆        min ┆      max │
╞═════════════╪════════╪═════════════╪════════════╪═════════════╪════════════╪══════════╡
│ _row_index_ ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆  9,999.0 │
│       index ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆  9,999.0 │
│        year ┆ 10,000 ┆           0 ┆ 2,017.9851 ┆    1.415937 ┆    2,016.0 ┆  2,020.0 │
│       month ┆ 10,000 ┆           0 ┆     6.5137 ┆    3.432141 ┆        1.0 ┆     12.0 │
│        var2 ┆ 10,000 ┆           0 ┆     4.9782 ┆    3.154508 ┆        0.0 ┆     10.0 │
│        var3 ┆ 10,000 ┆           0 ┆    25.1084 ┆   14.752302 ┆        0.0 ┆     50.0 │
│        var4 ┆ 10,000 ┆           0 ┆   0.505666 ┆    0.287861 ┆   0.000027 ┆ 0.999997 │
│        var5 ┆ 10,000 ┆           0 ┆     0.4999 ┆    0.500025 ┆        0.0 ┆      1.0 │
│ missing_hd1 ┆ 10,000 ┆           0 ┆   0.499985 ┆    0.288921 ┆   0.000129 ┆  0.99994 │
│ missing_hd2 ┆ 10,000 ┆           0 ┆   0.499133 ┆    0.288946 ┆   0.000133 ┆ 0.999972 │
│     var_hd1 ┆ 10,000 ┆           0 ┆      0.523 ┆    0.499496 ┆        0.0 ┆      1.0 │
│     var_hd2 ┆ 10,000 ┆           0 ┆  -2.416816 ┆   10.339524 ┆ -53.647677 ┆ 27.06033 │
└─────────────┴────────┴─────────────┴────────────┴─────────────┴────────────┴──────────┘


Look at the imputes

┌──────────────┬────────┬─────────────┬────────────┬─────────────┬────────────┬───────────┐
│     Variable ┆      n ┆ n (missing) ┆       mean ┆         std ┆        min ┆       max │
╞══════════════╪════════╪═════════════╪════════════╪═════════════╪════════════╪═══════════╡
│ ___rownumber ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│  _row_index_ ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│        index ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│         year ┆ 10,000 ┆           0 ┆ 2,017.9851 ┆    1.415937 ┆    2,016.0 ┆   2,020.0 │
│        month ┆ 10,000 ┆           0 ┆     6.5137 ┆    3.432141 ┆        1.0 ┆      12.0 │
│         var2 ┆ 10,000 ┆           0 ┆     4.9782 ┆    3.154508 ┆        0.0 ┆      10.0 │
│         var3 ┆ 10,000 ┆           0 ┆    25.1084 ┆   14.752302 ┆        0.0 ┆      50.0 │
│         var4 ┆ 10,000 ┆           0 ┆   0.505666 ┆    0.287861 ┆   0.000027 ┆  0.999997 │
│         var5 ┆ 10,000 ┆           0 ┆     0.4999 ┆    0.500025 ┆        0.0 ┆       1.0 │
│      var_hd1 ┆ 10,000 ┆           0 ┆     0.5246 ┆    0.499419 ┆        0.0 ┆       1.0 │
│      var_hd2 ┆ 10,000 ┆           0 ┆  -2.318508 ┆   10.219462 ┆ -51.051306 ┆ 26.725014 │
└──────────────┴────────┴─────────────┴────────────┴─────────────┴────────────┴───────────┘

┌──────────────┬────────┬─────────────┬────────────┬─────────────┬────────────┬───────────┐
│     Variable ┆      n ┆ n (missing) ┆       mean ┆         std ┆        min ┆       max │
╞══════════════╪════════╪═════════════╪════════════╪═════════════╪════════════╪═══════════╡
│ ___rownumber ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│  _row_index_ ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│        index ┆ 10,000 ┆           0 ┆    4,999.5 ┆ 2,886.89568 ┆        0.0 ┆   9,999.0 │
│         year ┆ 10,000 ┆           0 ┆ 2,017.9851 ┆    1.415937 ┆    2,016.0 ┆   2,020.0 │
│        month ┆ 10,000 ┆           0 ┆     6.5137 ┆    3.432141 ┆        1.0 ┆      12.0 │
│         var2 ┆ 10,000 ┆           0 ┆     4.9782 ┆    3.154508 ┆        0.0 ┆      10.0 │
│         var3 ┆ 10,000 ┆           0 ┆    25.1084 ┆   14.752302 ┆        0.0 ┆      50.0 │
│         var4 ┆ 10,000 ┆           0 ┆   0.505666 ┆    0.287861 ┆   0.000027 ┆  0.999997 │
│         var5 ┆ 10,000 ┆           0 ┆     0.4999 ┆    0.500025 ┆        0.0 ┆       1.0 │
│      var_hd1 ┆ 10,000 ┆           0 ┆      0.526 ┆    0.499349 ┆        0.0 ┆       1.0 │
│      var_hd2 ┆ 10,000 ┆           0 ┆  -2.303672 ┆   10.243192 ┆ -51.051306 ┆ 26.725014 │
└──────────────┴────────┴─────────────┴────────────┴─────────────┴────────────┴───────────┘


Look at the imputes | var_hd1 == 0

┌──────────────┬───────┬─────────────┬──────────────┬──────────────┬──────────┬──────────┐
│     Variable ┆     n ┆ n (missing) ┆         mean ┆          std ┆      min ┆      max │
╞══════════════╪═══════╪═════════════╪══════════════╪══════════════╪══════════╪══════════╡
│ ___rownumber ┆ 4,754 ┆           0 ┆ 5,041.065839 ┆ 2,874.309024 ┆      1.0 ┆  9,997.0 │
│  _row_index_ ┆ 4,754 ┆           0 ┆ 5,041.065839 ┆ 2,874.309024 ┆      1.0 ┆  9,997.0 │
│        index ┆ 4,754 ┆           0 ┆ 5,041.065839 ┆ 2,874.309024 ┆      1.0 ┆  9,997.0 │
│         year ┆ 4,754 ┆           0 ┆ 2,017.984013 ┆     1.420803 ┆  2,016.0 ┆  2,020.0 │
│        month ┆ 4,754 ┆           0 ┆     6.564367 ┆     3.414207 ┆      1.0 ┆     12.0 │
│         var2 ┆ 4,754 ┆           0 ┆     4.490745 ┆     3.246844 ┆      0.0 ┆     10.0 │
│         var3 ┆ 4,754 ┆           0 ┆    24.912495 ┆    12.861361 ┆      0.0 ┆     50.0 │
│         var4 ┆ 4,754 ┆           0 ┆     0.506516 ┆     0.289372 ┆ 0.000027 ┆ 0.999997 │
│         var5 ┆ 4,754 ┆           0 ┆     0.877787 ┆     0.327566 ┆      0.0 ┆      1.0 │
│      var_hd1 ┆ 4,754 ┆           0 ┆          0.0 ┆          0.0 ┆      0.0 ┆      0.0 │
│      var_hd2 ┆ 4,754 ┆           0 ┆          0.0 ┆          0.0 ┆      0.0 ┆      0.0 │
└──────────────┴───────┴─────────────┴──────────────┴──────────────┴──────────┴──────────┘

┌──────────────┬───────┬─────────────┬──────────────┬──────────────┬──────────┬──────────┐
│     Variable ┆     n ┆ n (missing) ┆         mean ┆          std ┆      min ┆      max │
╞══════════════╪═══════╪═════════════╪══════════════╪══════════════╪══════════╪══════════╡
│ ___rownumber ┆ 4,740 ┆           0 ┆ 5,043.214557 ┆ 2,878.670922 ┆      1.0 ┆  9,997.0 │
│  _row_index_ ┆ 4,740 ┆           0 ┆ 5,043.214557 ┆ 2,878.670922 ┆      1.0 ┆  9,997.0 │
│        index ┆ 4,740 ┆           0 ┆ 5,043.214557 ┆ 2,878.670922 ┆      1.0 ┆  9,997.0 │
│         year ┆ 4,740 ┆           0 ┆ 2,017.989241 ┆     1.418419 ┆  2,016.0 ┆  2,020.0 │
│        month ┆ 4,740 ┆           0 ┆     6.540717 ┆     3.413945 ┆      1.0 ┆     12.0 │
│         var2 ┆ 4,740 ┆           0 ┆     4.505696 ┆     3.244384 ┆      0.0 ┆     10.0 │
│         var3 ┆ 4,740 ┆           0 ┆    24.888608 ┆    12.864506 ┆      0.0 ┆     50.0 │
│         var4 ┆ 4,740 ┆           0 ┆     0.507112 ┆     0.289628 ┆ 0.000027 ┆ 0.999997 │
│         var5 ┆ 4,740 ┆           0 ┆     0.882068 ┆     0.322562 ┆      0.0 ┆      1.0 │
│      var_hd1 ┆ 4,740 ┆           0 ┆          0.0 ┆          0.0 ┆      0.0 ┆      0.0 │
│      var_hd2 ┆ 4,740 ┆           0 ┆          0.0 ┆          0.0 ┆      0.0 ┆      0.0 │
└──────────────┴───────┴─────────────┴──────────────┴──────────────┴──────────┴──────────┘


Look at the imputes | var_hd1 == 1

┌──────────────┬───────┬─────────────┬──────────────┬──────────────┬────────────┬───────────┐
│     Variable ┆     n ┆ n (missing) ┆         mean ┆          std ┆        min ┆       max │
╞══════════════╪═══════╪═════════════╪══════════════╪══════════════╪════════════╪═══════════╡
│ ___rownumber ┆ 5,246 ┆           0 ┆ 4,961.832444 ┆ 2,898.013585 ┆        0.0 ┆   9,999.0 │
│  _row_index_ ┆ 5,246 ┆           0 ┆ 4,961.832444 ┆ 2,898.013585 ┆        0.0 ┆   9,999.0 │
│        index ┆ 5,246 ┆           0 ┆ 4,961.832444 ┆ 2,898.013585 ┆        0.0 ┆   9,999.0 │
│         year ┆ 5,246 ┆           0 ┆ 2,017.986085 ┆     1.411649 ┆    2,016.0 ┆   2,020.0 │
│        month ┆ 5,246 ┆           0 ┆     6.467785 ┆     3.447995 ┆        1.0 ┆      12.0 │
│         var2 ┆ 5,246 ┆           0 ┆     5.419939 ┆     3.001108 ┆        0.0 ┆      10.0 │
│         var3 ┆ 5,246 ┆           0 ┆    25.285932 ┆    16.276467 ┆        0.0 ┆      50.0 │
│         var4 ┆ 5,246 ┆           0 ┆     0.504895 ┆      0.28651 ┆   0.000104 ┆  0.999885 │
│         var5 ┆ 5,246 ┆           0 ┆     0.157453 ┆     0.364262 ┆        0.0 ┆       1.0 │
│      var_hd1 ┆ 5,246 ┆           0 ┆          1.0 ┆          0.0 ┆        1.0 ┆       1.0 │
│      var_hd2 ┆ 5,246 ┆           0 ┆    -4.419574 ┆    13.777174 ┆ -51.051306 ┆ 26.725014 │
└──────────────┴───────┴─────────────┴──────────────┴──────────────┴────────────┴───────────┘