import ctypes
import pickle

import gc
import numpy as np
import scipy.sparse

from Main.Parameter import Parameter
from Model.EstimatorForPeakModels import EstimatorForPeakModels
from Utils import FilenameGenerator
from Utils.Multiplecore import mp_raw_array_data, run_multiple_process


def _func_worker_score_with_peak_model(func_run, func_para_share, q_input, q_output):
    feature_data = open(FilenameGenerator.merged_peak_model_data(Parameter.model), 'rb')
    all_mol_fingerprints = func_para_share[0]
    while not q_input.empty():
        try:
            q_item = q_input.get(block=True, timeout=1)
        except:
            continue
        if q_item is None:
            break

        try:
            i, para = q_item
            file_offset, str_len = para
            feature_data.seek(file_offset)
            estimator_str = feature_data.read(str_len)
            possibility = EstimatorForPeakModels.predict_from_fingerprint_to_peak(all_mol_fingerprints, estimator_str)
            result = possibility
        except Exception as e:
            result = None
            print(e)
        q_output.put((i, result))
    return


def _func_worker_score_part_with_peak_model(func_run, func_para_share, q_input, q_output):
    feature_data = open(FilenameGenerator.merged_peak_model_data(Parameter.model), 'rb')
    all_mol_fingerprints = func_para_share[0]
    while not q_input.empty():
        try:
            q_item = q_input.get(block=True, timeout=1)
        except:
            continue

        if q_item is None:
            break

        i, para = q_item

        try:
            # if 1:
            (file_offset, str_len), wanted_mol_id = para
            feature_data.seek(file_offset)
            estimator_str = feature_data.read(str_len)
            if len(wanted_mol_id) > 0:
                possibility = EstimatorForPeakModels.predict_from_selected_fingerprint_to_peak(
                    all_mol_fingerprints, wanted_mol_id, estimator_str)
            else:
                possibility = np.array([], np.float32)
            result = possibility
        except Exception as e:
            result = None
            print(e)
        q_output.put((i, result))
    return


