# -*- coding: utf-8 -*-

import pickle
import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem,DataStructs


def get_classes(path):
    f = open(path, 'rb')
    dict_ = pickle.load(f)
    f.close()
    classes = sorted(dict_.items(), key=lambda d: d[1],reverse=True)
    classes = [(x,y) for x,y in classes]
    return classes

def create_rxn_Morgan2FP_concatenate(rsmi, psmi, rxnfpsize=16384, pfpsize=16384, useFeatures=False, calculate_rfp=True, useChirality=True):
    # Similar as the above function but takes smiles separately and returns pfp and rfp separately

    rsmi = rsmi.encode('utf-8')
    psmi = psmi.encode('utf-8')
    try:
        mol = Chem.MolFromSmiles(rsmi)
    except Exception as e:
        print(e)
        return
    try:
        fp_bit = AllChem.GetMorganFingerprintAsBitVect(
            mol=mol, radius=2, nBits=rxnfpsize, useFeatures=useFeatures, useChirality=useChirality)
        fp = np.empty(rxnfpsize, dtype='float32')
        DataStructs.ConvertToNumpyArray(fp_bit, fp)
    except Exception as e:
        print("Cannot build reactant fp due to {}".format(e))
        return
    rfp = fp

    try:
        mol = Chem.MolFromSmiles(psmi)
    except Exception as e:
        return
    try:
        fp_bit = AllChem.GetMorganFingerprintAsBitVect(
            mol=mol, radius=2, nBits=pfpsize, useFeatures=useFeatures, useChirality=useChirality)
        fp = np.empty(pfpsize, dtype='float32')
        DataStructs.ConvertToNumpyArray(fp_bit, fp)
    except Exception as e:
        print("Cannot build product fp due to {}".format(e))
        return
    pfp = fp
    rxn_fp = pfp - rfp
    final_fp = np.concatenate((pfp, rxn_fp))
    return final_fp
[(11, 'pickle.load', 'pickle.load', False, 'import pickle\n'), (51, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (23, 'rdkit.Chem.MolFromSmiles', 'Chem.MolFromSmiles', False, 'from rdkit import Chem\n'), (28, 'rdkit.Chem.AllChem.GetMorganFingerprintAsBitVect', 'AllChem.GetMorganFingerprintAsBitVect', False, 'from rdkit.Chem import AllChem, DataStructs\n'), (30, 'numpy.empty', 'np.empty', True, 'import numpy as np\n'), (31, 'rdkit.Chem.DataStructs.ConvertToNumpyArray', 'DataStructs.ConvertToNumpyArray', False, 'from rdkit.Chem import AllChem, DataStructs\n'), (38, 'rdkit.Chem.MolFromSmiles', 'Chem.MolFromSmiles', False, 'from rdkit import Chem\n'), (42, 'rdkit.Chem.AllChem.GetMorganFingerprintAsBitVect', 'AllChem.GetMorganFingerprintAsBitVect', False, 'from rdkit.Chem import AllChem, DataStructs\n'), (44, 'numpy.empty', 'np.empty', True, 'import numpy as np\n'), (45, 'rdkit.Chem.DataStructs.ConvertToNumpyArray', 'DataStructs.ConvertToNumpyArray', False, 'from rdkit.Chem import AllChem, DataStructs\n')]