#!/usr/bin/env python
import sys
import numpy as np
import pandas as pd
import timeit

from Library import Score
from Library import SpectraDBAdaptor
from Library.Arguments import Arguments
from Library.Input import File
from Library.Parameter import Parameter

example_parameters = """
-spectra /home/yli/project/identification/benchmark_data/casmi_2016/spectrum_neg/test-001.mgf
-db_path /d/test/library/
-out /d/test/score-001_2.txt
-ion 0
-threads 24
"""


def main():
    print('Reading spectra.')
    candidate_spectra_all = File.get_spectrum_info(Parameter.spectrum)
    # candidate_spectra_all.sort(key=lambda x: x.precursor_mz)
    result = []
    print('Got {} spectra.'.format(len(candidate_spectra_all)))

    time_start = timeit.default_timer()
    print('Scoring spectra.')
    for candidate_spectra_i in candidate_spectra_all:
        cur_result = scoring_spectrum(candidate_spectra_i)
        result += cur_result

    time_end = timeit.default_timer()
    print('Finished, used time: {} s.'.format(time_end - time_start))
    print('Output result.')
    df = pd.DataFrame(result)
    # df.to_csv(Parameter.output_file, index=False, header=True)
    df.to_csv(Parameter.output_file, sep=",", index=False, header=True)
    return


def scoring_spectrum(candidate_spectra_i):
    spec_neu_mass = candidate_spectra_i.precursor_mz - \
                    {'+H': 1.007276, '-H': -1.007276}[Parameter.addition]
    db_info = SpectraDBAdaptor.get_spectra(
        Parameter.database,
        Parameter.model,
        spec_neu_mass,
        Parameter.ms1_ppm)
    spec_data = candidate_spectra_i.data.T
    spec_data = spec_data[np.bitwise_and(
        spec_data[:, 1] > 0,
        spec_data[:, 0] <= candidate_spectra_i.precursor_mz - 0.5)]
    spec_data = spec_data[spec_data[:, 0].argsort()]
    cur_result = []
    for info_i in db_info:
        mol_name, mol_inchi, mol_inchi_key, mol_database, spec_predict_i = info_i
        spec_predict_i = spec_predict_i[spec_predict_i[:, 0].argsort()]
        score = Score.score_spectrum(
            spec_data,
            spec_predict_i,
            Parameter.ms2_ppm)

        cur_result.append({
            'scan': candidate_spectra_i.raw_scan_num,
            'name': mol_name,
            'inchi': mol_inchi,
            'inchi_key': mol_inchi_key,
            'database': mol_database,
            'score': score,
            'retention_time': candidate_spectra_i.retention_time,
            'm/z': candidate_spectra_i.precursor_mz
        })
    if len(cur_result) > 0 and not Parameter.all_result:
        max_score = max([s['score'] for s in cur_result])
        if max_score > 0:
            cur_result = [x for x in cur_result if x['score'] == max_score]
        else:
            cur_result = []

    return cur_result


if __name__ == '__main__':
    parser = Arguments()
    args = example_parameters.split()
    if len(sys.argv) == 1:
        parser.parse_args(args=args)
    else:
        parser.parse_args()

    main()
