import numpy as np
import polars as pl

from survey_kit.imputation.srmi import SRMI
from survey_kit.utilities.dataframe import summary
from survey_kit import logger, config

# Draw some random data with a few variables missing

n_rows = 5_000
rng = np.random.default_rng(20260913)

x1 = rng.normal(size=n_rows)
x2 = rng.normal(size=n_rows)

income_latent = 40_000 + 8_000 * x1 - 3_000 * x2 + rng.normal(scale=5_000, size=n_rows)
has_disability = (rng.normal(size=n_rows) + 0.5 * x1 > 1.0).astype(int)

df = pl.DataFrame(
    dict(
        person_id=range(n_rows),
        age=rng.integers(18, 90, size=n_rows),
        education_years=rng.integers(8, 20, size=n_rows),
        income=income_latent,
        has_disability=has_disability,
    )
)

#   Punch some holes in the two variables we want imputed
missing_income = rng.random(n_rows) < 0.2
missing_disability = rng.random(n_rows) < 0.15
df = df.with_columns(
    [
        pl.when(pl.Series(missing_income)).then(None).otherwise(pl.col("income")).alias(
            "income"
        ),
        pl.when(pl.Series(missing_disability))
        .then(None)
        .otherwise(pl.col("has_disability"))
        .alias("has_disability"),
    ]
)

logger.info(f"Missing income:      {missing_income.sum()} of {n_rows} rows")
logger.info(f"Missing disability:  {missing_disability.sum()} of {n_rows} rows")

Missing income:      952 of 5000 rows

Missing disability:  698 of 5000 rows

logger.info("That's it - point SRMI.simple_model() at the dataframe and an id column")
logger.info(
    "It finds every column with missing values, figures out whether each one is "
    "binary or continuous, and picks a sensible model for each - no Variable objects "
    "to build by hand."
)

srmi = SRMI.simple_model(
    df=df,
    index="person_id",
    replication=SRMI.Replication(n_implicates=3, n_iterations=3),
    parallel=SRMI.Parallel(enabled=False),
    bootstrap=SRMI.Bootstrap(enabled=True),
    storage=SRMI.Storage(
        path_model=f"{config.path_temp_files}/tutorial_simple_model_getting_started",
        force_start=True,
    ),
)

That's it - point SRMI.simple_model() at the dataframe and an id column

It finds every column with missing values, figures out whether each one is binary or continuous, and picks a sensible model for each - no Variable objects to build by hand.

auto_detect: 'income' -> class=continuous, modeltype=LightGBM, predictors=['age', 'education_years', 'has_disability']

auto_detect: 'has_disability' -> class=binary, modeltype=LightGBM, predictors=['age', 'education_years', 'income']

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

logger.info("Inspect what it decided before running anything")
for v in srmi.variables:
    logger.info(f"  {v.impute_var}: modeltype={v.modeltype.name}, predictors={v.model}")

Inspect what it decided before running anything

  income: modeltype=LightGBM, predictors=['age', 'education_years', 'has_disability']

  has_disability: modeltype=LightGBM, predictors=['age', 'education_years', 'income']

logger.info("Now actually run the imputation")
srmi.run()

Now actually run the imputation

Variable selection before SRMI run, if necessary

     income: Method.No

     has_disability: Method.No

Hyperparameter tuning before SRMI run, if necessary

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

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

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

     Imputation using LightGBM

Running lightgbm model with parameters: {'objective': 'regression', 'metric': 'rmse', 'boosting': 'gbdt', 'min_data_per_group': 25, 'num_threads': 1, 'verbose': -1, 'seed': 3231957680}

logger.info("Look at one completed implicate")
_ = srmi.df_implicates.pipe(summary)

Look at one completed implicate

┌─────────────────┬───────┬─────────────┬───────────────┬──────────────┬──────────────┬───────────────┐
│        Variable ┆     n ┆ n (missing) ┆          mean ┆          std ┆          min ┆           max │
╞═════════════════╪═══════╪═════════════╪═══════════════╪══════════════╪══════════════╪═══════════════╡
│       person_id ┆ 5,000 ┆           0 ┆       2,499.5 ┆ 1,443.520003 ┆          0.0 ┆       4,999.0 │
│             age ┆ 5,000 ┆           0 ┆       53.6372 ┆     20.60971 ┆         18.0 ┆          89.0 │
│ education_years ┆ 5,000 ┆           0 ┆       13.5238 ┆      3.46194 ┆          8.0 ┆          19.0 │
│          income ┆ 5,000 ┆           0 ┆ 40,143.741992 ┆ 9,741.858426 ┆ 6,435.355391 ┆ 73,924.778328 │
│  has_disability ┆ 5,000 ┆           0 ┆        0.1902 ┆     0.392498 ┆          0.0 ┆           1.0 │
└─────────────────┴───────┴─────────────┴───────────────┴──────────────┴──────────────┴───────────────┘

┌─────────────────┬───────┬─────────────┬───────────────┬──────────────┬──────────────┬───────────────┐
│        Variable ┆     n ┆ n (missing) ┆          mean ┆          std ┆          min ┆           max │
╞═════════════════╪═══════╪═════════════╪═══════════════╪══════════════╪══════════════╪═══════════════╡
│       person_id ┆ 5,000 ┆           0 ┆       2,499.5 ┆ 1,443.520003 ┆          0.0 ┆       4,999.0 │
│             age ┆ 5,000 ┆           0 ┆       53.6372 ┆     20.60971 ┆         18.0 ┆          89.0 │
│ education_years ┆ 5,000 ┆           0 ┆       13.5238 ┆      3.46194 ┆          8.0 ┆          19.0 │
│          income ┆ 5,000 ┆           0 ┆ 40,140.109014 ┆ 9,723.732023 ┆ 6,435.355391 ┆ 73,924.778328 │
│  has_disability ┆ 5,000 ┆           0 ┆        0.1922 ┆     0.394069 ┆          0.0 ┆           1.0 │
└─────────────────┴───────┴─────────────┴───────────────┴──────────────┴──────────────┴───────────────┘

┌─────────────────┬───────┬─────────────┬───────────────┬──────────────┬──────────────┬───────────────┐
│        Variable ┆     n ┆ n (missing) ┆          mean ┆          std ┆          min ┆           max │
╞═════════════════╪═══════╪═════════════╪═══════════════╪══════════════╪══════════════╪═══════════════╡
│       person_id ┆ 5,000 ┆           0 ┆       2,499.5 ┆ 1,443.520003 ┆          0.0 ┆       4,999.0 │
│             age ┆ 5,000 ┆           0 ┆       53.6372 ┆     20.60971 ┆         18.0 ┆          89.0 │
│ education_years ┆ 5,000 ┆           0 ┆       13.5238 ┆      3.46194 ┆          8.0 ┆          19.0 │
│          income ┆ 5,000 ┆           0 ┆ 40,360.822393 ┆ 9,743.867376 ┆ 6,435.355391 ┆ 73,924.778328 │
│  has_disability ┆ 5,000 ┆           0 ┆        0.1892 ┆     0.391707 ┆          0.0 ┆           1.0 │
└─────────────────┴───────┴─────────────┴───────────────┴──────────────┴──────────────┴───────────────┘