import logging

import numpy as np


def _parse_info(info):
    result = {}

    item = info.strip().split('=')
    if len(item) < 2:
        return result

    info = item[0].strip()
    value = item[1].strip() if len(item) == 2 else '='.join(item[1:]).strip()

    if info == 'PEPMASS':
        result = dict(precursor_mz=float(value))

    elif info == 'CHARGE':
        _charge_dict = {'-': -1, '+': 1}
        if value[-1] in _charge_dict:
            charge = _charge_dict[value[-1]] * int(value[:-1])
        elif value[0] in _charge_dict:
            charge = _charge_dict[value[0]] * int(value[1:])
        else:
            charge = int(value)
        result = dict(charge=charge)
    else:
        result = {info.lower(): value}

    return result


def read(filename_input: str):
    experiment = Experiment()

    with open(filename_input) as fi:
        is_spectra_started = False
        for line in fi:
            if line[0] == '#':
                continue

            if line.find('BEGIN ION') >= 0:
                is_spectra_started = True
                spectrum = dict(
                    charge=1,
                    precursor_mz=0.,
                    spectrum=[],
                )
                info = {}
                continue

            if line.find('END IONS') >= 0:
                experiment.add_ms2(spectrum_data=np.array(spectrum['spectrum'], np.float32),
                                   precursor_mz=spectrum['precursor_mz'],
                                   precursor_charge=spectrum['charge'],
                                   scan_num=spectrum['scans'],
                                   )
                continue

            if not is_spectra_started:
                # info.update(_parse_info(line))
                continue

            if line.find('=') >= 0:
                spectrum.update(_parse_info(line))
            else:
                items = line.split()
                assert len(items) == 2, "Error found when parse {}".format(line)
                spectrum['spectrum'].append((float(items[0]), float(items[1])))

        return experiment

    pass


def write(exp, filename_output: str):
    with open(filename_output, 'wt') as fo:
        for spec in exp:
            _write_spec(spec, fo)
    pass


def write_spec(spec_info: dict, spec_data: np.ndarray, filename_output: str):
    """
    Write spectrum to a mgf file.
    :param spec_info:
    :param spec_data:
    :param filename_output:
    :return: None
    """
    if filename_output[-3:] != "mgf":
        logging.warning("Output filename is not a mgf file!")

    with open(filename_output, 'wt') as fo:
        _write_spec(spec_info, spec_data, fo)


def _write_spec(spec_info, spec_data, fileout):
    out_str = ['BEGIN IONS']

    def __add_to_output_str_if_exist(str_pre, str_suffix, item_name):
        if item_name in spec_info:
            out_str.append(str_pre + str(spec_info[item_name]) + str_suffix)

    __add_to_output_str_if_exist('PEPMASS=', '', 'precursor_mz')
    __add_to_output_str_if_exist('SCANS=', '', 'raw_scan_num')
    __add_to_output_str_if_exist('MSLEVEL=', '', 'ms_level')
    __add_to_output_str_if_exist('RTINSECONDS=', '', 'retention_time')

    if 'precursor_charge' in spec_info:
        out_str.append('CHARGE=' + str(abs(spec_info['precursor_charge'])) + '+'
                       if spec_info['precursor_charge'] >= 0 else '-')

    for peak in spec_data:
        out_str.append(str(peak[0]) + '\t' + str(peak[1]))

    out_str.append('END IONS\n')
    fileout.write('\n'.join(out_str))
