import sqlite3

import numpy as np

from Core.Molecule import Molecule
from Utils.Multiplecore import run_multiple_process


def _test_molecules(mol_identifier):
    mol = Molecule()
    is_mol_right = mol.set_mol_by_identifier(mol_identifier)
    if not is_mol_right:
        return None
    inchi = mol.get_inchi()
    return inchi


def _func_worker_get_spectra(func_run, func_para_share, q_input, q_output):
    filename_db, filename_spectra = func_para_share

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

    while not q_input.empty():
        try:
            q_item = q_input.get(block=True, timeout=1)
        except:
            continue

        if q_item is None:
            break
        else:
            i, para = q_item

        cur_inchi = _test_molecules(para[0])
        spectrum = None
        if cur_inchi 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==?)'
            db_c.execute(sql, (cur_inchi,))
            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)
                    pass
                pass
            pass
        pass

        q_output.put((i, spectrum))
    return


class SpectraDBAdaptor(object):
    def __init__(self, filename_db, filename_spectra_library):
        self._filename_db = filename_db
        self._filename_spectra_library = filename_spectra_library
        pass

    def filter_molecules(self, filename_input):
        with open(filename_input) as fi:
            cont = fi.readlines()

        cont = [x.strip() for x in cont]
        molecule_test_result = run_multiple_process(_test_molecules, all_para_individual=[(x,) for x in cont])
        # molecule_test_result = run_function(_test_molecules, [(x,) for x in cont], 3, chunksize=1000)

        wrong_molecule = [x for x, y in zip(cont, molecule_test_result) if y is None]
        right_molecule = [x for x in molecule_test_result if x is not None]
        right_molecule = list(set(right_molecule))

        # Filter to leave molecules we don't know the spectra.
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()

        db_c.executemany('INSERT OR IGNORE INTO InfoCandidateMol (InChI,Batch) VALUES (?,-1) ',
                         [(x,) for x in right_molecule])
        db_c.execute('CREATE UNIQUE INDEX InfoCandidateMolInChI ON InfoCandidateMol(InChI)')
        db_c.execute('CREATE UNIQUE INDEX InfoCandidateSpecMolInChI ON InfoCandidateSpec(InChI)')
        db_c.execute(
            'SELECT InfoCandidateMol.ID,InfoCandidateSpec.ID FROM InfoCandidateMol '
            'JOIN InfoCandidateSpec ON InfoCandidateMol.InChI==InfoCandidateSpec.InChI')
        duplicate_id = db_c.fetchall()
        db_c.execute('DROP INDEX  InfoCandidateMolInChI')
        db_c.execute('DROP INDEX  InfoCandidateSpecMolInChI')
        db_c.executemany('DELETE FROM InfoCandidateMol WHERE InfoCandidateMol.ID==?',
                         [(x[0],) for x in duplicate_id])
        db_conn.commit()

        result = {
            'wrong': wrong_molecule,
            'in_db': [(x[1],) for x in duplicate_id],
        }
        return result

    def re_cluster_candidate_molecules(self, batch_size):
        cur_batch_num = 0
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()

        sql = "UPDATE InfoCandidateMol SET Batch=-1"
        db_c.execute(sql)

        while True:
            sql = 'SELECT ID FROM InfoCandidateMol WHERE Batch ==-1 LIMIT {}'.format(batch_size)
            db_c.execute(sql)
            not_assigned_mol = db_c.fetchall()
            if not not_assigned_mol:
                break

            sql = "UPDATE InfoCandidateMol SET Batch={} WHERE ID==?".format(cur_batch_num)
            db_c.executemany(sql, not_assigned_mol)
            cur_batch_num += 1

        db_conn.commit()

        return cur_batch_num

    def get_to_be_processed_molecules(self, batch_num):
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()
        sql = "SELECT ID,InChI FROM InfoCandidateMol WHERE Batch == {}".format(batch_num)
        db_c.execute(sql)
        result = db_c.fetchall()
        result = [(x[0], x[1]) for x in result]
        return result

    def storage_spectra(self, info_to_db, fi, fo):
        cur_fo_start = fo.tell()
        for item in info_to_db:
            item[1] += cur_fo_start

        fo.write(fi.read())

        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()
        sql_insert = 'INSERT OR IGNORE INTO InfoCandidateSpec (InChI, FileOffset, PeakNum) VALUES ' \
                     '((SELECT InChI FROM InfoCandidateMol WHERE InfoCandidateMol.ID==?),?,?);'
        db_c.executemany(sql_insert, info_to_db)

        info_to_del = [(x[0],) for x in info_to_db]
        sql_delete = 'DELETE FROM InfoCandidateMol WHERE InfoCandidateMol.ID==?;'
        db_c.executemany(sql_delete, info_to_del)
        db_conn.commit()
        pass

    def storage_spectra_compact(self, info_to_db, fi, fo, select_peak_num=1000):
        db_conn = sqlite3.connect(self._filename_db)

        # Get spectra information.
        result_spec = []
        for spec_info in info_to_db:
            mol_id, spec_start, spec_len = spec_info

            if spec_len > 0:
                file_offset, peak_num = self._compress_one_spectrum(fi, fo, spec_start, spec_len, select_peak_num)
                result_spec.append([mol_id, file_offset, peak_num])

        # Insert spectra info into database
        db_c = db_conn.cursor()
        sql_insert = 'INSERT OR IGNORE INTO SpecInfo (ID, FileOffset, PeakNum) VALUES (?,?,?);'
        db_c.executemany(sql_insert, result_spec)

        # Insert molecular info into database
        db_c = db_conn.cursor()
        sql_insert = 'INSERT OR IGNORE INTO MolInfo (ID, InChI, InChIKey) VALUES ' \
                     '(?,(SELECT InChI FROM SpectraInfo WHERE ID==?),"-");'
        db_c.executemany(sql_insert, [(x[0], x[0]) for x in result_spec])

        info_to_del = [(x[0],) for x in result_spec]
        sql_delete = 'DELETE FROM CandidateMol WHERE CandidateMol.ID==?;'
        db_c.executemany(sql_delete, info_to_del)
        db_conn.commit()
        pass

    def _compress_one_spectrum(self, fi, fo, spec_start, spec_len, select_peak_num):
        fi.seek(spec_start)
        spectrum = np.fromfile(fi, dtype=np.float32, count=spec_len * 2).reshape(-1, 2)
        spec_p = spectrum[:, 1]
        if len(spec_p) > select_peak_num:
            spectrum_compact = spectrum[spec_p >= spec_p[spec_p.argsort()[-select_peak_num]],]
        else:
            spectrum_compact = spectrum
        # Write to result files.
        file_offset = fo.tell()
        peak_num = spectrum_compact.shape[0]
        if not spectrum_compact.flags['C_CONTIGUOUS']:
            spectrum_compact = np.ascontiguousarray(spectrum_compact)
        spectrum_compact.tofile(fo)
        return file_offset, peak_num

    def get_all_raw_spectra_info(self):
        db_conn = sqlite3.connect(self._filename_db)
        sql = "SELECT * FROM InfoCandidateSpec;"
        db_c = db_conn.cursor()
        db_c.execute(sql)
        return db_c.fetchall()

    def write_mol_info(self, db_info_mol):
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()
        sql_insert = 'INSERT INTO MolInfo (ID, InChI, InChIKey,Name,Mass) VALUES (?,?,?,"-",?);'
        db_c.executemany(sql_insert, db_info_mol)
        db_conn.commit()

    def write_spec_info(self, db_info_spec):
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()
        sql_insert = 'INSERT INTO SpecInfo (ID, FileOffset, PeakNum) VALUES (?,?,?);'
        db_c.executemany(sql_insert, db_info_spec)
        db_conn.commit()

    def add_final_index(self):
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()
        db_c.execute('CREATE UNIQUE INDEX SpecInfo_ID_uindex ON SpecInfo (ID);')
        db_c.execute('CREATE UNIQUE INDEX MolInfo_ID_uindex ON MolInfo (ID);')
        db_c.execute('CREATE INDEX MolInfo_Mass_index ON MolInfo (Mass);')
        db_c.execute('CREATE UNIQUE INDEX MolInfo_InChI_uindex ON MolInfo (InChI);')
        db_c.execute('CREATE UNIQUE INDEX MolInfo_InChIKey_uindex ON MolInfo (InChIKey);')
        db_c.execute('DROP TABLE InfoCandidateMol;')
        db_c.execute('DROP TABLE InfoCandidateSpec;')
        db_conn.commit()
        db_c.execute('VACUUM;')

    def compress_spectra(self, fi, fo, select_peak_num=1000):
        db_conn = sqlite3.connect(self._filename_db)
        sql = "SELECT * FROM SpectraInfo;"
        db_c = db_conn.cursor()
        db_c.execute(sql)

        result_spec = []
        result_mol = []
        for spec_info in db_c:
            mol_id, inchi, spec_start, spec_len = spec_info
            if spec_len > 0:
                file_offset, peak_num = self._compress_one_spectrum(fi, fo, spec_start, spec_len, select_peak_num)
                result_spec.append([mol_id, file_offset, peak_num])
                result_mol.append([mol_id, inchi])

        # Insert spectra info into database
        db_c = db_conn.cursor()
        sql_insert = 'INSERT OR IGNORE INTO SpecInfo (ID, FileOffset, PeakNum) VALUES (?,?,?);'
        db_c.executemany(sql_insert, result_spec)

        # Insert molecular info into database
        db_c = db_conn.cursor()
        sql_insert = 'INSERT OR IGNORE INTO MolInfo (ID, InChI, InChIKey) VALUES ' \
                     '(?,?,"-");'
        db_c.executemany(sql_insert, result_mol)
        db_conn.commit()

        pass

    def get_max_cluster_num(self):
        db_conn = sqlite3.connect(self._filename_db)
        db_c = db_conn.cursor()
        sql = "SELECT MAX(Batch) FROM InfoCandidateMol"
        db_c.execute(sql)
        result = db_c.fetchall()[0][0]
        if result is None:
            return -1
        else:
            return result
        pass

    def get_spectra(self, mol_inchi):
        result = run_multiple_process(None, func_merge=None, func_worker=_func_worker_get_spectra,
                                      all_para_individual=[(x,) for x in mol_inchi],
                                      para_share=(
                                          self._filename_db, self._filename_spectra_library))

        return result

    pass
