import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import missingno as msno
from sklearn.impute import SimpleImputer, KNNImputer, IterativeImputer
from sklearn.experimental import enable_iterative_imputer # Required for IterativeImputer
from sklearn.linear_model import BayesianRidge # Default for IterativeImputer
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import LabelEncoder, OneHotEncoder, StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.pipeline import Pipeline
from scipy.stats import kstest, chi2_contingency
from scipy.spatial.distance import jensenshannon
# For Little's MCAR test (install if you don't have it)
# pip install pingouin
try:
    import pingouin as pg
except ImportError:
    print("Pingouin not found. Install it for Little's MCAR test: pip install pingouin")
    pg = None
# Create a sample DataFrame with missing values
np.random.seed(42)
data = {
    'FeatureA': np.random.rand(100),
    'FeatureB': np.random.randint(1, 100, 100),
    'FeatureC': np.random.choice(['X', 'Y', 'Z'], 100),
    'FeatureD': np.random.normal(0, 1, 100),
    'Target': np.random.normal(50, 10, 100)
}
df = pd.DataFrame(data)
# Introduce missing values - simulating different patterns
# MCAR-like
df.loc[np.random.choice(df.index, 10, replace=False), 'FeatureA'] = np.nan
df.loc[np.random.choice(df.index, 5, replace=False), 'FeatureC'] = np.nan
# MAR-like: FeatureD missing if FeatureA is high
df.loc[df['FeatureA'] > 0.7, 'FeatureD'] = np.nan
# MNAR-like: Target missing if FeatureB is low (true value influencing missingness)
df.loc[df['FeatureB'] < 20, 'Target'] = np.nan
print("Original DataFrame head with NaNs:")
print(df.head())
print("\nMissing values per column:")
print(df.isnull().sum())
print("\nPercentage of missing values per column:")
print(df.isnull().sum() / len(df) * 100)

# Visualize missing data patterns
print("\nMissing Data Matrix:")
msno.matrix(df, figsize=(10, 5))
plt.title("Missing Data Matrix")
plt.show()
print("\nMissing Data Bar Plot:")
msno.bar(df, figsize=(10, 5))
plt.title("Missing Data Bar Plot")
plt.show()
# If you suspect relationships between missingness
print("\nMissing Data Heatmap (Correlations of missingness):")
msno.heatmap(df, figsize=(10, 5))
plt.title("Missing Data Heatmap")
plt.show()
# Dendrogram for hierarchical clustering of missingness
print("\nMissing Data Dendrogram:")
msno.dendrogram(df, figsize=(10, 6))
plt.title("Missing Data Dendrogram")
plt.show()

if pg:
    # Little's MCAR test (only works on numerical data by default, you might need to convert/exclude categorical)
    # It requires complete cases for some columns for the test to run.
    # We will only run it on numerical columns that have missing data.
    numerical_cols_with_nan = df.select_dtypes(include=np.number).columns[df.select_dtypes(include=np.number).isnull().any()].tolist()
    if numerical_cols_with_nan:
        print("\nLittle's MCAR test:")
        # We need to drop rows with NaNs in categorical columns first if we want to include those in the test,
        # or just run it on numerical subset
        df_for_mcar_test = df[numerical_cols_with_nan].copy()
        # Drop rows where ALL relevant columns are NaN, as it confuses the test
        df_for_mcar_test = df_for_mcar_test.dropna(how='all')
        if not df_for_mcar_test.empty:
            mcar_test_result = pg.mcar(df_for_mcar_test)
            print(mcar_test_result)
            if mcar_test_result['p-value'].iloc[0] < 0.05:
                print("Result: Reject null hypothesis of MCAR. Data is likely MAR or MNAR.")
            else:
                print("Result: Fail to reject null hypothesis of MCAR. Data could be MCAR.")
        else:
            print("Not enough non-missing data in numerical columns for Little's MCAR test.")
    else:
        print("No numerical columns with missing data to perform Little's MCAR test.")
else:
    print("Little's MCAR test skipped because 'pingouin' library is not installed.")

