# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/core/03_ligand.ipynb.

# %% auto 0
__all__ = ['plot_drug', 'rdkit_conformer', 'tanimoto', 'get_rdkit', 'get_rdkit_3d', 'get_rdkit_all', 'remove_hi_corr',
           'preprocess', 'get_rdkit_df', 'get_fp', 'compress_fp', 'tanimoto_numba', 'hash_fp', 'get_same_mol_group']

# %% ../../nbs/core/03_ligand.ipynb 3
import pandas as pd, numpy as np
from rdkit import Chem,RDLogger,DataStructs
from rdkit.Chem import Draw, rdDepictor, Descriptors,Descriptors3D, AllChem
from skfp.fingerprints import AtomPairFingerprint, ECFPFingerprint, MACCSFingerprint
import hashlib,numba
from .data import *

from tqdm.contrib.concurrent import process_map
from tqdm import tqdm
RDLogger.DisableLog('rdApp.warning')

# %% ../../nbs/core/03_ligand.ipynb 5
def plot_drug(drug_dict, flip_list=None, save_path=None):
    flip_list = flip_list or []
    mols = []

    for name, smi in drug_dict.items():
        mol = Chem.MolFromSmiles(smi)
        rdDepictor.Compute2DCoords(mol)

        # Flip horizontally if requested
        if name in flip_list:
            conf = mol.GetConformer()
            for i in range(conf.GetNumAtoms()):
                x, y, z = conf.GetAtomPosition(i)
                conf.SetAtomPosition(i, (-x, y, z))

        mol.SetProp("_Name", name)
        mols.append(mol)

    # Draw to SVG
    svg_obj = Draw.MolsToGridImage(
        mols,
        molsPerRow=3,
        subImgSize=(300, 250),
        legends=[m.GetProp("_Name") for m in mols],
        useSVG=True
    )

    svg_str = getattr(svg_obj, "data", getattr(svg_obj, "GetDrawingText", lambda: svg_obj)())

    if save_path:
        with open(save_path, "w", encoding="utf-8") as f:
            f.write(svg_str)
        print(f"✅ Saved structure grid to: {save_path}")
    else:
        # from IPython.display import SVG, display
        # display(SVG(svg_str))
        return svg_obj

# %% ../../nbs/core/03_ligand.ipynb 8
def rdkit_conformer(SMILES, # SMILES string
                    output=None, # file ".sdf" to be saved
                    method='ETKDG', # Optimization method, can be 'UFF', 'MMFF' or 'ETKDGv3'
                    visualize=True, #whether or not to visualize the compound
                    seed = 3,# randomness of the 3D conformation
                    ):

    "Gemerate 3D conformers from SMILES"
    
    np.random.seed(seed) 
    mol = Chem.MolFromSmiles(SMILES)
    
    # Generate a 3D conformation of the molecule
    AllChem.EmbedMolecule(mol)
    

    # Optimize the 3D conformation using the specified force field method
    if method == 'UFF':
        AllChem.UFFOptimizeMolecule(mol)
    elif method == 'MMFF':
        AllChem.MMFFOptimizeMolecule(mol)
    elif method == 'ETKDG':
        AllChem.EmbedMultipleConfs(mol, numConfs=1, useExpTorsionAnglePrefs=True, 
                                   useBasicKnowledge=True, enforceChirality=True,randomSeed=seed)
        AllChem.ETKDGv3()
        AllChem.UFFOptimizeMolecule(mol)

    else:
        raise ValueError('Invalid method specified')
        

    # Remove hydrogens from the molecule
    # mol = Chem.RemoveHs(mol)

    if output is not None:
        Path(output).parent.mkdir(parents=True,exist_ok=True)
    
        w = Chem.SDWriter(output)
        w.write(mol)
        w.close()
    return mol

