## ----echo = FALSE-------------------------------------------------------------
knitr::opts_chunk$set(collapse = TRUE, comment = "#>")

## ----setup, message=FALSE-----------------------------------------------------
library(smriti)

set.seed(20250601)
n <- 200; t_points <- 4

# ── Generate clean data with a linear growth process ──────────────────────
generate_data <- function(n, add_outliers = FALSE) {
  latent_intercept <- rnorm(n, 6, 1)
  latent_slope     <- rnorm(n, 2, 1)
  data_mat <- matrix(0, n, t_points)
  for (j in seq_len(t_points)) {
    data_mat[, j] <- latent_intercept + (j - 1) * latent_slope + rnorm(n, 0, 1)
  }
  if (add_outliers) {
    idx <- sample(seq_len(n), floor(0.05 * n))
    data_mat[idx, ] <- data_mat[idx, ] + 5.0  # +5 SD shift
  }
  colnames(data_mat) <- paste0("T", seq_len(t_points))
  as.data.frame(data_mat)
}

df_clean   <- generate_data(n, add_outliers = FALSE)
df_outlier <- generate_data(n, add_outliers = TRUE)

# ── Induce 15% MAR missingness (same pattern for both) ────────────────────
set.seed(42)
apply_mar <- function(df) {
  df_miss <- df
  for (t in 1:(t_points - 1)) {
    idx <- which(!is.na(df_miss[, t]))
    x_prev <- scale(df_miss[idx, t])
    p_miss <- 1 / (1 + exp(-(x_prev - qnorm(1 - 0.15))))
    drop_idx <- idx[rbinom(length(idx), 1, p_miss) == 1]
    df_miss[drop_idx, t + 1] <- NA
  }
  df_miss
}

df_clean_miss   <- apply_mar(df_clean)
df_outlier_miss <- apply_mar(df_outlier)

cat("Clean data missingness:  ", sum(is.na(df_clean_miss)),   "cells\n")
cat("Outlier data missingness:", sum(is.na(df_outlier_miss)), "cells\n")

## ----clean-comparison, eval=FALSE---------------------------------------------
# # On clean Normal data, default and robust modes produce similar results.
# imp_clean_default <- smriti_impute(df_clean_miss, time_cols = 1:4, robust = FALSE)
# imp_clean_robust  <- smriti_impute(df_clean_miss, time_cols = 1:4, robust = TRUE)

## ----outlier-comparison, eval=FALSE-------------------------------------------
# # On outlier-contaminated data, the robust mode preserves the true structure.
# imp_outlier_default <- smriti_impute(df_outlier_miss, time_cols = 1:4,
#                                      robust = FALSE)
# imp_outlier_robust  <- smriti_impute(df_outlier_miss, time_cols = 1:4,
#                                      robust = TRUE)
# 
# # Compare recovered covariance against the true (clean) population matrix.
# true_cov <- cov(df_clean[, 1:4])  # no missingness, no outliers
# 
# cat("Default mode Frobenius distance from truth:",
#     sqrt(sum((cov(imp_outlier_default[, 1:4]) - true_cov)^2)), "\n")
# cat("Robust  mode Frobenius distance from truth:",
#     sqrt(sum((cov(imp_outlier_robust[, 1:4])  - true_cov)^2)), "\n")

## ----wrappers, eval=FALSE-----------------------------------------------------
# # missForest + smriti robust refinement
# smriti_forest(df, time_cols = 1:4, robust = TRUE)
# 
# # missRanger + smriti robust refinement
# smriti_ranger(df, time_cols = 1:4, robust = TRUE)
# 
# # Multiple imputation with robust targets on every replicate
# smriti_mi(df, time_cols = 1:4, m = 10, robust = TRUE)