import numpy as np
import scipy.sparse


class MassFeatureFinder(object):
    def __init__(self, extract_type):
        self._type = extract_type

        self._feature = {}
        self._feature_num = []
        self._intensity = []
        pass

    def count_features(self, all_mol_spec_info):
        for mol_num, mol_info in enumerate(all_mol_spec_info):
            all_mass = [mol_info['spectrum'].precursor_mass]
            for peak in mol_info['spectrum']:
                if not peak['is_product_ion']:
                    continue

                if peak['calibrated_mass']:
                    all_mass.append(peak['calibrated_mass'])
                    pass
                pass
            pass

        for feature_num, feature_mol_id in enumerate(self._feature_num):
            self._feature_num[feature_num] = (len(set(feature_mol_id)), len(feature_mol_id))

        pass

    def extract_features(self, all_mol_spec_info):
        total_mol_num = len(all_mol_spec_info)
        point_array = [-1 for x in range(len(self._feature_num))]
        intensity_array = [np.zeros(x, dtype=np.float32) for x in self._feature_num]
        mol_num_array = [np.zeros(x, dtype=np.int32) - 1 for x in self._feature_num]

        for mol_num, mol_info in enumerate(all_mol_spec_info):
            for peak in mol_info['spectrum']:
                for cur_feature in peak[self._type]:
                    if cur_feature in self._feature:
                        cur_num = self._feature[cur_feature]
                        # Add the intensity to self._intensity
                        point = point_array[cur_num]
                        if mol_num_array[cur_num][point] == mol_num:
                            intensity_array[cur_num][point] += peak['intensity']
                            pass
                        else:
                            point_array[cur_num] += 1
                            point += 1
                            intensity_array[cur_num][point] += peak['intensity']
                            mol_num_array[cur_num][point] = mol_num
                        pass
                    pass
                pass
            pass

        for mol_num, intensity in zip(mol_num_array, intensity_array):
            self._intensity.append(
                scipy.sparse.csr_matrix((intensity, (mol_num, [0 for x in mol_num])),
                                        shape=(total_mol_num, 1), dtype=np.float32)
            )

        pass

    def select_feature_by_number(self, cutoff):
        new_feature = {}
        new_feature_num = []

        remain_num = {}
        for item_num, (feature_mol_num, feature_spec_num) in enumerate(self._feature_num):
            if feature_mol_num >= cutoff:
                remain_num[item_num] = len(new_feature_num)
                new_feature_num.append(feature_spec_num)

        for feature in self._feature:
            if self._feature[feature] in remain_num:
                new_feature[feature] = remain_num[self._feature[feature]]

        self._feature = new_feature
        self._feature_num = new_feature_num
        print('Total find {} features in extract type: {}'.format(len(new_feature), self._type))
        pass

    def reselect_feature(self, all_mol_spec_info, remain_feature_num):
        new_feature_num = np.zeros(len(self._feature), np.int64)

        for mol_num, mol_info in enumerate(all_mol_spec_info):
            for peak in mol_info['spectrum']:
                for anno_feature in peak['anno']:
                    max_num = 0
                    for cur_feature in anno_feature[self._type_place]:
                        if cur_feature in self._feature:
                            cur_num = self._feature_num[self._feature[cur_feature]]
                            if cur_num > max_num:
                                max_num = cur_num
                    for cur_feature in anno_feature[self._type_place]:
                        if cur_feature in self._feature:
                            cur_feature_id = self._feature[cur_feature]
                            if self._feature_num[cur_feature_id] == max_num:
                                new_feature_num[cur_feature_id] += 1
                    pass

        # Decide cutoff
        if new_feature_num.shape[0] <= remain_feature_num:
            cutoff = 0
        else:
            cutoff = -np.max(np.partition(0 - new_feature_num, remain_feature_num)[:remain_feature_num])

        old_feature_num = self._feature_num
        self._feature_num = []
        new_feature = {}
        for feature, feature_id in self._feature.viewitems():
            if new_feature_num[feature_id] >= cutoff:
                cur_total_feature_num = len(new_feature)
                new_feature[feature] = cur_total_feature_num
                self._feature_num.append(old_feature_num[feature_id])

        self._feature = new_feature

        print('Remain feature: {}, the cutoff for the feature is {}.'.format(len(self._feature), cutoff))
        pass

    def get_intensity_by_feature(self, feature):
        if feature in self._feature:
            feature_num = self._feature[feature]
            return self._intensity[feature_num]
        return None

    def get_all_features(self):
        return self._feature.keys()

    def get_intensity(self):
        return self._intensity

    pass