# Separate features and target for demonstration (though we might impute target as well, carefully)
X = df.drop('Target', axis=1)
y = df['Target']
# Identify column types
numerical_cols = X.select_dtypes(include=np.number).columns
categorical_cols = X.select_dtypes(include='object').columns
# Create preprocessors
numerical_transformer = Pipeline(steps=[
    ('scaler', StandardScaler())
])
categorical_transformer = Pipeline(steps=[
    ('onehot', OneHotEncoder(handle_unknown='ignore'))
])
# Combine preprocessors (before imputation for most advanced methods)
from sklearn.compose import ColumnTransformer
preprocessor = ColumnTransformer(
    transformers=[
        ('num', numerical_transformer, numerical_cols),
        ('cat', categorical_transformer, categorical_cols)
    ],
    remainder='passthrough' # Keep other columns (like our target initially)
)
# Apply preprocessing to X
X_processed = preprocessor.fit_transform(X)
# Get feature names after one-hot encoding for easier interpretation
new_numerical_cols = numerical_cols.tolist()
if len(categorical_cols) > 0:
    onehot_feature_names = preprocessor.named_transformers_['cat']['onehot'].get_feature_names_out(categorical_cols)
    all_feature_names = new_numerical_cols + list(onehot_feature_names)
else:
    all_feature_names = new_numerical_cols
# Convert back to DataFrame for imputation, retaining column names
X_processed_df = pd.DataFrame(X_processed, columns=all_feature_names, index=X.index)
# Re-introduce NaNs into the processed DataFrame for imputation
# Note: For real scenarios, you'd apply preprocessing and imputation *together* in a pipeline,
# but for demonstration, we re-create the NaNs here after the initial preprocessing pass.
# The NaNs are in the original X, so we need to map them back
for col in numerical_cols:
    X_processed_df.loc[X[col].isnull(), col] = np.nan
for col in categorical_cols:
    # Need to find the one-hot encoded columns corresponding to the original categorical column
    # and set them all to NaN if the original was NaN
    if X[col].isnull().any():
        for onehot_col in [f for f in onehot_feature_names if f.startswith(col + '_')]:
            X_processed_df.loc[X[col].isnull(), onehot_col] = np.nan
# Ensure the target variable has NaNs as in the original df
y_imputable = y.copy()

print("\n--- Simple Imputation (Mean/Median/Mode) ---")
# --- Mean Imputation (for numerical) ---
imputer_mean = SimpleImputer(strategy='mean')
X_mean_imputed = X_processed_df.copy() # Make a copy to avoid modifying original
# Apply only to numerical columns that had NaNs
cols_to_impute_mean = [col for col in X_processed_df.columns if col in numerical_cols and X_processed_df[col].isnull().any()]
X_mean_imputed[cols_to_impute_mean] = imputer_mean.fit_transform(X_mean_imputed[cols_to_impute_mean])
print("Mean Imputed Numerical Data head (first 5 rows):")
print(X_mean_imputed.head())
print("Missing values after mean imputation:", X_mean_imputed[cols_to_impute_mean].isnull().sum().sum())
# --- Median Imputation (for numerical) ---
imputer_median = SimpleImputer(strategy='median')
X_median_imputed = X_processed_df.copy()
cols_to_impute_median = [col for col in X_processed_df.columns if col in numerical_cols and X_processed_df[col].isnull().any()]
X_median_imputed[cols_to_impute_median] = imputer_median.fit_transform(X_median_imputed[cols_to_impute_median])
print("\nMedian Imputed Numerical Data head (first 5 rows):")
print(X_median_imputed.head())
print("Missing values after median imputation:", X_median_imputed[cols_to_impute_median].isnull().sum().sum())
# --- Mode Imputation (for numerical or one-hot encoded categorical) ---
# For one-hot encoded, mode typically means finding the most frequent '1' or '0'
imputer_mode = SimpleImputer(strategy='most_frequent')
X_mode_imputed = X_processed_df.copy()
cols_to_impute_mode = [col for col in X_processed_df.columns if X_processed_df[col].isnull().any()]
X_mode_imputed[cols_to_impute_mode] = imputer_mode.fit_transform(X_mode_imputed[cols_to_impute_mode])
print("\nMode Imputed Data head (first 5 rows):")
print(X_mode_imputed.head())
print("Missing values after mode imputation:", X_mode_imputed[cols_to_impute_mode].isnull().sum().sum())

