import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem

def smiles_to_fp(s, fp_dim=2048, pack=False):
    mol = Chem.MolFromSmiles(s)
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=fp_dim)
    onbits = list(fp.GetOnBits())
    # import pdb; pdb.set_trace()
    arr = np.zeros(fp.GetNumBits(), dtype=bool)
    # arr = np.zeros(fp.GetNumBits(), dtype=np.bool)
    arr[onbits] = 1

    if pack:
        arr = np.packbits(arr)

    return arr

def batch_smiles_to_fp(s_list, fp_dim):
    fps = []
    for s in s_list:
        fps.append(smiles_to_fp(s, fp_dim))
    fps = np.array(fps)

    assert fps.shape[0] == len(s_list) and fps.shape[1] == fp_dim

    return fps