from sklearn.preprocessing import MultiLabelBinarizer


class DataframePreprocessing:
    
    DEFAULT_TARGET_THEMES = [0, 5, 6, 26, 33, 139, 163, 232, 313, 339, 
                             350, 406, 409, 555, 589, 597, 634, 660, 695, 
                             729, 766, 773, 793, 800, 810, 852, 895, 951, 975]

    OTHER_THEMES_VALUE = 4242

    def __init__(self, df, x_column_name='page_text_extract', y_column_name='tema', target_themes=DEFAULT_TARGET_THEMES, other_themes_value=OTHER_THEMES_VALUE):
        self.x_column_name = x_column_name
        self.y_column_name = y_column_name
        self.other_themes_value = other_themes_value
        self.target_themes = target_themes

        self._group_samples_by_process(df.copy())

        self.target_themes.sort()
        self.labels_freq = {}

        self._set_labels_frequency()
        self.processed_df = self._process_dataframe()

    def _group_samples_by_process(self, df):
        print('Grouping processes...')
        self.df = df.groupby('process_id').apply(self._aggregate)

    def _aggregate(self, series):
        reduced = {}
        series[self.x_column_name].drop_duplicates(inplace=True)
       
        reduced[self.x_column_name] = ' '.join(str(x) for x in series[self.x_column_name].values)
        temas = np.unique(series[self.y_column_name].values)
        reduced[self.y_column_name] = temas[~np.isnan(temas)]
        return pd.Series(reduced, index=[self.x_column_name, *[self.y_column_name]])

    def _remove_rare_samples(self, series, threshold=2):
        if self.labels_freq.get(tuple(series.tolist())) < threshold:
            return False
        return True 

    def _switch_other_themes_values(self, actual_label):
        """
            Replace the values of themes that are not in target themes
            to a unique value
        """

        modified_label = set()
        for theme in actual_label:
            if theme not in self.target_themes:
                modified_label.add(self.other_themes_value)
            else:
                modified_label.add(theme)
        return sorted(modified_label)

    def _normalize_labels(self, series):
        return np.asarray(self._switch_other_themes_values(series.tolist()))
    
    def _set_labels_frequency(self):
        print('Setting labels frequency...')
        for label in self.df[self.y_column_name]:
            normalized_label = tuple(self._switch_other_themes_values(label))

            if not self.labels_freq.get(normalized_label):
                self.labels_freq[normalized_label] = 1
            else:
                self.labels_freq[normalized_label] += 1

    def _get_distinct_themes(self):
        distinct_themes = set()
        for label in self.df['labels_with_others']:
            for theme in label:
                distinct_themes.add(theme)
        self.distinct_themes = list(sorted(distinct_themes))

    def _binarize_labels(self):
        print('Binarizing labels...')
        mlb = MultiLabelBinarizer()
        binarized_columns = mlb.fit_transform(self.df['labels_with_others'].to_numpy())
        
        self._get_distinct_themes()
        columns_names = {ix:binarized_columns[:,i] for i, ix in enumerate(self.distinct_themes)}

        return pd.concat([self.df.reset_index(drop=True), pd.DataFrame(columns_names)], axis=1) 


    def _process_dataframe(self):
        self.df['labels_with_others'] = self.df[self.y_column_name].apply(self._normalize_labels)
        print('Removing rare samples...')

        self.df = self.df[self.df['labels_with_others'].apply(self._remove_rare_samples)]
        return self._binarize_labels()
