import pickle

import numpy as np
import pandas as pd
import scipy.sparse

from Core.Molecule import Molecule
from Main.MainArguments import check_folder
from Main.Parameter import Parameter
from Model.FingerprintGenerator import generate_fingerprint
from Model.FormulaFeatureGenerator import FormulaFeatureGenerator
from Utils import FilenameGenerator
from Utils.Multiplecore import run_multiple_process


# from Model.FastForestClassifier import FastForestClassifier


def _extract_features_from_model(model):
    all_features = model.get_all_features()
    feature_info = {}
    for feature in all_features:
        item = feature[1]
        value = (feature[0], feature[2], all_features[feature])
        if item in feature_info:
            feature_info[item].append(value)
        else:
            feature_info[item] = [value]

    return feature_info


def _predict_mass(mol_id, mol_identifier, all_model_formula_feature):
    try:
        mol = Molecule()
        is_mol_right = mol.set_mol_by_identifier(mol_identifier.strip())
        if not is_mol_right:
            return mol_id, None, None, None, None

        mol_fingerprint = generate_fingerprint(mol)

        mol_ff = FormulaFeatureGenerator(mol)

        all_possible_peaks = mol_ff.find_all_possible_peaks(all_model_formula_feature)

        spectra_mass = []
        spectra_ff_num = []
        for peak_feature in all_possible_peaks:
            peak_mass = peak_feature[0].get_ionized_mass()
            wanted_formula_feature_num = peak_feature[1]

            spectra_mass.append(peak_mass)
            spectra_ff_num.append(wanted_formula_feature_num)
            pass

        spectra_mass = np.array(spectra_mass, np.float32)
        spectra_ff_num = np.array(spectra_ff_num, np.float32)

        peak_order = np.argsort(spectra_ff_num)
        spectra_mass = spectra_mass[peak_order]
        spectra_ff_num = spectra_ff_num[peak_order]

        spectra_mass = np.array(spectra_mass, np.float32, copy=False)
        spectra_ff_num = np.array(spectra_ff_num, np.int32, copy=False)
        used_ff_num = np.array(sorted(list(set(spectra_ff_num))), np.int32)
    except Exception as e:
        return mol_id, None, None, None, None
    return mol_id, mol_fingerprint, spectra_mass, spectra_ff_num, used_ff_num


def _generate_spectrum(mol_peak_mass, mol_peak_mass_ff_num, mol_peak_p, mol_peak_p_ff_num):
    if len(mol_peak_p_ff_num) == 0:
        return np.array([], np.float32)

    all_peak_num = mol_peak_mass.shape[0]
    result_spectrum = np.zeros((all_peak_num, 2), np.float32)
    result_spectrum[:, 0] = mol_peak_mass

    p_i = 0
    p_num = mol_peak_p_ff_num[p_i]
    for i, ff_num in enumerate(mol_peak_mass_ff_num):
        while ff_num > p_num:
            p_i += 1
            p_num = mol_peak_p_ff_num[p_i]
        result_spectrum[i, 1] = mol_peak_p[p_i]

    result_spectrum = result_spectrum[result_spectrum[:, 1] > 0]
    result_spectrum = result_spectrum[np.argsort(result_spectrum[:, 0])]

    result = pd.DataFrame(result_spectrum, columns=['mass', 'p'])
    result['mass_int'] = (result_spectrum[:, 0] * 1e4 + 0.5).astype(int)
    result = result.sort_values(by='p', ascending=False)
    result = result.drop_duplicates(subset='mass_int')
    result = result.sort_values(by='mass')
    result = result[result['mass'] >= 40.]

    result_spectrum = result[['mass', 'p']].values
    return result_spectrum


def _load_spectrum_data(spectra_path):
    fi_mass = open(FilenameGenerator.in_silico_spectra_peak_mass_data(spectra_path), 'rb')
    all_right_mol_num, all_peak_mass_num, all_mass_offset = \
        pickle.load(open(FilenameGenerator.in_silico_spectra_peak_mass_info(spectra_path)))

    fi_p = open(FilenameGenerator.in_silico_spectra_peak_p_data(spectra_path), 'rb')
    all_peak_p_num, all_p_offset = \
        pickle.load(open(FilenameGenerator.in_silico_spectra_peak_p_info(spectra_path)))

    for mol_num, mol_id in enumerate(all_right_mol_num):
        fi_mass.seek(all_mass_offset[mol_num])
        mol_peak_mass = np.fromfile(fi_mass, dtype=np.float32, count=all_peak_mass_num[mol_num])
        mol_peak_mass_ff_num = np.fromfile(fi_mass, dtype=np.int32, count=all_peak_mass_num[mol_num])

        fi_p.seek(all_p_offset[mol_num])
        mol_peak_p = np.fromfile(fi_p, dtype=np.float32, count=all_peak_p_num[mol_num])
        mol_peak_p_ff_num = np.fromfile(fi_p, dtype=np.int32, count=all_peak_p_num[mol_num])

        mol_spectrum = _generate_spectrum(
            mol_peak_mass, mol_peak_mass_ff_num, mol_peak_p, mol_peak_p_ff_num)
        yield mol_id, mol_spectrum


