from Library.Parameter import Parameter


def _merge_peak_cluster(raw_spectrum, peak_start, peak_end):
    intensity_sum = 0.
    mz_intensity_sum = 0.
    for i in range(peak_start, peak_end):
        intensity_sum += raw_spectrum[i][1]
        mz_intensity_sum += raw_spectrum[i][0] * raw_spectrum[i][1]
        pass

    return mz_intensity_sum / intensity_sum, intensity_sum


def peak_picker(raw_spectrum):
    new_spectrum = []

    # Find peak clusters.
    peak_num_cur = 0
    while peak_num_cur < len(raw_spectrum) - 1:
        peak_num_next = peak_num_cur + 1
        peak_cur = raw_spectrum[peak_num_cur]
        peak_next = raw_spectrum[peak_num_next]
        mass_error = 1. * peak_cur[0] / 1e6 * Parameter.ms2_ppm

        if peak_next[0] - peak_cur[0] < mass_error:
            # There is peak clusters.
            cluster_num_start = peak_num_cur

            # Find cluster cores.
            while (peak_next[1] > peak_cur[1]) and \
                    (peak_next[0] - peak_cur[0] < mass_error) and \
                    peak_num_next < len(raw_spectrum):
                peak_num_cur += 1
                peak_num_next += 1
                if peak_num_next == len(raw_spectrum):
                    break
                peak_cur = raw_spectrum[peak_num_cur]
                peak_next = raw_spectrum[peak_num_next]

            # Find cluster end.
            while (peak_next[1] <= peak_cur[1]) and \
                    (peak_next[0] - peak_cur[0] < mass_error) and \
                    peak_num_next < len(raw_spectrum):
                peak_num_cur += 1
                peak_num_next += 1
                if peak_num_next == len(raw_spectrum):
                    break
                peak_cur = raw_spectrum[peak_num_cur]
                peak_next = raw_spectrum[peak_num_next]
                pass
            cluster_num_end = peak_num_cur + 1

            new_peak = _merge_peak_cluster(raw_spectrum, cluster_num_start, cluster_num_end)
            new_spectrum.append(new_peak)
            pass
        else:
            new_spectrum.append(peak_cur)
            pass

        # Move to the next
        peak_num_cur += 1

    new_spectrum.append(raw_spectrum[-1])
    return new_spectrum


class Spectrum(list):
    def __init__(self, mass_intensity, is_normalized=False, need_refine=True):
        """
        Build Spectrum object from mass and intensity.
        :param mass_intensity: [[m/z,intensity],[],...]
        """
        super(Spectrum, self).__init__()

        self.precursor_mass = None
        if is_normalized:
            normalized_intensity = mass_intensity
        else:
            normalized_intensity = self.normalize_spec_intensity_by_sum(mass_intensity)

        # For each spectrum, refine peaks
        intensity_cutoff = 0.
        normalized_intensity.sort(key=lambda x: x[0], reverse=False)

        if need_refine:
            refined_intensity = peak_picker(normalized_intensity)
        else:
            refined_intensity = normalized_intensity

        if Parameter.remain_peaks and len(refined_intensity) > Parameter.remain_peaks:
            refined_intensity.sort(key=lambda x: x[1], reverse=True)
            intensity_cutoff = refined_intensity[Parameter.remain_peaks - 1][1]
            pass
        normalized_intensity = refined_intensity

        for mass, intensity in normalized_intensity:
            if intensity >= intensity_cutoff and intensity > 0.:
                s = [mass, intensity]
                self.append(s)
        if len(self) == 0:
            self.append([0.0, 0.0])

        # sort spec first
        self.sort(key=lambda x: x[0])

    @staticmethod
    def normalize_spec_intensity_by_sum(spec):
        """
        Normalized by sum of all the intensity
        :param spec: [[m/z,intensity],[],...]
        :return: [[m/z,intensity],[],...]
        """
        new_spec = []
        intensity_sum = sum([x[1] for x in spec]) / 100.0
        if intensity_sum != 0.0:
            for mz, intensity in spec:
                new_spec.append([mz, intensity / intensity_sum])
        return new_spec