print("\n--- KNN Imputation ---")
# K-NN imputation works well on scaled numerical data.
# It can handle one-hot encoded categorical data.
imputer_knn = KNNImputer(n_neighbors=5) # n_neighbors is a hyperparameter to tune
X_knn_imputed = pd.DataFrame(imputer_knn.fit_transform(X_processed_df), columns=X_processed_df.columns, index=X_processed_df.index)
print("KNN Imputed Data head (first 5 rows):")
print(X_knn_imputed.head())
print("Missing values after KNN imputation:", X_knn_imputed.isnull().sum().sum())

print("\n--- Iterative Imputation (MICE-like) ---")
# IterativeImputer requires numerical input, handles one-hot encoded categorical data well.
# We can specify the estimator, e.g., BayesianRidge, RandomForestRegressor, etc.
# Default estimator is BayesianRidge.
imputer_iterative = IterativeImputer(
    estimator=RandomForestRegressor(n_estimators=10, random_state=42), # Or BayesianRidge(), ExtraTreesRegressor(), etc.
    max_iter=10, # Number of imputation rounds
    random_state=42
)
X_iterative_imputed = pd.DataFrame(imputer_iterative.fit_transform(X_processed_df), columns=X_processed_df.columns, index=X_processed_df.index)
print("Iterative Imputed Data head (first 5 rows):")
print(X_iterative_imputed.head())
print("Missing values after Iterative imputation:", X_iterative_imputed.isnull().sum().sum())

# Select a feature that had missing values for detailed comparison
feature_to_compare = 'FeatureA' # This was MCAR-like
feature_to_compare_mar = 'FeatureD' # This was MAR-like
feature_to_compare_target = 'Target' # This was MNAR-like
# --- 1. Distribution Comparison ---
print("\n--- Evaluation: Distribution Comparison ---")
# For numerical features
def plot_imputation_distributions(original_df, imputed_dfs, feature_name, title_prefix=""):
    plt.figure(figsize=(10, 6))
    sns.kdeplot(original_df[feature_name].dropna(), label='Original Observed', color='blue', linestyle='--')
    for name, df_imputed in imputed_dfs.items():
        # Only plot if the feature was actually imputed (i.e., had NaNs)
        if original_df[feature_name].isnull().any():
            sns.kdeplot(df_imputed[feature_name], label=f'Imputed ({name})', alpha=0.7)
    plt.title(f"{title_prefix} Distribution Comparison for '{feature_name}'")
    plt.legend()
    plt.show()
# Create a dictionary of imputed DataFrames for easy iteration
imputed_data_frames = {
    'Mean': X_mean_imputed,
    'Median': X_median_imputed,
    'Mode': X_mode_imputed,
    'KNN': X_knn_imputed,
    'Iterative': X_iterative_imputed
}
# Plot for FeatureA (MCAR-like numerical)
# Need to transform original X to match the scaling/encoding of imputed_data_frames
# Let's create a full original_processed_df for comparison, where only actual NaNs exist
original_processed_full_df = X_processed_df.copy()
plot_imputation_distributions(original_processed_full_df, imputed_data_frames, preprocessor.named_transformers_['num'].get_feature_names_out(['FeatureA'])[0], "Numerical (MCAR-like)")
plot_imputation_distributions(original_processed_full_df, imputed_data_frames, preprocessor.named_transformers_['num'].get_feature_names_out(['FeatureD'])[0], "Numerical (MAR-like)")
# For categorical features (using one-hot encoded columns)
# Let's pick one of the one-hot encoded columns for FeatureC
feature_c_onehot_col = 'FeatureC_Y' # Or 'FeatureC_X', 'FeatureC_Z'
plt.figure(figsize=(12, 6))
# Original distribution of '1's for FeatureC_Y where not NaN
original_count = original_processed_full_df[feature_c_onehot_col].dropna().value_counts(normalize=True)
plt.bar(original_count.index, original_count.values, width=0.2, label='Original Observed', align='center', alpha=0.7, color='blue')
offset = -0.2
for name, df_imputed in imputed_data_frames.items():
    if original_processed_full_df[feature_c_onehot_col].isnull().any():
        imputed_count = df_imputed[feature_c_onehot_col].value_counts(normalize=True)
        # Shift bars for better visibility
        plt.bar(np.array(imputed_count.index) + offset, imputed_count.values, width=0.15, label=f'Imputed ({name})', align='center', alpha=0.6)
        offset += 0.15
