import zlib

import base64
import numpy as np
import xml.etree.ElementTree as ET

from .Core import Experiment


def _decode_spec(line, peak_info=None):
    type_code = '>f4'
    shape = (-1, 2)
    need_de_compress = False

    if peak_info is not None:
        assert peak_info['byteOrder'] == 'network'
        assert peak_info['contentType'] == 'm/z-int'
        if peak_info['precision'] == '64':
            type_code = '>f8'
        if peak_info['compressionType'] == 'zlib':
            need_de_compress = True

    if line is None or len(line) == 0:
        return np.array([], dtype=np.dtype(type_code))
    else:
        decoded = base64.b64decode(line.strip())
        if need_de_compress:
            decoded = zlib.decompress(decoded)

        spec = np.frombuffer(decoded, dtype=np.dtype(type_code)).reshape(shape)
        spec = spec[spec[:, 1] > 0.]
        spec = spec[np.argsort(spec[:, 0])]
        return spec.T


def read(filename_input, wanted_ms_level=None):
    experiment = Experiment.Experiment()

    parser = ET.iterparse(filename_input, ('start', 'end'))
    for event, elem in parser:
        if event == 'start' and elem.tag.endswith('}scan'):
            mslevel = elem.attrib.get('msLevel', '0')
            if mslevel == '1' and \
                    ((wanted_ms_level is None) or (wanted_ms_level is not None and 1 in wanted_ms_level)):
                scan_info = elem.attrib
                while event == 'start':
                    event, elem = next(parser)

                assert elem.tag.endswith('}peaks')
                cur_spec = _decode_spec(elem.text, peak_info=elem.attrib)
                experiment.add_ms1(
                    spectrum_data=cur_spec,
                    scan_num=int(scan_info['num']),
                    retention_time=float(scan_info['retentionTime'][2:-1])
                )
                event, elem = next(parser)
                pass

            elif mslevel == '2' and \
                    ((wanted_ms_level is None) or (wanted_ms_level is not None and 2 in wanted_ms_level)):
                scan_info = elem.attrib

                event, elem = next(parser)
                event, elem = next(parser)
                precursor_info = elem.attrib
                precursor_mz = elem.text

                if 'precursorCharge' in precursor_info:
                    precursor_charge = int(precursor_info['precursorCharge'])
                else:
                    precursor_charge = 1

                if scan_info['polarity'] == '-':
                    precursor_charge = -precursor_charge

                event, elem = next(parser)
                event, elem = next(parser)

                cur_spec = _decode_spec(elem.text, peak_info=elem.attrib)

                experiment.add_ms2(spectrum_data=cur_spec,
                                   precursor_mz=float(precursor_mz),
                                   precursor_charge=precursor_charge,
                                   scan_num=int(scan_info['num']),
                                   retention_time=float(scan_info['retentionTime'][2:-1]))
                event, elem = next(parser)
        pass
    return experiment

    pass
