# -*- 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, 12, 11, 26), 'pickle.load', 'pickle.load', ({(11, 24, 11, 25): 'f'}, {}), '(f)', False, 'import pickle\n'), ((51, 15, 51, 44), 'numpy.concatenate', 'np.concatenate', ({(51, 30, 51, 43): '(pfp, rxn_fp)'}, {}), '((pfp, rxn_fp))', True, 'import numpy as np\n'), ((23, 14, 23, 38), 'rdkit.Chem.MolFromSmiles', 'Chem.MolFromSmiles', ({(23, 33, 23, 37): 'rsmi'}, {}), '(rsmi)', False, 'from rdkit import Chem\n'), ((28, 17, 29, 99), 'rdkit.Chem.AllChem.GetMorganFingerprintAsBitVect', 'AllChem.GetMorganFingerprintAsBitVect', (), '', False, 'from rdkit.Chem import AllChem, DataStructs\n'), ((30, 13, 30, 49), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((31, 8, 31, 51), 'rdkit.Chem.DataStructs.ConvertToNumpyArray', 'DataStructs.ConvertToNumpyArray', ({(31, 40, 31, 46): 'fp_bit', (31, 48, 31, 50): 'fp'}, {}), '(fp_bit, fp)', False, 'from rdkit.Chem import AllChem, DataStructs\n'), ((38, 14, 38, 38), 'rdkit.Chem.MolFromSmiles', 'Chem.MolFromSmiles', ({(38, 33, 38, 37): 'psmi'}, {}), '(psmi)', False, 'from rdkit import Chem\n'), ((42, 17, 43, 97), 'rdkit.Chem.AllChem.GetMorganFingerprintAsBitVect', 'AllChem.GetMorganFingerprintAsBitVect', (), '', False, 'from rdkit.Chem import AllChem, DataStructs\n'), ((44, 13, 44, 47), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((45, 8, 45, 51), 'rdkit.Chem.DataStructs.ConvertToNumpyArray', 'DataStructs.ConvertToNumpyArray', ({(45, 40, 45, 46): 'fp_bit', (45, 48, 45, 50): 'fp'}, {}), '(fp_bit, fp)', False, 'from rdkit.Chem import AllChem, DataStructs\n')]