plt.title(f"Distribution Comparison for '{feature_c_onehot_col}' (Categorical)")
plt.xticks([0, 1], ['0', '1'])
plt.legend()
plt.show()
# --- 2. Relationship Preservation (Correlation Matrix) ---
print("\n--- Evaluation: Relationship Preservation (Correlation Matrix) ---")
# Calculate original correlations (using listwise deletion for simplicity in comparison)
original_corr = original_processed_full_df.dropna().corr()
plt.figure(figsize=(18, 5))
plt.subplot(1, 3, 1)
sns.heatmap(original_corr, annot=True, cmap='coolwarm', fmt=".2f", linewidths=.5)
plt.title('Original (Listwise Deletion) Correlations')
# Compare with one of the advanced imputation methods, e.g., Iterative
imputed_corr_iterative = X_iterative_imputed.corr()
plt.subplot(1, 3, 2)
sns.heatmap(imputed_corr_iterative, annot=True, cmap='coolwarm', fmt=".2f", linewidths=.5)
plt.title('Iterative Imputed Correlations')
# You can also look at the difference matrix
diff_corr = abs(original_corr - imputed_corr_iterative)
plt.subplot(1, 3, 3)
sns.heatmap(diff_corr, annot=True, cmap='viridis', fmt=".2f", linewidths=.5)
plt.title('Absolute Difference in Correlations')
plt.show()
# --- 3. Model Performance ---
print("\n--- Evaluation: Model Performance ---")
# Re-create our target variable `y` including its NaNs, and impute it for the model comparison
y_imputed_mean = SimpleImputer(strategy='mean').fit_transform(y_imputable.to_frame())
y_imputed_knn = KNNImputer(n_neighbors=5).fit_transform(y_imputable.to_frame())
y_imputed_iterative = IterativeImputer(random_state=42).fit_transform(y_imputable.to_frame())
# Define a simple regression model
from sklearn.linear_model import LinearRegression
def evaluate_model_on_imputed_data(X_imputed, y_imputed, method_name):
    # Split data AFTER imputation for consistent train/test sets
    X_train, X_test, y_train, y_test = train_test_split(X_imputed, y_imputed, test_size=0.2, random_state=42)
    model = LinearRegression()
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    r2 = r2_score(y_test, y_pred)
    print(f"  {method_name} Imputation - RMSE: {rmse:.2f}, R2: {r2:.2f}")
    return rmse, r2
# Evaluate models with different imputation strategies
print("Evaluating Linear Regression model with different imputation methods:")
# Baseline (Listwise Deletion if we want to compare)
X_listwise = original_processed_full_df.dropna()
y_listwise = y.loc[X_listwise.index].dropna() # Ensure target also dropped for same rows
if not X_listwise.empty and not y_listwise.empty:
    evaluate_model_on_imputed_data(X_listwise, y_listwise, "Listwise Deletion")
else:
    print("  Listwise Deletion baseline skipped due to too much missing data.")
# Mean imputed (using Features with NaNs, and Target with NaNs)
evaluate_model_on_imputed_data(imputed_data_frames['Mean'], y_imputed_mean, "Mean Imputation")
# KNN imputed (using Features with NaNs, and Target with NaNs)
evaluate_model_on_imputed_data(imputed_data_frames['KNN'], y_imputed_knn, "KNN Imputation")
# Iterative imputed (using Features with NaNs, and Target with NaNs)
evaluate_model_on_imputed_data(imputed_data_frames['Iterative'], y_imputed_iterative, "Iterative Imputation")
# --- 4. Sensitivity Analysis ---
print("\n--- Evaluation: Sensitivity Analysis ---")
print("The model performance comparison above serves as a basic sensitivity analysis.")
print("In a real scenario, you would repeat this process with different models (e.g., RandomForest, XGBoost)")
print("and different hyperparameters for the imputers to see how robust your conclusions are.")
print("If model performance drastically changes with different imputers, it signals that the imputation strategy is highly influential,")
print("and you should investigate further (e.g., collect more data, apply more sophisticated MNAR handling).")

# Example: Adding a missingness indicator for 'FeatureA'
    df['FeatureA_missing'] = df['FeatureA'].isnull().astype(int)