# %% ../../nbs/core/03_ligand.ipynb 10
def tanimoto(df, # df with SMILES and ID columns
             smiles_col='SMILES', # colname of SMILES
             id_col='ID', # colname of compound ID
             target_col=None, # colname of compound values (e.g., IC50)
             radius=2, # radius of the Morgan fingerprint.
             ):
    "Calculates the Tanimoto similarity scores between all pairs of molecules in a pandas DataFrame."
    
    df = df.copy()
    # Convert SMILES to molecule objects
    df['Molecule'] = df[smiles_col].apply(lambda x: Chem.MolFromSmiles(x))

    # Calculate fingerprints
    df['Fingerprint'] = df['Molecule'].apply(lambda x: AllChem.GetMorganFingerprintAsBitVect(x, radius))

    # Calculate similarity scores
    similarity_scores = []
    for i in range(len(df)):
        for j in range(i+1, len(df)):
            sim_score = DataStructs.TanimotoSimilarity(df['Fingerprint'][i], df['Fingerprint'][j])
            if target_col is not None:
                similarity_scores.append((df[id_col][i], df[id_col][j], df[smiles_col][i], df[smiles_col][j], sim_score, df[target_col][i], df[target_col][j]))
            else:
                similarity_scores.append((df[id_col][i], df[id_col][j], df[smiles_col][i], df[smiles_col][j], sim_score))

    # Create a new DataFrame with the similarity scores
    if target_col is not None:
        result_df = pd.DataFrame(similarity_scores, columns=['ID1', 'ID2', 'SMILES1', 'SMILES2', 'SimilarityScore', 'Target1', 'Target2'])
    else:
        result_df = pd.DataFrame(similarity_scores, columns=['ID1', 'ID2', 'SMILES1', 'SMILES2', 'SimilarityScore'])

    # Sort by similarity score in descending order
    result_df.sort_values('SimilarityScore', ascending=False, inplace=True)
    result_df = result_df.reset_index(drop=True)

    return result_df

# %% ../../nbs/core/03_ligand.ipynb 14
def get_rdkit(SMILES:str):
    """
    Extract chemical features from SMILES
    Reference: https://greglandrum.github.io/rdkit-blog/posts/2022-12-23-descriptor-tutorial.html
    """
    mol = Chem.MolFromSmiles(SMILES)
    return Descriptors.CalcMolDescriptors(mol)

# %% ../../nbs/core/03_ligand.ipynb 15
def get_rdkit_3d(SMILES:str):
    "Extract 3d features from SMILES"
    mol = Chem.MolFromSmiles(SMILES)
    mol = Chem.AddHs(mol)
    AllChem.EmbedMolecule(mol, AllChem.ETKDG())
    AllChem.UFFOptimizeMolecule(mol)
    return Descriptors3D.CalcMolDescriptors3D(mol)

# %% ../../nbs/core/03_ligand.ipynb 16
def get_rdkit_all(SMILES:str):
    "Extract chemical features and 3d features from SMILES"
    feat = get_rdkit(SMILES)
    feat_3d = get_rdkit_3d(SMILES)
    return feat|feat_3d

# %% ../../nbs/core/03_ligand.ipynb 17
def remove_hi_corr(df: pd.DataFrame, 
                   thr=0.99 # threshold
                   ):
    "Remove highly correlated features in a dataframe given a pearson threshold"
    corr_matrix = df.corr().abs()
    upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))
    to_drop = [column for column in upper.columns if any(upper[column] > thr)]
    return df.drop(to_drop, axis=1), to_drop

# %% ../../nbs/core/03_ligand.ipynb 18
def preprocess(df: pd.DataFrame, thr=0.99):
    "Remove features with no variance, and highly correlated features based on threshold."
    col_ori = df.columns

    # Remove columns with std == 0
    std_zero_cols = df.columns[df.std() == 0].tolist()
    
    if std_zero_cols:
        n=len(std_zero_cols)
        print(f"\n {n} Columns with zero std: {std_zero_cols}")
    df = df.loc[:, df.std() != 0].copy()

    # Remove highly correlated columns
    df, high_corr_cols = remove_hi_corr(df, thr)
    if high_corr_cols:
        n=len(high_corr_cols)
        print(f"\n {n} Columns removed due to high similarity (pearson>{thr}): {high_corr_cols}")

    dropping_col = set(col_ori) - set(df.columns)
    n = len(dropping_col)
    print(f"\n Total removed columns: {n}")
    
    return df

