import numpy as np
import sqlite3


def simplify_inchi(raw_inchi):
    raw_inchi = raw_inchi.strip()
    inchi_list = raw_inchi.split('/')
    new_inchi_list = inchi_list[:2]
    for cont in inchi_list:
        if cont[0] == 'c':
            new_inchi_list.append(cont)
        elif cont[0] == 'h':
            new_inchi_list.append(cont)
    new_inchi = '/'.join(new_inchi_list)
    # new_inchi = new_inchi.encode('ascii', 'ignore')
    return new_inchi


def get_spectra(filename_db, filename_spectra, mol_inchi):
    db_conn = sqlite3.connect(filename_db)
    db_c = db_conn.cursor()
    fi = open(filename_spectra, 'rb')

    result = []
    for inchi_i in mol_inchi:
        inchi_i = simplify_inchi(inchi_i)
        spectrum = None
        if inchi_i is not None:
            sql = 'SELECT FileOffset,PeakNum FROM SpectraInfo WHERE InChI==?'
            sql = 'SELECT FileOffset,PeakNum FROM SpecInfo WHERE ID==(SELECT ID FROM MolInfo WHERE InChI==?)'
            print(inchi_i)
            db_c.execute(sql, (inchi_i,))
            cur_result = db_c.fetchone()
            if cur_result is not None:
                offset, peak_num = cur_result
                if peak_num > 0:
                    fi.seek(offset)
                    spectrum = np.fromfile(fi, dtype=np.float32, count=peak_num * 2).reshape(-1, 2)

        result.append(spectrum)

    return result


def get_spectra(filename_db, filename_spectra, mol_mass, ms1_ppm):
    db_conn = sqlite3.connect(filename_db)
    db_c = db_conn.cursor()
    fi = open(filename_spectra, 'rb')

    sql = 'SELECT SpecInfo.*,MolInfo.InChI,MolInfo.InChIKey,MolInfo.Name,MolInfo.Database FROM SpecInfo ' \
          'INNER JOIN MolInfo ON SpecInfo.ID=MolInfo.ID ' \
          'WHERE MolInfo.Mass > ? AND MolInfo.Mass < ?;'
    db_c.execute(sql, (mol_mass * (1 - ms1_ppm * 1e-6), mol_mass * (1 + ms1_ppm * 1e-6)))

    result = []
    for info_i in db_c:
        mol_id, offset, peak_num, inchi, inchi_key, name, database = info_i
        if peak_num > 0:
            fi.seek(offset)
            spectrum = np.fromfile(fi, dtype=np.float32, count=peak_num * 2).reshape(-1, 2)
            result.append((name, inchi, inchi_key, database, spectrum))

    return result
