# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/core/01_utils.ipynb.

# %% auto 0
__all__ = ['rglob', 'copy_files', 'get_rec_lig', 'get_box', 'view_mol', 'view_complex']

# %% ../../nbs/core/01_utils.ipynb 3
from pathlib import Path
import subprocess,shutil,zipfile
import numpy as np

import py3Dmol
from rdkit import Chem

# %% ../../nbs/core/01_utils.ipynb 6
def rglob(path, pattern, max_depth):
    "Get a file list given folder depths"
    base_path = Path(path).resolve()
    for path in base_path.rglob(pattern):
        if len(path.relative_to(base_path).parts) <= max_depth:
            yield path

# %% ../../nbs/core/01_utils.ipynb 8
def copy_files(file_list, dest_dir):
    "Copy a list of files to the destination directory, or zip them if dest_dir ends with .zip."
    dest_path = Path(dest_dir)

    if dest_path.suffix == ".zip":
        with zipfile.ZipFile(dest_path, 'w') as zipf:
            for file_path in file_list:
                file_path = Path(file_path)
                zipf.write(file_path, arcname=file_path.name)
        print(f'Zipped {len(file_list)} files to {dest_path}')
    else:
        dest_path.mkdir(parents=True, exist_ok=True)
        for file_path in file_list:
            file_path = Path(file_path)
            shutil.copy2(file_path, dest_path / file_path.name)
        print(f'Copied {len(file_list)} files to {dest_path}')

# %% ../../nbs/core/01_utils.ipynb 11
def get_rec_lig(pdb_id: str, # pdb id for download
                            lig_id: str, # ligand id shown on the protein page
                            out_dir = '.', # directory path to save pdb files
                            ):
    "Download pdb and extract receptor and ligand from a PDB ID."
    out_dir = Path(out_dir).expanduser().resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    pdb_file = out_dir / f"{pdb_id}.pdb"
    rec_file = out_dir / f"{pdb_id}_receptor.pdb"
    lig_pdb_file = out_dir / f"{pdb_id}_lig.pdb"
    lig_sdf_file = out_dir / f"{pdb_id}_lig.sdf"

    # Download if not exists
    if not pdb_file.exists():
        url = f"http://files.rcsb.org/download/{pdb_id}.pdb"
        print(f'Downloading pdb: {pdb_id}')
        subprocess.run(["wget", url, "-O", str(pdb_file)], check=True,stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    print(f'{pdb_id}.pdb is detected!')

    # Extract protein (all ATOM lines excluding ligand ID)
    with open(pdb_file) as infile, open(rec_file, 'w') as out_rec:
        for line in infile:
            if line.startswith("ATOM") and lig_id not in line:
                out_rec.write(line)

    # Extract ligand
    with open(pdb_file) as infile, open(lig_pdb_file, 'w') as out_lig:
        for line in infile:
            if lig_id in line and line.startswith(("HETATM", "ATOM")):
                out_lig.write(line)

    # Convert ligand PDB to SDF using RDKit
    mol = Chem.MolFromPDBFile(str(lig_pdb_file), removeHs=False)
    if mol is None:
        raise ValueError("Failed to parse ligand from PDB.")
    
    writer = Chem.SDWriter(str(lig_sdf_file))
    writer.write(mol)
    writer.close()

    return str(rec_file), str(lig_sdf_file)

# %% ../../nbs/core/01_utils.ipynb 14
def get_box(sdf_file, autobox_add=4.0,tolist=False):
    "Get the box coordinates of ligand.sdf; mimic GNINA's --autobox_ligand behavior."
    mol = Chem.SDMolSupplier(str(sdf_file), removeHs=False)[0]
    if mol is None:
        raise ValueError(f"Failed to read molecule from {sdf_file}")
    
    conf = mol.GetConformer()
    coords = np.array([list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())])
    
    min_coords = coords.min(axis=0)
    max_coords = coords.max(axis=0)
    
    center = (min_coords + max_coords) / 2
    size = (max_coords - min_coords) + autobox_add

    box_dict = {
        "center_x": round(float(center[0]), 3),
        "center_y": round(float(center[1]), 3),
        "center_z": round(float(center[2]), 3),
        "size_x": round(float(size[0]), 3),
        "size_y": round(float(size[1]), 3),
        "size_z": round(float(size[2]), 3)
    }
    return list(box_dict.values()) if tolist else box_dict

# %% ../../nbs/core/01_utils.ipynb 18
def view_mol(file, #sdf or pdb file
            ):
    "Visualize pdb or sdf file"
    
    v = py3Dmol.view()
    v.addModel(open(file).read())
    v.setStyle({'stick':{}})
    v.zoomTo()
    v.show()

# %% ../../nbs/core/01_utils.ipynb 20
def view_complex(receptor,           # protein file
                 ligand,             # ligand (green), or docked ligand
                 ori_ligand=None,    # original ligand (yellow)
                 box=None            # optional box: [x, y, z, sizeX, sizeY, sizeZ]
                ):

    "Visualize the receptor, ligand, optional original ligand, and optional box via py3Dmol."
    v = py3Dmol.view()
    
    # Load receptor
    v.addModel(open(receptor).read())
    v.setStyle({'cartoon': {}, 'stick': {'radius': 0.15}})
    
    # Load docked ligand
    v.addModel(open(ligand).read())
    v.setStyle({'model': 1}, {'stick': {'colorscheme': 'greenCarbon'}})

    # Load original ligand if provided
    if ori_ligand is not None:
        v.addModel(open(ori_ligand).read())
        v.setStyle({'model': 2}, {'stick': {'colorscheme': 'yellowCarbon'}})

    # Add bounding box if specified
    if box is not None and len(box) == 6:
        x, y, z, sizeX, sizeY, sizeZ = box
        v.addBox({
            'center': {'x': x, 'y': y, 'z': z},
            'dimensions': {'w': sizeX, 'h': sizeY, 'd': sizeZ},
            'color': 'red',
            'opacity': 1,
            'wireframe': True
        })

    v.zoomTo({'model': 1})
    v.show()
