# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/gnina/04_gnina_docking.ipynb.

# %% auto 0
__all__ = ['setup_gnina_local', 'setup_gnina_docker', 'extract_gnina_dock', 'gnina_dock']

# %% ../../nbs/gnina/04_gnina_docking.ipynb 3
# basics
import re,subprocess, py3Dmol
from tqdm import tqdm
from pathlib import Path
import pandas as pd,numpy as np

# rdkit
from rdkit import Chem
from rdkit.Chem import AllChem

# %% ../../nbs/gnina/04_gnina_docking.ipynb 8
def setup_gnina_local(version='v1.3'):
    "Download and install gnina in the current directory"
    # Check CUDA availability
    # try:
    #     subprocess.run(["nvidia-smi"], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    # except subprocess.CalledProcessError:
    #     raise EnvironmentError("CUDA not detected. Please make sure a CUDA-capable GPU is available and drivers are installed.")
    # except FileNotFoundError:
    #     raise EnvironmentError("nvidia-smi not found. Make sure NVIDIA drivers and CUDA are installed.")

    subprocess.run(["sudo", "apt-get", "update", "-yq"], check=True,stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    subprocess.run(["sudo", "apt-get", "install", "-yq", "openbabel"], check=True)

    gnina_url = f"https://github.com/gnina/gnina/releases/download/{version}/gnina"
    print(f'Downloading {version} gnina')
    subprocess.run(["wget",gnina_url], check=True,stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    subprocess.run(["chmod", "+x", 'gnina'], check=True)
    
    print('Finish setup!')

# %% ../../nbs/gnina/04_gnina_docking.ipynb 13
def setup_gnina_docker():
    "Pull gnina docker image"
    print("Pulling GNINA Docker image: gnina/gnina")
    subprocess.run(["docker", "pull", "gnina/gnina"], check=True,stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    print("GNINA Docker image is ready.")

# %% ../../nbs/gnina/04_gnina_docking.ipynb 18
def extract_gnina_dock(gnina_output):
    "Extract values from gnina output"
    mode1_line = re.search(r'\b1\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\b', gnina_output)
    
    if mode1_line:
        affinity = float(mode1_line.group(1))
        cnn_pose_score = float(mode1_line.group(2))
        cnn_affinity = float(mode1_line.group(3))
        
        return affinity, cnn_pose_score, cnn_affinity
    
    return None

# %% ../../nbs/gnina/04_gnina_docking.ipynb 19
def gnina_dock(receptor, # receptor file
              ligand, # ligand file
              autobox_ligand, # ligand file isolated from the complex
              output = 'docked.sdf', # output file (sdf or sdf.gz) to be saved
              seed=0, # set seeds
              exhaustiveness=None, # number of MC chains, default is 8 if None, the higher the better (16,32); for whole protein, use 64
              ):
    
    command = ['./gnina', 
               '-r', receptor, 
               '-l', ligand, 
               '--autobox_ligand', autobox_ligand,
               '-o', output,
               '--seed', str(seed)]

    if exhaustiveness is not None:
        command.extend(['--exhaustiveness', str(exhaustiveness)])

    output_txt = subprocess.run(command, capture_output=True, text=True).stdout
    
    print(f'save the docked file as {output}')
    
    values = extract_gnina_dock(output_txt)
    
    print(f'affinity, cnn_pose_score, and cnn_affinity are: {values}')

    return values

# %% ../../nbs/gnina/04_gnina_docking.ipynb 21
def gnina_dock(df, 
                 ID_col = 'ID',
                 smi_col = 'SMILES',
              output_dir = 'gnina_docked'
                 ):
    affinity_values = []
    cnn_pose_score_values = []
    cnn_affinity_values = []

    
    Path(output_dir).mkdir(parents=True,exist_ok=True)

    for i, r in tqdm(df.iterrows(),total=len(df),desc='Docking'):
        rdkit_conformer(SMILES=r[smi_col], output = f'ligand/{r[ID_col]}.sdf', visualize=False)
        affinity, cnn_pose_score, cnn_affinity = gnina_dock('rec.pdb',f'ligand/{r[ID_col]}.sdf', 'lig.pdb',f'docked/docked_{r[ID_col]}.sdf')

        affinity_values.append(affinity)
        cnn_pose_score_values.append(cnn_pose_score)
        cnn_affinity_values.append(cnn_affinity)

    df = df.copy()
    df['Affinity'] = affinity_values
    df['CNN_Pose_Score'] = cnn_pose_score_values
    df['CNN_Affinity'] = cnn_affinity_values
    
    return df
