# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/protenix/06_proteinix_dock.ipynb.

# %% auto 0
__all__ = ['capture_output', 'get_protenix_dock', 'get_protenix_vina_dock', 'json2sdf']

# %% ../../nbs/protenix/06_proteinix_dock.ipynb 11
# kdock
from ..core.utils import *

# basics
import json,shutil
import os,sys,logging,contextlib
from datetime import datetime
from pathlib import Path

# import protenix-dock
try:
    from pxdock import ProtenixDock
    from pxdock.common.rmsd_calculator import write_ligand_to_sdf

    # Prevent multiprocessing-related errors in pxdock
    import multiprocessing as mp
    mp.set_start_method("spawn", force=True)

except ImportError as e:
    print(f"pxdock import failed: {e}")
    ProtenixDock = None
    write_ligand_to_sdf = None


# %% ../../nbs/protenix/06_proteinix_dock.ipynb 15
@contextlib.contextmanager
def capture_output(log_path):
    with open(log_path, 'a') as log_file:  # Append mode prevents overwriting
        # Redirect stdout/stderr
        old_stdout_fd = os.dup(1)
        old_stderr_fd = os.dup(2)
        os.dup2(log_file.fileno(), 1)
        os.dup2(log_file.fileno(), 2)

        # Setup Python logger
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)
        logger.handlers = [h for h in logger.handlers if not isinstance(h, logging.StreamHandler)]

        file_handler = logging.FileHandler(log_path)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

        try:
            yield
        finally:
            os.dup2(old_stdout_fd, 1)
            os.dup2(old_stderr_fd, 2)
            os.close(old_stdout_fd)
            os.close(old_stderr_fd)

# %% ../../nbs/protenix/06_proteinix_dock.ipynb 16
def get_protenix_dock(receptor_pdb, # pdb path
                      ligand_sdf, # sdf path
                      box:list, # center xyz + size xyz
                      out_dir='.', # output directory
                      num_walker=20, # number of Monte Carlo threads, the more walkers the more space to explore
                      mc_prune_energy_threshold=500, # maximum allowed energy for a pose; if exceeds, skipped
                      include_affinity=True, # include bscore
                      **kwargs):
    "Use protenix-dock to dock ligand sdf to receptor pdb using protenix-dock"
    Path(out_dir).mkdir(parents=True,exist_ok=True)
    
    dock_instance = ProtenixDock(receptor_pdb)
    box_center,box_size=box[:3],box[3:]
    dock_instance.set_box(box_center, box_size)
    
    log_path = Path(out_dir)/f'{datetime.now():%Y%m%d_%H%M%S}.log'
    with capture_output(log_path):
        # cache_dir = dock_instance.generate_cache_maps(spacing=0.175)
        result_dir = dock_instance.run_docking(ligand_sdf,out_dir=out_dir,num_walker=num_walker,mc_prune_energy_threshold=mc_prune_energy_threshold,include_affinity=include_affinity,**kwargs)
        # dock_instance.drop_cache_maps()
    
    print(f'Docking complete. Results in: {result_dir}')
    print('Cache dir:', cache_dir)

# %% ../../nbs/protenix/06_proteinix_dock.ipynb 20
def get_protenix_vina_dock(receptor_pdb, # pdb path
                      ligand_sdf, # sdf path
                      box:list, # center xyz + size xyz
                      out_dir='.', # output directory
                      **kwargs):
    "Dock ligand sdf to receptor pdb using protenix-dock"
    Path(out_dir).mkdir(parents=True,exist_ok=True)
    
    dock_instance = ProtenixDock(receptor_pdb)
    box_center,box_size=box[:3],box[3:]
    dock_instance.set_box(box_center, box_size)
    
    log_path = Path(out_dir)/f'{datetime.now():%Y%m%d_%H%M%S}.log'
    with capture_output(log_path):
        result_dir = dock_instance.run_pose_opt(ligand_sdf,out_dir=out_dir,**kwargs)

    print(f'Docking complete. Results in: {result_dir}')

# %% ../../nbs/protenix/06_proteinix_dock.ipynb 23
def json2sdf(json_path,
             sdf_path=None, # .sdf to be saved
             ):
    with open(json_path, "r") as f:
        json_out_data = json.load(f)

    best_pose = json_out_data["best_pose"]
    best_id = best_pose["index"]
    bscore = best_pose["bscore"]
    best_lig = json_out_data["poses"][best_id]
    xyz=best_lig["ligand"]["xyz"]
    if sdf_path is not None:
        write_ligand_to_sdf(json_out_data["mapped_smiles"], [xyz], sdf_path)
    return {'energy':best_lig["energy"],'pscore':best_lig["pscore"],'bscore':bscore}