def _func_worker_generate_spectrum(func_run, func_para_share, q_input, q_output):
    spectra_path = func_para_share[0]
    fi_mass = open(FilenameGenerator.in_silico_spectra_peak_mass_data(spectra_path), 'rb')
    all_peak_mass_num, all_mass_offset = \
        pickle.load(open(FilenameGenerator.in_silico_spectra_peak_mass_info(spectra_path), 'rb'))

    fi_p = open(FilenameGenerator.in_silico_spectra_peak_p_data(spectra_path), 'rb')
    all_peak_p_num, all_p_offset = \
        pickle.load(open(FilenameGenerator.in_silico_spectra_peak_p_info(spectra_path), 'rb'))

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

        mol_num, mol_id = q_item
        fi_mass.seek(all_mass_offset[mol_num])
        mol_peak_mass = np.fromfile(fi_mass, dtype=np.float32, count=all_peak_mass_num[mol_num])
        mol_peak_mass_ff_num = np.fromfile(fi_mass, dtype=np.int32, count=all_peak_mass_num[mol_num])

        fi_p.seek(all_p_offset[mol_num])
        mol_peak_p = np.fromfile(fi_p, dtype=np.float32, count=all_peak_p_num[mol_num])
        mol_peak_p_ff_num = np.fromfile(fi_p, dtype=np.int32, count=all_peak_p_num[mol_num])

        mol_spectrum = _generate_spectrum(
            mol_peak_mass, mol_peak_mass_ff_num, mol_peak_p, mol_peak_p_ff_num)
        q_output.put((mol_num, (mol_id, mol_spectrum)))
    return


def generate_spectra_path(batch_num=None):
    if batch_num is None:
        batch_num = Parameter.batch_num
    out_spectra_path = Parameter.spectra.strip() + '/' + str(batch_num) + '/'
    check_folder(out_spectra_path)
    return out_spectra_path


def predict_mass(model, molecules_to_be_processed, output_path):
    model_feature = _extract_features_from_model(model)

    all_mol_num = len(molecules_to_be_processed)

    all_fingerprints = np.zeros((all_mol_num, sum([x[1] for x in Parameter.fp_type_bits])), np.int8)
    all_offset = np.zeros(all_mol_num, np.uint64)
    all_peak_num = np.zeros(all_mol_num, np.uint32)
    all_right_mol_id = []
    all_wrong_mol_id = []
    fo = open(FilenameGenerator.in_silico_spectra_peak_mass_data(output_path), 'wb')
    all_wanted_ff = [[] for _ in range(model.get_total_feature_num())]

    def func_merge(right_item_num, f_cur_result):
        if right_item_num is None:
            right_item_num = 0
        cur_mol_id, mol_fingerprint, spectra_mass, spectra_ff_num, used_ff_num = f_cur_result
        if mol_fingerprint is None:
            all_wrong_mol_id.append(cur_mol_id)
        else:
            all_fingerprints[right_item_num] = mol_fingerprint
            all_offset[right_item_num] = fo.tell()
            all_peak_num[right_item_num] = spectra_mass.shape[0]
            spectra_mass.tofile(fo)
            spectra_ff_num.tofile(fo)
            all_right_mol_id.append(cur_mol_id)
            for ff_num in used_ff_num:
                all_wanted_ff[ff_num].append(right_item_num)
            right_item_num += 1
            pass
        return right_item_num

    right_item_num = run_multiple_process(
        func_run=_predict_mass, func_merge=func_merge,
        all_para_individual=[(x[0], x[1],) for x in molecules_to_be_processed],
        para_share=(model_feature,))

    all_peak_num = all_peak_num[:right_item_num]
    all_offset = all_offset[:right_item_num]

    pickle.dump([all_peak_num, all_offset],
                open(FilenameGenerator.in_silico_spectra_peak_mass_info(output_path), 'wb'))

    error_mol = [str(x[0]) + '\t' + str(x[1]) for x in molecules_to_be_processed if x[0] in all_wrong_mol_id]
    # error_mol = [molecules_to_be_processed[x] for x in all_wrong_mol_id]
    # error_mol = [str(x[0]) + '\t' + str(x[1]) for x in error_mol]
    with open(FilenameGenerator.in_silico_spectra_error_mol(output_path), 'wt') as fo:
        fo.writelines('\n'.join(error_mol))

    all_fingerprints = all_fingerprints[:right_item_num]
    for ff_num in range(len(all_wanted_ff)):
        all_wanted_ff[ff_num] = np.array(all_wanted_ff[ff_num], dtype=np.int32)
    return all_right_mol_id, all_fingerprints, all_wanted_ff


