1__author__ = "Yuri E. Corilo"
  2__date__ = "Nov 11, 2019"
  3
  4from copy import deepcopy
  5from io import BytesIO
  6from pathlib import Path
  7
  8import chardet
  9from bs4 import BeautifulSoup
 10from pandas import read_csv, read_excel, read_pickle
 11from pandas.core.frame import DataFrame
 12from s3path import S3Path
 13
 14from corems.encapsulation.constant import Labels
 15from corems.encapsulation.factory.parameters import default_parameters
 16from corems.encapsulation.factory.processingSetting import DataInputSetting
 17from corems.encapsulation.input.parameter_from_json import (
 18    load_and_set_parameters_class,
 19    load_and_set_parameters_ms,
 20    load_and_set_toml_parameters_class,
 21)
 22
 23
 24class MassListBaseClass:
 25    """The MassListBaseClass object reads mass list data types and returns the mass spectrum obj
 26
 27    Parameters
 28    ----------
 29    file_location : Path or S3Path
 30        Full data path.
 31    isCentroid : bool, optional
 32        Determines the mass spectrum data structure. If set to True, it assumes centroid mode. If set to False, it assumes profile mode and attempts to peak pick. Default is True.
 33    analyzer : str, optional
 34        The analyzer used for the mass spectrum. Default is 'Unknown'.
 35    instrument_label : str, optional
 36        The label of the instrument used for the mass spectrum. Default is 'Unknown'.
 37    sample_name : str, optional
 38        The name of the sample. Default is None.
 39    header_lines : int, optional
 40        The number of lines to skip in the file, including the column labels line. Default is 0.
 41    isThermoProfile : bool, optional
 42        Determines the number of expected columns in the file. If set to True, only m/z and intensity columns are expected. Signal-to-noise ratio (S/N) and resolving power (RP) will be calculated based on the data. Default is False.
 43    headerless : bool, optional
 44        If True, assumes that there are no headers present in the file (e.g., a .xy file from Bruker) and assumes two columns: m/z and intensity. Default is False.
 45
 46    Attributes
 47    ----------
 48    parameters : DataInputSetting
 49        The data input settings for the mass spectrum.
 50    data_type : str
 51        The type of data in the file.
 52    delimiter : str
 53        The delimiter used to read text-based files.
 54
 55    Methods
 56    -------
 57    * set_parameter_from_toml(parameters_path). Sets the data input settings from a TOML file.
 58    * set_parameter_from_json(parameters_path). Sets the data input settings from a JSON file.
 59    * get_dataframe(). Reads the file and returns the data as a pandas DataFrame.
 60    * load_settings(mass_spec_obj, output_parameters). Loads the settings for the mass spectrum.
 61    * get_output_parameters(polarity, scan_index=0). Returns the output parameters for the mass spectrum.
 62    * clean_data_frame(dataframe). Cleans the data frame by removing columns that are not in the expected columns set.
 63
 64    """
 65
 66    def __init__(
 67        self,
 68        file_location: Path | S3Path,
 69        isCentroid: bool = True,
 70        analyzer: str = "Unknown",
 71        instrument_label: str = "Unknown",
 72        sample_name: str = None,
 73        header_lines: int = 0,
 74        isThermoProfile: bool = False,
 75        headerless: bool = False,
 76    ):
 77        self.file_location = (
 78            Path(file_location) if isinstance(file_location, str) else file_location
 79        )
 80
 81        if not self.file_location.exists():
 82            raise FileExistsError("File does not exist: %s" % file_location)
 83
 84        # (newline="\n")
 85
 86        self.header_lines = header_lines
 87
 88        if isThermoProfile:
 89            self._expected_columns = {Labels.mz, Labels.abundance}
 90
 91        else:
 92            self._expected_columns = {
 93                Labels.mz,
 94                Labels.abundance,
 95                Labels.s2n,
 96                Labels.rp,
 97            }
 98
 99        self._delimiter = None
