#!/usr/bin/env python
import random

import numpy as np

from Main.MainArguments import MainArguments
from Main.Parameter import Parameter
from Model.AnnotateMolecule import annotate_spectrum_from_db, func_merge_annotated_spectrum
from Model.FormulaFeatureFinder import FormulaFeatureFinder
from Model.ModelBuilder import ModelBuilder
from Utils import FilenameGenerator
from Utils.GenerateMolecule import generate_mol_for_training
from Utils.Multiplecore import run_multiple_process

"""
-db /e/Work/Project/identification/sf-matching/BuildModel_CalculateLibrary/Data/Spectral.db
-out_path /d/test/tmp
-threads 8
-extract_cutoff 5
-ion 0                                                 
-d
"""


def main():
    deal_with_arguments()

    training_id_inchi = generate_mol_for_training()
    para = [(x[0], x[1], False) for x in training_id_inchi]

    select_mol_num = int(len(para) * Parameter.select)
    select_mol = random.sample(range(len(para)), select_mol_num)
    select_mol.sort()
    select_para = [para[x] for x in select_mol]
    para = select_para

    all_annotated_results = run_multiple_process(func_run=annotate_spectrum_from_db,
                                                 func_merge=func_merge_annotated_spectrum,
                                                 all_para_individual=para,
                                                 threads=None)

    _, _, known_ff_list, known_ff_num, _, all_spec_info = all_annotated_results
    print('Total find {} spectra.'.format(len(all_spec_info)))

    model_builder = ModelBuilder()
    # Get wanted feature.
    formula_feature_finder = FormulaFeatureFinder()
    formula_feature_finder.select_feature_by_number(Parameter.extract_cutoff, known_ff_list, known_ff_num)

    formula_feature_finder.extract_features(all_spec_info)
    model_builder.add_intensity(formula_feature_finder)

    model_builder.mol_id = np.array([x['mol_id'] for x in all_spec_info], np.int32)
    model_builder.spec_id = np.array([x['spec_id'] for x in all_spec_info], np.int32)
    model_builder.uniq_mol_id = np.array(list({x['mol_id'] for x in all_spec_info}), np.int32)
    model_builder.add_fingerprint(np.array([x['fingerprint'] for x in all_spec_info], np.int8))

    model_builder.save_to_file(FilenameGenerator.feature_data(Parameter.output_path))
    pass


def deal_with_arguments():
    parser = MainArguments(description='')

    parser.add_argument('-exclude', required=False, type=str, dest='exclude', default='', metavar='', help='')
    parser.add_argument('-exclude_id', required=False, type=str, dest='exclude_id', default='', metavar='', help='')

    parser.add_argument('-extract_cutoff', required=False, type=int, dest='extract_cutoff', default=5,
                        metavar='', help='')

    parser.add_argument('-select', required=False, type=float, dest='select', default=1., metavar='', help='')

    parser.parse_args()


if __name__ == '__main__':
    main()
