import pickle
import zlib
import bz2

import numpy as np
import sklearn

from Main.MainArguments import check_folder
from Model.FastForest.FastForest import FastForest
from Utils.Multiplecore import run_multiple_process


def _fit_from_fingerprint_to_peak(intensity, model_num, estimator_base, fingerprints):
    estimator = sklearn.base.clone(estimator_base)
    weighted_y = intensity.toarray().flatten()
    y = weighted_y > 0.
    weight = np.zeros_like(weighted_y)
    weight[y] = weighted_y[y]
    sum_y = np.sum(weighted_y == 0.)
    if sum_y > 0.:
        weight[np.bitwise_not(y)] = np.sum(weighted_y) / np.sum(weighted_y == 0.)

    y = np.array(y, np.int8)

    if sum(y) == 0:
        estimator = 0
    elif sum(y) == len(y):
        estimator = 1
    else:
        estimator.fit(fingerprints, y, weight)

    """
    estimator_str = pickle.dumps(estimator)
    estimator_str_old = zlib.compress(estimator_str)
    # """

    # optimize estimator
    estimator_optimized = FastForest()
    estimator_optimized.import_from_forest(estimator)
    estimator_str = pickle.dumps(estimator_optimized)
    estimator_str = zlib.compress(estimator_str)
    return model_num, estimator_str


class EstimatorForPeakModels(object):
    def __init__(self, estimator, model_folder):
        self._estimator_base = estimator
        self._model_folder = check_folder(model_folder) + 'peak_model_'
        pass

    def fit_from_fingerprint_to_peak(self, fingerprint, intensity, model_start_num):
        all_output_nums = intensity.shape[0]
        filename_output = self._model_folder + str(model_start_num)
        fo_data = open(filename_output + '.data', 'wb')
        fo_overview = open(filename_output + '.overview', 'wb')
        fo_overview_list = []

        def _merge_result(all_result, cur_result):
            model_num, estimator_str = cur_result

            fo_data_start = fo_data.tell()
            fo_data.write(estimator_str)
            fo_data_end = fo_data.tell()
            fo_overview_list.append((model_num, fo_data_start, fo_data_end - fo_data_start))
            pass

        run_multiple_process(func_run=_fit_from_fingerprint_to_peak, func_merge=_merge_result,
                             all_para_individual=[(intensity[i], model_start_num + i,) for i in range(all_output_nums)],
                             para_share=(self._estimator_base, fingerprint,))

        pickle.dump(fo_overview_list, fo_overview)
        pass

    @classmethod
    def predict_from_fingerprint_to_peak(cls, fingerprint, estimator_str):
        estimator = pickle.loads(zlib.decompress(estimator_str))
        if type(estimator) is int:
            return [estimator for _ in range(len(fingerprint))]

        result = estimator.predict_proba(fingerprint)
        return result

    @classmethod
    def predict_from_selected_fingerprint_to_peak(cls, fingerprint, selected_mol, estimator_str):
        estimator = pickle.loads(zlib.decompress(estimator_str))
        if type(estimator) is int:
            return [estimator for _ in range(len(selected_mol))]

        result = estimator.predict_proba(fingerprint[selected_mol])

        return result

    @classmethod
    def compress_estimator(cls, estimator_str):
        estimator = pickle.loads(zlib.decompress(estimator_str))

        estimator_optimized = FastForest()
        estimator_optimized.import_from_forest(estimator)

        estimator_str = pickle.dumps(estimator_optimized)
        estimator_str = zlib.compress(estimator_str)
        return estimator_str

    pass
