#!/usr/bin/env python
import numpy as np
from rdkit import Chem
from rdkit.Chem import Descriptors

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


def _func_worker(func_run, func_para_share, q_input, q_output):
    filename_in, select_peak_num = func_para_share
    fi = open(filename_in, "rb")
    while 1:
        try:
            q_item = q_input.get(block=True, timeout=1)
        except:
            continue

        if q_item is None:
            break

        i, para = q_item
        try:
            result = func_run(*para, fi, select_peak_num)
        except Exception as e:
            result = None
            print(e)
        q_output.put((i, result))
    return


def _compress_spec(mol_id, file_spec_start, file_spec_len, fi, select_peak_num):
    # Compress spectrum
    fi.seek(file_spec_start)
    spectrum = np.fromfile(fi, dtype=np.float32, count=file_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
    if not spectrum_compact.flags['C_CONTIGUOUS']:
        spectrum_compact = np.ascontiguousarray(spectrum_compact)
    return mol_id, spectrum_compact


def _extract_mol_info(mol_id, inchi, file_spec_start, file_spec_len):
    # Get molecule information
    mol = Chem.MolFromInchi(Molecule.simplify_inchi(inchi))
    if mol is None:
        print(inchi)
        raise RuntimeError("Error in parser mol_id: {}, inchi: {}".format(mol_id, inchi))
    inchikey = Chem.InchiToInchiKey(inchi)
    inchikey = inchikey.split('-')[0]
    mol_mass = Descriptors.ExactMolWt(mol)

    return mol_id, inchi, inchikey, mol_mass, file_spec_start, file_spec_len


def compress_spectra(db, filename_in, filename_out, select_peak_num=1000):
    # Get molecular information from raw database.
    mol_info_all = db.get_all_raw_spectra_info()
    mol_info_all = [x for x in mol_info_all if x[-1] > 0]

    # Generate molecular information.
    mol_info_all_new = run_multiple_process(_extract_mol_info, all_para_individual=mol_info_all)
    mol_info_all_new = [x for x in mol_info_all_new if x is not None]
    mol_info_all_new.sort(key=lambda x: x[2])
    mol_info_all_new.sort(key=lambda x: x[3])
    mol_info_all_new = [[i] + list(x[1:]) for i, x in enumerate(mol_info_all_new)]

    db_info_mol = [[x[0], x[1], x[2], x[3]] for x in mol_info_all_new]
    db.write_mol_info(db_info_mol)

    # Generate spectral information.
    fo = open(filename_out, 'wb')

    def _func_merge(result_all, result_i):
        if result_all is None:
            result_all = []

        mol_id, spectrum_compact = result_i
        file_spec_len = spectrum_compact.shape[0]
        file_spec_start = fo.tell()
        spectrum_compact.tofile(fo)

        result_all.append([mol_id, file_spec_start, file_spec_len])
        return result_all

    db_info_spec = run_multiple_process(
        _compress_spec,
        func_merge=_func_merge,
        func_worker=_func_worker,
        all_para_individual=[[x[0], x[4], x[5]] for x in mol_info_all_new],
        para_share=(filename_in, select_peak_num),
        threads=Parameter.threads
    )
    db.write_spec_info(db_info_spec)
    db.add_final_index()
