import copy

from FormulaClass.FormulaCalculator import *


class SpectrumAnnotator(object):
    def __init__(self, mol):
        self._mol = mol
        self._formula = mol.get_formula()
        self._precursor_mass = mol.get_ionized_mass()
        pass

    def annotate_spectrum(self, spec, match_large_error=False):
        spec.precursor_mass = self._mol.get_ionized_mass()
        annotated_peak_num = 0
        for peak in spec:
            peak['anno'] = self._assign_formula_to_peak(peak, match_large_error)
            peak['calibrated_mass'] = self._recalibrate_mass(peak)
            if peak['anno'] and peak['is_product_ion']:
                annotated_peak_num += 1

        if not Parameter.no_iso:
            self._label_isotopic_peak(spec)
            self._remove_isotopic_peak(spec)

        return annotated_peak_num

    @staticmethod
    def _remove_isotopic_peak(spec):
        for peak in spec[::-1]:
            iso_father_peak = peak['iso_father_peak']
            if iso_father_peak is not None:
                spec[iso_father_peak]['intensity'] += peak['intensity']
                peak['is_product_ion'] = False
            pass
        return spec

    @staticmethod
    def _label_isotopic_peak(spec):
        for peak in spec:
            peak['iso_father_peak'] = None

        for peak_num, peak in enumerate(spec):
            if not peak['is_product_ion']:
                continue

            for anno in peak['anno']:
                isotopic_info = anno[0].get_isotopic_info()

                for iso_info in isotopic_info:
                    delta_mass = iso_info[0] / 1e6 * Parameter.ms2_ppm
                    min_mass = iso_info[0] - delta_mass
                    max_mass = iso_info[0] + delta_mass
                    max_intensity = iso_info[1] * peak['intensity']
                    for iso_peak_num in range(peak_num + 1, len(spec)):
                        if min_mass <= spec[iso_peak_num]['mass']:
                            if spec[iso_peak_num]['mass'] <= max_mass:
                                if spec[iso_peak_num]['intensity'] <= max_intensity:
                                    if spec[iso_peak_num]['iso_father_peak'] is None:
                                        spec[iso_peak_num]['iso_father_peak'] = peak_num
                                    else:
                                        if peak['intensity'] > spec[spec[iso_peak_num]['iso_father_peak']]['intensity']:
                                            spec[iso_peak_num]['iso_father_peak'] = peak_num
                                        else:
                                            # Ignore when the another explain of the peak has high intensity.
                                            pass
                                        pass
                                    pass
                                else:
                                    # Ignore when the intensity of this peak is higher than expected.
                                    pass
                                pass
                            else:
                                break
                            pass
                    pass
                pass
        return spec

    @staticmethod
    def extract_info(spectrum):
        for peak in spectrum:
            peak['frag_formula'] = []
            peak['root_loss'] = []
            if ('r_group' in peak) and (peak['r_group']):
                for r_group in peak['r_group']:
                    for frag, root_loss in r_group:
                        peak['frag_formula'].append(frag)
                        peak['root_loss'].append(root_loss)
                del peak['r_group']
            else:
                for frag, root_loss in peak['anno']:
                    peak['frag_formula'].append(frag)
                    peak['root_loss'].append(root_loss)

            del peak['anno']
        pass

    @staticmethod
    def _cal_frag_formula_root_loss(peak):
        peak['frag_formula'] = []
        peak['root_loss'] = []
        for possible_anno in peak['anno']:
            remain, root_loss = possible_anno
            peak['frag_formula'].append(remain)
            peak['root_loss'].append(root_loss)
        pass

    @staticmethod
    def _recalibrate_mass(peak):
        all_possible_mass = []
        for formula in peak['anno']:
            cur_mass = formula[0].get_ionized_mass()
            all_possible_mass.append(cur_mass)

        if all_possible_mass:
            return sum(all_possible_mass) / len(all_possible_mass)
        else:
            return None

    def _assign_formula_to_peak(self, peak, match_large_error):
        if peak['mass'] > self._precursor_mass - 0.5:
            peak['is_product_ion'] = False
            return []
        else:
            peak['is_product_ion'] = True

        # First, use ms2 ppm to map _formula
        mass_error = peak['mass'] * Parameter.ms2_ppm * 1e-6
        result_formula = self._map_mass_to_formula(mass_error, peak)
        if len(result_formula) > 0:
            return result_formula
        if not match_large_error:
            return []

        while 1:
            # If use ms2 ppm failed, increase ppm for 50 every step.
            mass_error += peak.mass * Parameter.ms2_ppm_step * 1e-6
            if mass_error > Parameter.ms2_max_delta:
                return []

            result_formula = self._map_mass_to_formula(mass_error, peak)
            if len(result_formula) > 0:
                mass_error += peak.mass * Parameter.ms2_ppm_step * 1e-6
                result2 = self._map_mass_to_formula(mass_error, peak)
                # check is the same
                for item2 in result2:
                    for item1 in result_formula:
                        if item2[0] == item1[0]:
                            break
                    else:
                        result_formula.append(item2)
                return result_formula
        pass

    def _map_mass_to_formula(self, mass_error, peak):
        # Label precursor ion
        if peak['mass'] >= self._precursor_mass - mass_error:
            return [[copy.deepcopy(self._formula), MolecularFormula()]]

        all_possible_formula = mass_to_formula(peak['mass'], mass_error, self._formula)
        cur_result = []
        # Check saturation.
        for f_remain in all_possible_formula:
            is_remain_right = test_unsaturation(f_remain)
            if not is_remain_right:
                continue

            f_loss = self._formula.sub_not_allow_zero(f_remain)
            if f_loss is None:
                continue

            is_loss_right = test_unsaturation(f_loss)
            if not is_loss_right:
                continue

            cur_result.append([f_remain, f_loss])
            pass
        return cur_result

    pass