class ModelBuilder(object):
    def __init__(self):
        self._feature = {}
        self._feature_data_info = []
        self._feature_data = None

        self._intensity = []
        self._fingerprint = []
        pass

    def fit_peak_model(self, estimator):
        model_num = self._intensity.shape[1]
        batch_size = model_num // Parameter.total_part + 1
        model_start_num = batch_size * Parameter.cal_part
        model_end_num = batch_size * (Parameter.cal_part + 1)
        if model_end_num > model_num:
            model_end_num = model_num

        fingerprint = mp_raw_array_data(self._fingerprint, ctypes.c_int8)
        intensity = self._intensity.T.tocsr()
        intensity = intensity[model_start_num:model_end_num, :]

        self_keys = list(self.__dict__.keys())
        for att in self_keys:
            delattr(self, att)
        gc.collect()

        estimator.fit_from_fingerprint_to_peak(fingerprint, intensity, model_start_num)
        pass

    def score_mol_info_with_peak_model(self, all_mol_intensity, all_mol_features, all_mol_fingerprints):
        all_wanted_features = []
        for item in all_mol_features:
            for item2 in item:
                if item2:
                    all_wanted_features += item2
        all_wanted_features = list(set(all_wanted_features))
        all_wanted_features_dict = {x: i for i, x in enumerate(all_wanted_features)}

        wanted_data_info = [self._feature_data_info[feature_num] for feature_num in all_wanted_features]
        if len(wanted_data_info) == 0:
            return [0 for _ in range(len(all_mol_intensity))]

        all_mol_fingerprints = np.array(all_mol_fingerprints, np.float32)
        all_mol_fingerprints = mp_raw_array_data(all_mol_fingerprints)
        all_mol_intensity_predict_wanted = run_multiple_process(None, None,
                                                                func_worker=_func_worker_score_with_peak_model,
                                                                para_share=(all_mol_fingerprints,),
                                                                all_para_individual=wanted_data_info)

        all_peak_score = []
        for i in range(len(all_mol_features)):
            mol_intensity_real = all_mol_intensity[i]

            cur_score = 0.
            mol_feature = all_mol_features[i]
            for peak_num, peak_feature in enumerate(mol_feature):
                peak_score = 0.
                for model_num in peak_feature:
                    cur_peak_score = all_mol_intensity_predict_wanted[all_wanted_features_dict[model_num]][i]
                    peak_score = cur_peak_score if cur_peak_score > peak_score else peak_score

                cur_score += mol_intensity_real[peak_num] * peak_score

            all_peak_score.append(cur_score)

        score = []
        for mol_num, mol_intensity in enumerate(all_mol_intensity):
            peak_score = all_peak_score[mol_num]
            total_intensity = np.sum(mol_intensity)
            if total_intensity > 0.:
                score.append(peak_score / total_intensity)
            else:
                score.append(0.)
        return score

    def predict_probability_with_peak_model(self, all_fingerprints, all_wanted_ff):
        all_probability = run_multiple_process(
            None, None,
            func_worker=_func_worker_score_part_with_peak_model,
            para_share=(all_fingerprints,),
            all_para_individual=list(zip(self._feature_data_info, all_wanted_ff)))
        return all_probability

    def convert_mol_info_to_feature_intensity_matrix(self, all_mol_info):
        vector_mol_num = []
        vector_feature_num = []
        vector_intensity = []
        for mol_num, mol_info in enumerate(all_mol_info):
            for peak in mol_info['spectrum']:
                if not peak['is_product_ion']:
                    continue

                for feature in peak['ff']:
                    if feature in self._feature:
                        vector_mol_num.append(mol_num)
                        vector_feature_num.append(self._feature[feature])
                        vector_intensity.append(peak['intensity'])

        feature_intensity_matrix = scipy.sparse.csr_matrix(
            (vector_intensity, (vector_mol_num, vector_feature_num)),
            shape=(len(all_mol_info), self.get_total_feature_num()))
        return feature_intensity_matrix

    def convert_mol_info_to_peak_intensity_features(self, all_mol_info):
        all_peak_features = []
        all_peak_intensity = []
        for mol_info in all_mol_info:
            cur_all_peak_features = []
            cur_all_peak_intensity = []
            for peak in mol_info['spectrum']:
                if not peak['is_product_ion']:
                    continue
                cur_peak = []

                for feature in peak['ff']:
                    if feature in self._feature:
                        cur_peak.append(self._feature[feature])
                cur_all_peak_features.append(cur_peak)
                cur_all_peak_intensity.append(peak['intensity'])
            all_peak_features.append(cur_all_peak_features)
            all_peak_intensity.append(cur_all_peak_intensity)

        return all_peak_intensity, all_peak_features

    def get_total_feature_num(self):
        return len(self._feature)

    def get_all_features(self):
        return self._feature

    def add_intensity(self, formula_feature_finder):
        for feature in formula_feature_finder.get_all_features():
            self._feature[feature] = len(self._intensity)
            self._intensity.append(formula_feature_finder.get_intensity_by_feature(feature))
            pass
        pass

    def add_fingerprint(self, fingerprint):
        self._fingerprint = fingerprint

    def add_data_overview(self, overview):
        self._feature_data_info = overview

    def save_to_file(self, filename):
        fingerprint, intensity = self._fingerprint, self._intensity
        del self._fingerprint
        del self._intensity
        pickle.dump(self, open(filename, 'wb'))

        np.save(filename + '.fingerprint.npy', fingerprint)
        intensity = scipy.sparse.hstack(intensity)
        pickle.dump(intensity, open(filename + '.intensity.npy', 'wb'), 1)

        self._fingerprint, self._intensity = fingerprint, intensity
        pass

    @classmethod
    def load_from_file(cls, filename, load_training_data=True):
        model = pickle.load(open(filename, 'rb'))
        if load_training_data:
            model._fingerprint = np.load(filename + '.fingerprint.npy')
            model._intensity = np.load(filename + '.intensity.npy')
        return model
