import copy

from Core.Molecule import Molecule
from Core.Spectrum import Spectrum
from FormulaClass.FormulaSelector import FormulaSelector
from FormulaClass.SpectrumAnnotator import SpectrumAnnotator
from Main.Parameter import Parameter
from Model.FingerprintGenerator import generate_fingerprint
from Model.FormulaFeatureGenerator import *
from Utils.DatabaseConnector import DatabaseConnector


def annotate_spectrum_from_db(spec_id, mol_inchi, use_known_model=None):
    # Generate molecules
    mol = Molecule()
    mol_result = mol.set_mol_by_identifier(mol_inchi)

    if not mol_result:
        print('Molecule error: {}'.format(mol_inchi))
        return None

    # Generate spectrum
    db = DatabaseConnector(Parameter.database)
    ori_spec = db.get_spec_from_spec_id(spec_id)

    if ori_spec is not None:
        spectrum = Spectrum(ori_spec, need_refine=True)
    else:
        return None

    # Annotate spectrum
    spec_info = annotate_spectrum(mol, spectrum, use_known_model)

    if spec_info is None:
        return None

    mol_id = db.get_mol_id_from_spec_id(spec_id)
    spec_info['mol_id'] = mol_id
    spec_info['spec_id'] = spec_id

    return spec_info


def annotate_spectrum(mol, spectrum, known_model=False, allow_null_spec=False):
    spectrum = copy.deepcopy(spectrum)

    # Calculate fingerprint
    fingerprint = generate_fingerprint(mol)

    # Annotate spectra.
    spectrum_annotator = SpectrumAnnotator(mol)
    annotated_peak_num = spectrum_annotator.annotate_spectrum(spectrum, match_large_error=False)
    if (annotated_peak_num == 0) and (not allow_null_spec):
        return None

    if known_model:
        formula_selector = FormulaSelector(fingerprint, known_model)
        formula_selector.refine_annotation(spectrum)
        pass

    else:
        # Annotate molecule.
        feature_generator = FormulaFeatureGenerator(mol)
        #feature_generator = SimpleFormulaFeatureGenerator(mol)
        feature_generator.convert_to_formula_feature(spectrum)

    annotated_spectrum = {'spectrum': spectrum,
                          'fingerprint': fingerprint}
    return annotated_spectrum


def func_merge_annotated_spectrum(f_final_result, f_cur_result):
    if f_cur_result is None:
        return f_final_result

    if f_final_result is None:
        known_formula_list = []
        known_formula_dict = {}

        known_ff_list = []
        known_ff_num = []
        known_ff_dict = {}

        known_spec_info = [f_cur_result]
    else:
        known_formula_list, known_formula_dict, \
        known_ff_list, known_ff_num, known_ff_dict, \
        known_spec_info = f_final_result

        known_spec_info.append(f_cur_result)

    cur_spec_known_ff = set()
    for peak in f_cur_result['spectrum']:
        for cur_ff_num, cur_ff in enumerate(peak['ff']):
            if cur_ff in known_ff_dict:
                cur_known_ff_num = known_ff_dict[cur_ff]
                peak['ff'][cur_ff_num] = known_ff_list[cur_known_ff_num]
                cur_spec_known_ff.add(known_ff_list[cur_known_ff_num])

            else:
                if cur_ff[-1] in known_formula_dict:
                    cur_ff = cur_ff[:-1] + (known_formula_list[known_formula_dict[cur_ff[-1]]],)
                else:
                    known_formula_dict[cur_ff[-1]] = len(known_formula_list)
                    known_formula_list.append(cur_ff[-1])

                peak['ff'][cur_ff_num] = cur_ff
                known_ff_dict[cur_ff] = len(known_ff_list)
                known_ff_list.append(cur_ff)
                cur_spec_known_ff.add(cur_ff)
                known_ff_num.append(0)

    for cur_known_ff in cur_spec_known_ff:
        known_ff_num[known_ff_dict[cur_known_ff]] += 1

    if Parameter.debug:
        print(len(known_spec_info))

    return known_formula_list, known_formula_dict, \
           known_ff_list, known_ff_num, known_ff_dict, \
           known_spec_info