100
101        self.isCentroid = isCentroid
102
103        self.isThermoProfile = isThermoProfile
104
105        self.headerless = headerless
106
107        self._data_type = None
108
109        self.analyzer = analyzer
110
111        self.instrument_label = instrument_label
112
113        self.sample_name = sample_name
114
115        self._parameters = deepcopy(DataInputSetting())
116
117    @property
118    def parameters(self):
119        return self._parameters
120
121    @parameters.setter
122    def parameters(self, instance_DataInputSetting):
123        self._parameters = instance_DataInputSetting
124
125    def set_parameter_from_toml(self, parameters_path):
126        self._parameters = load_and_set_toml_parameters_class(
127            "DataInput", self.parameters, parameters_path=parameters_path
128        )
129
130    def set_parameter_from_json(self, parameters_path):
131        self._parameters = load_and_set_parameters_class(
132            "DataInput", self.parameters, parameters_path=parameters_path
133        )
134
135    @property
136    def data_type(self):
137        return self._data_type
138
139    @data_type.setter
140    def data_type(self, data_type):
141        self._data_type = data_type
142
143    @property
144    def delimiter(self):
145        return self._delimiter
146
147    @delimiter.setter
148    def delimiter(self, delimiter):
149        self._delimiter = delimiter
150
151    def encoding_detector(self, file_location) -> str:
152        """
153        Detects the encoding of a file.
154
155        Parameters
156        --------
157        file_location : str
158            The location of the file to be analyzed.
159
160        Returns
161        --------
162        str
163            The detected encoding of the file.
164        """
165
166        with file_location.open("rb") as rawdata:
167            result = chardet.detect(rawdata.read(10000))
168        return result["encoding"]
169
170    def set_data_type(self):
171        """
172        Set the data type and delimiter based on the file extension.
173
174        Raises
175        ------
176        TypeError
177            If the data type could not be automatically recognized.
178        """
179        if self.file_location.suffix == ".csv":
180            self.data_type = "txt"
181            self.delimiter = ","
182        elif self.file_location.suffix == ".txt":
183            self.data_type = "txt"
184            self.delimiter = "\t"
185        elif self.file_location.suffix == ".tsv":
186            self.data_type = "txt"
187            self.delimiter = "\t"
188        elif self.file_location.suffix == ".xlsx":
189            self.data_type = "excel"
190        elif self.file_location.suffix == ".ascii":
191            self.data_type = "txt"
192            self.delimiter = "  "
193        elif self.file_location.suffix == ".pkl":
194            self.data_type = "dataframe"
195        elif self.file_location.suffix == ".pks":
196            self.data_type = "pks"
197            self.delimiter = "          "
198            self.header_lines = 9
199        elif self.file_location.suffix == ".xml":
200            self.data_type = "xml"
201            # self.delimiter = None
202            # self.header_lines = None
203        elif self.file_location.suffix == ".xy":
204            self.data_type = "txt"
205            self.delimiter = " "
206            self.header_lines = None
207        else:
208            raise TypeError(
209                "Data type could not be automatically recognized for %s; please set data type and delimiter manually."
210                % self.file_location.name
211            )
212
213    def get_dataframe(self) -> DataFrame:
214        """
215        Get the data as a pandas DataFrame.
216
217        Returns
218        -------
219        pandas.DataFrame
220            The data as a pandas DataFrame.
221
222        Raises
223        ------
224        TypeError
225            If the data type is not supported.
226        """
227
228        if not self.data_type or not self.delimiter:
229            self.set_data_type()
230
231        if isinstance(self.file_location, S3Path):
232            data = BytesIO(self.file_location.open("rb").read())
233        else:
234            data = self.file_location
235
236        if self.data_type == "txt":
237            if self.headerless:
238                dataframe = read_csv(
239                    data,
240                    skiprows=self.header_lines,
241                    delimiter=self.delimiter,
242                    header=None,
243                    names=["m/z", "I"],
244                    encoding=self.encoding_detector(self.file_location),
245                    engine="python",
246                )
247            else:
248                dataframe = read_csv(
249                    data,
250                    skiprows=self.header_lines,
251                    delimiter=self.delimiter,
252                    encoding=self.encoding_detector(self.file_location),
253                    engine="python",
254                )
255
256        elif self.data_type == "pks":
257            names = [
258                "m/z",
259                "I",
260                "Scaled Peak Height",
261                "Resolving Power",
262                "Frequency",
263                "S/N",
264            ]
265            clean_data = []
266            with self.file_location.open() as maglabfile:
267                for i in maglabfile.readlines()[8:-1]:
268                    clean_data.append(i.split())
269            dataframe = DataFrame(clean_data, columns=names)
270
271        elif self.data_type == "dataframe":
272            dataframe = read_pickle(data)
273
274        elif self.data_type == "excel":
275            dataframe = read_excel(data)
276
277        elif self.data_type == "xml":
278            dataframe = self.read_xml_peaks(data)
279
280        else:
281            raise TypeError("Data type %s is not supported" % self.data_type)
282
283        return dataframe
284
285    def load_settings(self, mass_spec_obj, output_parameters):
286        """
287        #TODO loading output parameters from json file is not functional
288        Load settings from a JSON file and apply them to the given mass_spec_obj.
289
290        Parameters
291        ----------
292        mass_spec_obj : MassSpec
293            The mass spectrum object to apply the settings to.
294
295        """
296        import json
297        import warnings
298
299        settings_file_path = self.file_location.with_suffix(".json")
300
301        if settings_file_path.exists():
302            self._parameters = load_and_set_parameters_class(
303                "DataInput", self._parameters, parameters_path=settings_file_path
304            )
305
306            load_and_set_parameters_ms(
307                mass_spec_obj, parameters_path=settings_file_path
308            )
309
310        else:
311            warnings.warn(
312                "auto settings loading is enabled but could not locate the file:  %s. Please load the settings manually"
313                % settings_file_path
314            )
315
316        # TODO this will load the setting from SettingCoreMS.json
317        # coreMSHFD5 overrides this function to import the attrs stored in the h5 file
318        # loaded_settings = {}
319        # loaded_settings['MoleculaSearch'] = self.get_scan_group_attr_data(scan_index,  time_index, 'MoleculaSearchSetting')
320        # loaded_settings['MassSpecPeak'] = self.get_scan_group_attr_data(scan_index,  time_index, 'MassSpecPeakSetting')
321
322        # loaded_settings['MassSpectrum'] = self.get_scan_group_attr_data(scan_index, time_index, 'MassSpectrumSetting')
323        # loaded_settings['Transient'] = self.get_scan_group_attr_data(scan_index, time_index, 'TransientSetting')
324
325    def get_output_parameters(self, polarity: int, scan_index: int = 0) -> dict:
326        """
327        Get the output parameters for the mass spectrum.
328
329        Parameters
330        ----------
331        polarity : int
332            The polarity of the mass spectrum +1 or -1.
333        scan_index : int, optional
334            The index of the scan. Default is 0.
335
336        Returns
337        -------
338        dict
339            A dictionary containing the output parameters.
340
341        """
342        from copy import deepcopy
343
344        output_parameters = default_parameters(self.file_location)
345
346        if self.isCentroid:
347            output_parameters["label"] = Labels.corems_centroid
348        else:
349            output_parameters["label"] = Labels.bruker_profile
350
351        output_parameters["analyzer"] = self.analyzer
352
353        output_parameters["instrument_label"] = self.instrument_label
354
355        output_parameters["sample_name"] = self.sample_name
356
357        output_parameters["Aterm"] = None
358
359        output_parameters["Bterm"] = None
360
361        output_parameters["Cterm"] = None
362
363        output_parameters["polarity"] = polarity
364
365        # scan_number and rt will be need to lc ms====
366
367        output_parameters["mobility_scan"] = 0
368
369        output_parameters["mobility_rt"] = 0
370
371        output_parameters["scan_number"] = scan_index
372
373        output_parameters["rt"] = 0
374
375        return output_parameters
376
377    def clean_data_frame(self, dataframe):
378        """
379        Clean the input dataframe by removing columns that are not expected.
380
381        Parameters
382        ----------
383        pandas.DataFrame
384            The input dataframe to be cleaned.
385
386        """
387
388        for column_name in dataframe.columns:
389            expected_column_name = self.parameters.header_translate.get(column_name)
390            if expected_column_name not in self._expected_columns:
391                del dataframe[column_name]
392
393    def check_columns(self, header_labels: list[str]):
394        """
395        Check if the given header labels match the expected columns.
396
397        Parameters
398        ----------
399        header_labels : list
400            The header labels to be checked.
401
402        Raises
403        ------
404        Exception
405            If any expected column is not found in the header labels.
406        """
407        found_label = set()
408
409        for label in header_labels:
410            if not label in self._expected_columns:
411                user_column_name = self.parameters.header_translate.get(label)
412                if user_column_name in self._expected_columns:
413                    found_label.add(user_column_name)
414            else:
415                found_label.add(label)
416
417        not_found = self._expected_columns - found_label
418
419        if len(not_found) > 0:
420            raise Exception(
421                "Please make sure to include the columns %s" % ", ".join(not_found)
422            )
423
424    def read_xml_peaks(self, data: str) -> DataFrame:
425        """
426        Read peaks from a Bruker .xml file and return a pandas DataFrame.
427
428        Parameters
429        ----------
430        data : str
431            The path to the .xml file.
432
433        Returns
434        -------
435        pandas.DataFrame
436            A DataFrame containing the peak data with columns: 'm/z', 'I', 'Resolving Power', 'Area', 'S/N', 'fwhm'.
437        """
438        from numpy import nan
439
440        with open(data, "r") as file:
441            content = file.readlines()
442            content = "".join(content)
443            bs_content = BeautifulSoup(content, features="xml")
444        peaks_xml = bs_content.find_all("pk")
445
446        # initialise lists of the peak variables
447        areas = []
448        fwhms = []
449        intensities = []
450        mzs = []
451        res = []
452        sn = []
453        # iterate through the peaks appending to each list
454        for peak in peaks_xml:
455            areas.append(
456                float(peak.get("a", nan))
457            )  # Use a default value if key 'a' is missing
458            fwhms.append(
459                float(peak.get("fwhm", nan))
460            )  # Use a default value if key 'fwhm' is missing
461            intensities.append(
462                float(peak.get("i", nan))
463            )  # Use a default value if key 'i' is missing
464            mzs.append(
465                float(peak.get("mz", nan))
466            )  # Use a default value if key 'mz' is missing
467            res.append(
468                float(peak.get("res", nan))
469            )  # Use a default value if key 'res' is missing
470            sn.append(
471                float(peak.get("sn", nan))
472            )  # Use a default value if key 'sn' is missing
473
474        # Compile pandas dataframe of these values
475        names = ["m/z", "I", "Resolving Power", "Area", "S/N", "fwhm"]
476        df = DataFrame(columns=names, dtype=float)
477        df["m/z"] = mzs
478        df["I"] = intensities
479        df["Resolving Power"] = res
480        df["Area"] = areas
481        df["S/N"] = sn
482        df["fwhm"] = fwhms
483        return df
484
485    def get_xml_polarity(self):
486        """
487        Get the polarity from an XML peaklist.
488
489        Returns
490        -------
491        int
492            The polarity of the XML peaklist. Returns -1 for negative polarity, +1 for positive polarity.
493
494        Raises
495        ------
496        Exception
497            If the data type is not XML peaklist in Bruker format or if the polarity is unhandled.
498        """
499
500        # Check its an actual xml
501        if not self.data_type or not self.delimiter:
502            self.set_data_type()
503
504        if isinstance(self.file_location, S3Path):
505            # data = self.file_location.open('rb').read()
506            data = BytesIO(self.file_location.open("rb").read())
507
508        else:
509            data = self.file_location
510
511        if self.data_type != "xml":
512            raise Exception("This function is only for XML peaklists (Bruker format)")
513
514        with open(data, "r") as file:
515            content = file.readlines()
516            content = "".join(content)
517            bs_content = BeautifulSoup(content, features="xml")
518        polarity = bs_content.find_all("ms_spectrum")[0]["polarity"]
519        if polarity == "-":
520            return -1
521        elif polarity == "+":
522            return +1
523        else:
524            raise Exception("Polarity %s unhandled" % polarity)