from rdkit import Chem

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


def _convert_identifier_to_inchi_key(identifier):
    if identifier.startswith('InChI='):
        inchikey = Chem.InchiToInchiKey(identifier)
    else:
        mol = Chem.MolFromSmiles(identifier)
        if mol is None:
            inchikey = None
        else:
            inchi = Chem.MolToInchi(mol)
            inchikey = Chem.InchiToInchiKey(inchi)

    if inchikey is None:
        return 'NA'
    inchikey = inchikey.split('-')[0]
    return inchikey


def generate_mol_for_training():
    ion_type = '[M' + Parameter.addition + ']' + ('+' if Parameter.ion_charge == 0 else '-')
    db = DatabaseConnector(Parameter.database)

    all_id_inchi = db.get_all_spec_id_inchi_inchikey(ion_type)
    if Parameter.debug:
        training_id_inchi = all_id_inchi[:200]
    else:
        training_id_inchi = all_id_inchi

    if Parameter.exclude:
        with open(Parameter.exclude) as fi:
            exclude_identifier = fi.read().split('\n')
            exclude_identifier = set([x for x in exclude_identifier if x.strip()])
            exclude_inchikey = run_multiple_process(_convert_identifier_to_inchi_key,
                                                    all_para_individual=[(x,) for x in exclude_identifier])
            exclude_inchikey = set(exclude_inchikey)
            training_id_inchi = [(x[0], x[1]) for x in training_id_inchi if x[2] not in exclude_inchikey]
        return training_id_inchi

    if Parameter.exclude_id:
        all_id_inchi_splash = db.get_all_spec_id_inchi_splash(ion_type)

        exclude_splash = open(Parameter.exclude_id).read().split('\n')
        exclude_splash = set([x.strip() for x in exclude_splash if x.strip()])

        training_id_inchi = [(x[0], x[1]) for x in all_id_inchi_splash if x[2] not in exclude_splash]
        return training_id_inchi

    return training_id_inchi


def get_mol_from_mol_str(mol_identifier, mass=None):
    mol = Molecule()
    set_mol_success = mol.set_mol_by_identifier(mol_identifier)

    if not set_mol_success:
        # print mol_identifier
        return None

    mol_mass = mol.get_neutral_mass()
    if mass is not None:
        if abs(mol_mass - mass) > 0.5:
            return None

    return mol
