%load_ext autoreload
%autoreload 2
from fedimpute.data_prep import load_data, display_data
data, data_config = load_data("codrna")
display_data(data)
print("Data Dimensions: ", data.shape)
print("Data Config:\n", data_config)

%load_ext autoreload
%autoreload 2
from fedimpute.scenario import ScenarioBuilder

scenario_builder = ScenarioBuilder()
scenario_data = scenario_builder.create_simulated_scenario(
    data, data_config, num_clients = 4, dp_strategy='iid-even', ms_scenario='mnar-heter'
)
print('Results Structure (Dict Keys):')
print(list(scenario_data.keys()))
scenario_builder.summarize_scenario()

from fedimpute.execution_environment.imputation.base import BaseMLImputer, ICEImputerMixin
from sklearn.neural_network import MLPRegressor
import numpy as np
from collections import OrderedDict

class MLPICEImputer(BaseMLImputer, ICEImputerMixin):

    def __init__(self):

        super().__init__(name='mlp_mice', model_persistable=False) # it needs two parameters: name of imputer and whether the model is persistable (can be saved to disk), we set it to False because ICE imptutation is not persistable

        self.imp_models = [] # list of imputation models (each for a feature)
        self.min_values = None # min values of features used for clipping
        self.max_values = None # max values of features used for clipping
        self.seed = None # seed for randomization
        self.fit_res_history = {} # fit results history

    def initialize(
            self, X: np.array, missing_mask: np.array, data_utils: dict, params: dict, seed: int
    ) -> None:
        """
        Initialize imputer - statistics imputation models etc.

        Args:
            X: data with intial imputed values
            missing_mask: missing mask of data
            data_utils: data utils dictionary - contains information about data
            params: params for initialization
            seed: int - seed for randomization
        """

        # initialized imputation models (from sklearn's MLPRegressor (fully connected neural network))
        self.imp_models = []
        for i in range(data_utils['n_features']):
            estimator = MLPRegressor(hidden_layer_sizes=(100, 100), max_iter=1000)
            X_train = X[:, np.arange(X.shape[1]) != i][0:10]
            y_train = X[:, i][0:10]
            estimator.fit(X_train, y_train)
            self.imp_models.append(estimator)

        # initialize min max values for a clipping threshold (this method is defined in `ICEImputerMixin`)
        self.min_values, self.max_values = self.get_clip_thresholds(data_utils)
        self.seed = seed

    def get_imp_model_params(self, params: dict) -> OrderedDict:
        """
        Return model parameters

        Args:
            params: dict contains parameters for get_imp_model_params

        Returns:
            OrderedDict - model parameters dictionary
        """
        # This method is used to get the parameters of the imputation model for a given feature
        # get feature index of imputation models
        feature_idx = params['feature_idx']
        imp_model = self.imp_models[feature_idx]

        # get parameters from sklearn model
        coefs = imp_model.coefs_
        intercept = imp_model.intercepts_

        # convert parameters to a dictionary (we need to convert parameters to ordered dictionary as required by `BaseMLImputer`)
        parameters = {}
        for i in range(len(coefs)):
            parameters[f'coef_{i}'] = coefs[i]
            parameters[f'intercept_{i}'] = intercept[i]

        return OrderedDict(parameters)

    def set_imp_model_params(self, updated_model_dict: OrderedDict, params: dict) -> None:
        """
        Set model parameters

        Args:
            updated_model_dict: global model parameters dictionary
            params: parameters for set parameters function
        """
        # This method is used to set the parameters of the imputation model for a given feature (update models)
        # get feature index of imputation models
        feature_idx = params['feature_idx']
        imp_model = self.imp_models[feature_idx]

        # set parameters to sklearn model
        coefs = []
        intercepts = []
        for i in range(len(imp_model.coefs_)):
            coefs.append(updated_model_dict[f'coef_{i}'])
            intercepts.append(updated_model_dict[f'intercept_{i}'])
        imp_model.coefs_ = coefs
        imp_model.intercepts_ = intercepts

    def fit(self, X: np.array, y: np.array, missing_mask: np.array, params: dict) -> dict:
        """
        Fit imputer to train local imputation models

        Args:
            X: np.array - float numpy array features
            y: np.array - target
            missing_mask: np.array - missing mask
            params: parameters for local training
        """
        # This method is used to fit the imputation model for a given feature
        # get complete data of the feature
        feature_idx = params['feature_idx']
        row_mask = missing_mask[:, feature_idx]
        X_train = X[~row_mask][:, np.arange(X.shape[1]) != feature_idx]
        y_train = X[~row_mask][:, feature_idx]

        # fit MLP imputation models
        estimator = self.imp_models[feature_idx]
        estimator.fit(X_train, y_train)
        y_pred = estimator.predict(X_train)
        loss = np.mean((y_pred - y_train) ** 2)
        self.fit_res_history[feature_idx].append({
            'loss': loss,
            'sample_size': X_train.shape[0]
        })

        return {
            'loss': loss,
            'sample_size': X_train.shape[0]
        }

    def impute(self, X: np.array, y: np.array, missing_mask: np.array, params: dict) -> np.ndarray:
        """
        Impute missing values using an imputation model

        Args:
            X (np.array): numpy array of features
            y (np.array): numpy array of target
            missing_mask (np.array): missing mask
            params (dict): parameters for imputation

        Returns:
            np.ndarray: imputed data - numpy array - same dimension as X
        """
        # This method is used to impute the missing values using the imputation model for a given feature
        # get incomplete data of the feature
        feature_idx = params['feature_idx']
        row_mask = missing_mask[:, feature_idx]
        if np.sum(row_mask) == 0:
            return X

        # impute missing values
        X_test = X[row_mask][:, np.arange(X.shape[1]) != feature_idx]
        estimator = self.imp_models[feature_idx]
        imputed_values = estimator.predict(X_test)
        imputed_values = np.clip(imputed_values, self.min_values[feature_idx], self.max_values[feature_idx])
        X[row_mask, feature_idx] = np.squeeze(imputed_values)

        return X

    def get_fit_res(self, params: dict) -> dict:

        # This method is used to get fit results for a given feature from the fit history
        try:
            feature_idx = params['feature_idx']
        except KeyError:
            raise ValueError("Feature index not found in params")

        return self.fit_res_history[feature_idx][-1]

%load_ext autoreload
%autoreload 2
from fedimpute.execution_environment import FedImputeEnv

env = FedImputeEnv(debug_mode=False)
env.register.register_imputer(
    name = 'mlp_mice',                     # name of we give to the new imputer
    imputer = MLPICEImputer,              # the class of the new imputer we just developed
    workflow = 'ice',                     # because it is ICE imputation, we use 'ice' workflow
    fed_strategy = ['local', 'fedmice']   # we support local and fedmice strategy for this imputer
)

# then we can use the new imputer in the environment and run the federated imputation
env.configuration(imputer = 'mlp_mice', fed_strategy='local', workflow_params={'log_metric': None})
env.setup_from_scenario_builder(scenario_builder = scenario_builder, verbose=1)
env.show_env_info()
env.run_fed_imputation(verbose=2)