from numpy import random
from assets.classification.classification_visualizations import plot_multiple_confusion_matrix
from mlapp.managers import ModelManager, pipeline
from sklearn.model_selection import StratifiedKFold, train_test_split
from mlapp.utils.automl import AutoMLPandas


class ClassificationModelManager(ModelManager):
    @pipeline
    def train_model(self, data):
        """
        This function should create a model  and train the model with the features generated by the function build_and_store_features in the Task_Manager Class
        :param data: pandas DataFrame that contains X, y for training the model
        :return local_model_path: string: pickle file path
        :return img_file_names: list of strings: of img file names paths
        :return model_results: dict: contains info of model such as accuracy, configs, etc.
        :return predictions_df: pandas DataFrame: estimates+trend+seasonality of the train period
        """
        # -------------------------------- preparation for train --------------------------------
        variable_to_predict = self.model_settings['variable_to_predict']

        # Sampling the data
        data = self._sample_data(data)

        # Preparing the data
        X_train, X_test, y_train, y_test = train_test_split(
            data.drop(variable_to_predict, axis=1), data[variable_to_predict],
            train_size=self.model_settings.get('train_percent', 0.8), stratify=data[variable_to_predict])

        # Run AutoML
        auto_ml_config = self.model_settings.get('auto_ml', {})
        estimator_family = auto_ml_config.get('estimator', 'binary')
        cv_splitter = StratifiedKFold(n_splits=5, shuffle=True, random_state=random.seed(0))
        result = AutoMLPandas(
            estimator_family, **auto_ml_config.get(estimator_family, {}),
            feature_selection=auto_ml_config.get('feature_selection')).run(
            X_train, y_train, X_test, y_test, cv=cv_splitter)

        # Print report
        result.print_report(full=False)

        # Save results
        self.save_automl_result(result)

        # Add confusion matrix plot
        self.save_image('confusion_matrix', plot_multiple_confusion_matrix(
            result.get_train_predictions(), y_train, result.get_test_predictions(), y_test))

        # Save probability thresholds
        proba_threshold = result.predict_proba_by_threshold(
            X_train, X_test, y_train, y_test, self.model_settings.get('predict_proba_threshold', 0.5))
        self.save_metadata('predict_proba_by_threshold', proba_threshold)

        # Add predictions
        self.add_predictions(X_train.index, result.get_train_predictions(), y_train, 'TRAIN')
        self.add_predictions(X_test.index, result.get_test_predictions(), y_test, 'TEST')

    @pipeline
    def forecast(self, data_df):
        result = self.get_automl_result()

        # get best model
        model = result.get_best_model()

        # get selected features from feature selection algorithm
        selected_features_names = result.get_selected_features()

        # filter data according to selected features
        filtered_data = data_df[selected_features_names]

        # predicting
        predictions = model.predict(filtered_data)

        self.add_predictions(primary_keys_columns=data_df.index, y_hat=predictions, y=[], prediction_type='forecast')

    def _sample_data(self, data):
        sample_data = self.model_settings.get('down_sample_method', None)
        if sample_data is not None and sample_data.get('flag', False):
            n_samples = sample_data.get('n_samples', data.shape[0])
            seed = sample_data.get('seed', 1500)
            if isinstance(n_samples, int) and n_samples <= data.shape[0]:
                return data.sample(n=n_samples, random_state=seed)
            elif isinstance(n_samples, float) and (0.0 < n_samples <= 1.0):
                n_samples = int(data.shape[0] * n_samples)
                return data.sample(n=n_samples, random_state=seed)
            else:
                print('Warning: "n_samples" can be an Integer Type that <= number of data samples or Float type (ratio)'
                      ' that bigger than 0 and smaller or equal to 1 (0 < n_sample <= 1) \n Taking Full Data!')
            print('After data sampling: (' + repr(data.shape[0]) + ',' + repr(data.shape[1]) + ')')

        return data

    def refit(self, data):
        raise NotImplementedError()