# %% ../../nbs/core/03_ligand.ipynb 19
def get_rdkit_df(df: pd.DataFrame,
                 include_3d=False,
                 col='SMILES',
                 postprocess=False,
                 chunksize=128, # for parallel process_map 
                 ):
    "Extract rdkit features (optionally in parallel with 3D) from SMILES in a df"
    func = get_rdkit_all if include_3d else get_rdkit
    smiles_list = df[col].tolist()

    features = process_map(func, smiles_list,chunksize=chunksize)

    out = pd.DataFrame(features)

    if postprocess:
        out = StandardScaler().fit_transform(out)
        out = preprocess(out) # remove redundant

    return out

# %% ../../nbs/core/03_ligand.ipynb 25
def get_fp(SMILES, 
           name='ecfp',
           ELEMENTS_PER_WORKER = 1_000_000):
    "Super fast method to get molecule fingerprints using scikit-fingerprints"
    if name == 'ecfp':
        fp_transformer = ECFPFingerprint(fp_size=2048, radius=3, n_jobs=-1)
    elif name == 'ap':
        fp_transformer = AtomPairFingerprint(fp_size=2048, n_jobs=-1)
    elif name == 'maccs':
        fp_transformer = MACCSFingerprint(n_jobs=-1)
    else:
        raise Exception('Wrong fingerprint name!')
    
    middle_parts = []
    k_splits = len(SMILES) // ELEMENTS_PER_WORKER
        
    for i in tqdm(range(k_splits)):
        middle_parts.append(fp_transformer.transform(SMILES[i * ELEMENTS_PER_WORKER: (i + 1) * ELEMENTS_PER_WORKER]))
        
    if len(SMILES) % ELEMENTS_PER_WORKER > 0:   
        middle_parts.append(fp_transformer.transform(SMILES[k_splits * ELEMENTS_PER_WORKER:]))
    
    return np.concatenate(middle_parts)

# %% ../../nbs/core/03_ligand.ipynb 29
def compress_fp(array):
    "Compress binary finterprints using np.packbits"
    return np.packbits(array,axis=1)

# %% ../../nbs/core/03_ligand.ipynb 47
@numba.njit(parallel=True)
def tanimoto_numba(fps):
    "Get a NxN matrix of tanimoto similarity among N compounds."
    n = fps.shape[0]
    result = np.zeros((n, n), dtype=np.float32)
    
    for i in numba.prange(n):
        for j in range(i, n):  # only upper triangle
            inter = np.bitwise_and(fps[i], fps[j]).sum()
            union = np.bitwise_or(fps[i], fps[j]).sum()
            sim = inter / union if union > 0 else 0.0
            result[i, j] = sim
            result[j, i] = sim  # fill symmetric position
    return result

# %% ../../nbs/core/03_ligand.ipynb 50
def hash_fp(fp_row):
    "Hash a binary fingerprint row using SHA256"
    return hashlib.sha256(fp_row.tobytes()).hexdigest()

# %% ../../nbs/core/03_ligand.ipynb 52
def get_same_mol_group(df, smi_col='SMILES'):
    "Assign a group number to the same compounds by utilizing hash sha256 to encode morgan fp and find same molecule."
    df = df.copy()
    smiles = df[smi_col].tolist()
    
    fps = get_fp(smiles)
    
    # Hash each fingerprint
    # hashes = [hash_fp(fp) for fp in fps] # non-parallel
    hashes = process_map(hash_fp, fps, chunksize=256) # if parallel
    
    # Assign a group number based on hash buckets
    hash_to_group = {}
    group_ids = []
    group_counter = 0
    for h in hashes:
        if h not in hash_to_group:
            hash_to_group[h] = group_counter
            group_counter += 1
        group_ids.append(hash_to_group[h])
    
    df['group'] = group_ids
    return df
