from . import Spectrum


class Experiment(list):
    def __init__(self, *args, **kwargs):
        super(Experiment, self).__init__(*args, **kwargs)
        self._scan_num_index = {}
        pass

    def add_spectrum(self, spectrum):
        self.append(spectrum)

    def add_ms1(self, spectrum_data, scan_num=-1, retention_time=-1., ):
        spectrum = Spectrum.Spectrum(experiment=self, experiment_spec_id=len(self),
                                     data=spectrum_data,
                                     scan_num=scan_num, retention_time=retention_time,
                                     ms_level=1)
        self.append(spectrum)
        pass

    def add_ms2(self, spectrum_data,
                precursor_mz,
                precursor_charge=-1,
                scan_num=-1,
                retention_time=-1., ):
        spectrum = Spectrum.Spectrum(experiment=self, experiment_spec_id=len(self),
                                     data=spectrum_data,
                                     scan_num=scan_num, retention_time=retention_time,
                                     ms_level=2,
                                     precursor_mz=precursor_mz, precursor_charge=precursor_charge)
        self.append(spectrum)
        pass

    def build_scan_num_index(self):
        for i, spec in enumerate(self):
            self._scan_num_index[spec.raw_scan_num] = i
        pass

    def get_spectrum_by_scan_num(self, scan_num):
        return self[self._scan_num_index[scan_num]]

    def get_xic_from_mz(self, mz_low, mz_high, rt_low=None, rt_high=None):
        pass
