# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/gnina/05_gnina_AF3_rescore.ipynb.

# %% auto 0
__all__ = ['ChainSelect', 'rename_residues', 'split_cif', 'pdb2sdf', 'prepare_rec_lig', 'gnina_rescore_local',
           'gnina_rescore_docker', 'extract_gnina_rescore', 'get_gnina_rescore', 'get_gnina_rescore_folder']

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 3
import pandas as pd
import re, os, subprocess, py3Dmol
from Bio.PDB import MMCIFParser, PDBIO, Select
from rdkit import Chem
from rdkit.Chem import AllChem
from pathlib import Path
from fastcore.all import L
from tqdm.contrib.concurrent import process_map
from functools import partial

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 5
class ChainSelect(Select):
    "Select chain to save"
    def __init__(self, chain_ids):
        self.chain_ids = chain_ids
    def accept_chain(self, chain):
        return chain.get_id() in self.chain_ids

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 6
def rename_residues(structure, chain_id, new_resname='LIG'):
    "Rename residue name from LIG_L to LIG as LIG_L exceeds lengths and leads to error in RDKit"
    for model in structure:
        for chain in model:
            if chain.id == chain_id:
                for residue in chain:
                    residue.resname = new_resname

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 7
def split_cif(cif_path, rec_chain_id,lig_chain_id, rec_pdb_path, lig_pdb_path):
    "Split AF3 output CIF to protein and ligand PDBs"
    parser = MMCIFParser(QUIET=True)
    structure = parser.get_structure('complex', cif_path)
    rename_residues(structure, chain_id=lig_chain_id, new_resname='LIG')
    io = PDBIO()
    io.set_structure(structure)
    io.save(str(rec_pdb_path), ChainSelect(rec_chain_id))  # receptor
    io.save(str(lig_pdb_path), ChainSelect(lig_chain_id))  # ligand

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 8
def pdb2sdf(pdb_path, sdf_path):
    "Convert ligand pdb to sdf file"
    mol = Chem.MolFromPDBFile(pdb_path, sanitize=True, removeHs=False)
    if mol:
        writer = Chem.SDWriter(sdf_path)
        writer.write(mol)
        writer.close()
        return None
    else:
        print('Conversion failed for:', pdb_path)
        return pdb_path

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 9
def prepare_rec_lig(cif_path, rec_chain_id, lig_chain_id, rec_pdb_path, lig_pdb_path):
    "Split AF3 cif to protein.pdb (chainA) and ligand.sdf (chainL) "
    
    tmp_name = Path(cif_path).stem
    tmp_path = f'{tmp_name}_lig.pdb'
    split_cif(cif_path, rec_chain_id,lig_chain_id, rec_pdb_path, tmp_path)
    failed = pdb2sdf(tmp_path, lig_pdb_path)
    try:
        os.remove(tmp_path)
    except OSError:
        pass
    return failed

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 15
def gnina_rescore_local(protein_pdb,  # receptor file
                        ligand_sdf,   # ligand file
                        CNN_affinity=True,
                        vinardo=False, # if True, use vinardo instead of vina
                        ):

    command = ['./gnina', 
               '-r', protein_pdb, 
               '-l', ligand_sdf,
               '--minimize']  # always include this

    # Handle scoring options
    if not CNN_affinity:
        command += ['--cnn_scoring', 'none']
    if vinardo:
        command += ['--scoring', 'vinardo']

    result = subprocess.run(command, capture_output=True, text=True)
    return result.stdout

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 17
def gnina_rescore_docker(protein_pdb, 
                         ligand_sdf, 
                         CNN_affinity=True, 
                         vinardo=False):
    "Run GNINA rescoring using Docker. Supports receptor and ligand in different folders."

    protein_pdb = Path(protein_pdb).resolve()
    ligand_sdf = Path(ligand_sdf).resolve()

    # Mount points inside the Docker container
    rec_mount = '/recdata'
    lig_mount = '/ligdata'

    command = [
        'docker', 'run', '--rm',
        '-v', f'{protein_pdb.parent}:{rec_mount}',  # mount receptor directory
        '-v', f'{ligand_sdf.parent}:{lig_mount}',   # mount ligand directory
        'gnina/gnina',
        'gnina',
        '-r', f'{rec_mount}/{protein_pdb.name}',
        '-l', f'{lig_mount}/{ligand_sdf.name}',
        '--minimize',  # always include
    ]

    if not CNN_affinity:
        command += ['--cnn_scoring', 'none']
    if vinardo:
        command += ['--scoring', 'vinardo']

    result = subprocess.run(command, capture_output=True, text=True)
    return result.stdout

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 19
def extract_gnina_rescore(txt):
    """Extract GNINA output metrics into a dictionary (partial match allowed)."""
    result = {}

    patterns = {
        'binding_energy': r'Affinity:\s+([-.\d]+)',
        'uncertainty':    r'Affinity:\s+[-.\d]+\s+([-.\d]+)',
        'RMSD':           r'RMSD:\s+([-.\d]+)',
        'CNNscore':       r'CNNscore:\s+([-.\d]+)',
        'CNNaffinity':    r'CNNaffinity:\s+([-.\d]+)',
        'CNNvariance':    r'CNNvariance:\s+([-.\d]+)',
    }

    for key, pat in patterns.items():
        match = re.search(pat, txt)
        if match:
            result[key] = float(match.group(1))

    return result

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 24
def get_gnina_rescore(cif_path,
                      rec_chain_id='A', 
                      lig_chain_id='L', 
                      CNN_affinity=True,
                      vinardo=False,
                      is_local=True):
    "Split the CIF into receptor and ligand folders, then extract the GNINA rescored affinity score"
    cif_path = Path(cif_path).expanduser()
    parent,stem = cif_path.parent,cif_path.stem

    rec_dir,lig_dir = Path(str(parent) + '_receptor'),Path(str(parent) + '_ligand')
    
    rec_path,lig_path = rec_dir/f'{stem}.pdb',lig_dir/f'{stem}.sdf'
    
    rec_dir.mkdir(exist_ok=True)
    lig_dir.mkdir(exist_ok=True)
    
    prepare_rec_lig(cif_path,rec_chain_id, lig_chain_id,rec_path,lig_path)
    if is_local:
        gnina_output = gnina_rescore_local(rec_path,lig_path,CNN_affinity,vinardo)
    else:
        gnina_output = gnina_rescore_docker(rec_path,lig_path,CNN_affinity,vinardo)
    return extract_gnina_rescore(gnina_output)

# %% ../../nbs/gnina/05_gnina_AF3_rescore.ipynb 29
def get_gnina_rescore_folder(cif_folder,
                             rec_chain_id='A', 
                             lig_chain_id='L',
                             CNN_affinity=True,
                             vinardo=False,
                             is_local=True):
    "Parallel processing to get gnina rescore given folder path"
    cifs = L(Path(cif_folder).expanduser().glob("*.cif")) # just take cif file
    
    func = partial(get_gnina_rescore,
                   rec_chain_id=rec_chain_id, 
                   lig_chain_id=lig_chain_id,
                   CNN_affinity=CNN_affinity,
                   vinardo=vinardo,
                   is_local=is_local)
    results = process_map(func, cifs, max_workers=4)

    # use path.stem as df index
    results_dict = dict(zip([p.stem for p in cifs], results))
    result_df = pd.DataFrame(results_dict).T.reset_index(names='ID')
    
    prefix = "vinardo_" if vinardo else "vina_"
    result_df = result_df.rename(columns={
        "binding_energy": f"{prefix}binding_energy",
        "uncertainty": f"{prefix}uncertainty",
        "RMSD": f"{prefix}RMSD"
    })

    return result_df

    