def predict_p(model, all_fingerprints, all_wanted_ff, output_path):
    all_mol_num = all_fingerprints.shape[0]
    all_ff_num = len(all_wanted_ff)
    possibility = model.predict_probability_with_peak_model(all_fingerprints, all_wanted_ff)

    wanted_ff = np.concatenate([[i] * len(x) for i, x in enumerate(all_wanted_ff)])
    wanted_mol_list = np.concatenate(all_wanted_ff)
    possibility = np.concatenate(possibility)

    try:
        possibility_result = scipy.sparse.csr_matrix(
            (possibility, (wanted_mol_list, wanted_ff)),
            shape=(all_mol_num, all_ff_num), dtype=np.float32, copy=False
        )
    except Exception as e:
        print((possibility.shape, (wanted_mol_list.shape, wanted_ff.shape)), (all_mol_num, all_ff_num))
        raise e

    possibility_result.sort_indices()
    possibility_result_data = np.array(possibility_result.data, np.float32, copy=False)
    possibility_result_indices = np.array(possibility_result.indices, np.int32, copy=False)

    # storage the result.
    all_offset = np.zeros(all_mol_num, np.uint64)
    all_peak_num = np.zeros(all_mol_num, np.uint32)
    fo = open(FilenameGenerator.in_silico_spectra_peak_p_data(output_path), 'wb')
    for cur_mol_num in range(all_mol_num):
        cur_peak_start = possibility_result.indptr[cur_mol_num]
        cur_peak_end = possibility_result.indptr[cur_mol_num + 1]

        all_peak_num[cur_mol_num] = cur_peak_end - cur_peak_start
        all_offset[cur_mol_num] = fo.tell()
        possibility_result_data[cur_peak_start:cur_peak_end].tofile(fo)
        possibility_result_indices[cur_peak_start:cur_peak_end].tofile(fo)
    fo.close()

    pickle.dump([all_peak_num, all_offset],
                open(FilenameGenerator.in_silico_spectra_peak_p_info(output_path), 'wb'))
    pass


def generate_spectra(all_right_mol_id, spectra_path, select_peak_num=1000):
    file_output = open(FilenameGenerator.in_silico_spectra_data(spectra_path), 'wb')

    def _func_merge_spectrum(f_final_result, f_cur_result):
        if f_final_result is None:
            info_to_db = []
        else:
            info_to_db = f_final_result
        mol_id, mol_spectrum = f_cur_result

        file_offset = file_output.tell()

        if len(mol_spectrum) > select_peak_num:
            spec_p = mol_spectrum[:, 1]
            mol_spectrum = mol_spectrum[spec_p >= spec_p[spec_p.argsort()[-select_peak_num]],]

        if not mol_spectrum.flags['C_CONTIGUOUS']:
            mol_spectrum = np.ascontiguousarray(mol_spectrum)

        peak_num = mol_spectrum.shape[0]
        mol_spectrum.tofile(file_output)
        info_to_db.append([mol_id, file_offset, peak_num])
        return info_to_db

    info_to_db = run_multiple_process(None, func_merge=_func_merge_spectrum,
                                      func_worker=_func_worker_generate_spectrum,
                                      all_para_individual=all_right_mol_id, para_share=(spectra_path,))
    pickle.dump(info_to_db, open(FilenameGenerator.in_silico_spectra_info(spectra_path), 'wb'))
    return


def get_spectra(spectra_path):
    info_to_db = pickle.load(open(FilenameGenerator.in_silico_spectra_info(spectra_path), 'rb'))
    fi = open(FilenameGenerator.in_silico_spectra_data(spectra_path), 'rb')
    return info_to_db, fi
