# small script to extract the ligand and save it in a separate file because GNINA will use the ligand position as
# initial pose
import os
import shutil
import subprocess
import sys

import time
from argparse import ArgumentParser, FileType
from datetime import datetime

import numpy as np
import pandas as pd
from biopandas.pdb import PandasPdb
from rdkit import Chem
from rdkit.Chem import AllChem, MolToPDBFile
from rdkit.Geometry import Point3D
from scipy.spatial.distance import cdist

from datasets.pdbbind import read_mol
from utils.utils import read_strings_from_txt

parser = ArgumentParser()
parser.add_argument('--data_dir', type=str, default='data/PDBBIND_atomCorrected', help='')
parser.add_argument('--file_suffix', type=str, default='_baseline_ligand', help='Path to folder with trained model and hyperparameters')
parser.add_argument('--protein_file_suffix', type=str, default='_protein_processed', help='Path to folder with trained model and hyperparameters')
parser.add_argument('--results_path', type=str, default='results/gnina_predictions', help='')
parser.add_argument('--complex_names_path', type=str, default='data/splits/timesplit_test', help='')
parser.add_argument('--seed_molecules_path', type=str, default=None, help='Use the molecules at seed molecule path as initialization and only search around them')
parser.add_argument('--seed_molecule_filename', type=str, default='equibind_corrected.sdf', help='Use the molecules at seed molecule path as initialization and only search around them')
parser.add_argument('--seed_molecule_includes_name', action='store_true', default=False, help='additionally add the complex name at the beginning of the seedmolecule path')
parser.add_argument('--smina', action='store_true', default=False, help='')
parser.add_argument('--no_gpu', action='store_true', default=False, help='')
parser.add_argument('--exhaustiveness', type=int, default=8, help='')
parser.add_argument('--num_cpu', type=int, default=16, help='')
parser.add_argument('--pocket_mode', action='store_true', default=False, help='')
parser.add_argument('--pocket_cutoff', type=float, default=5, help='')
parser.add_argument('--num_modes', type=int, default=10, help='')
parser.add_argument('--autobox_add', type=int, default=4, help='')
parser.add_argument('--use_p2rank_pocket', action='store_true', default=False, help='')
parser.add_argument('--skip_p2rank', action='store_true', default=False, help='')
parser.add_argument('--prank_path', type=str, default='', help='')
parser.add_argument('--skip_existing', action='store_true', default=False, help='')
parser.add_argument('--flexdist', type=float, default=-1, help='If this is -1 then no sidechains are flexible')
parser.add_argument('--flex_max', type=int, default=-1, help='If this is -1 then no sidechains are flexible')
parser.add_argument('--gnina_path', type=str, default='', help='Path to the gnina executable')


args = parser.parse_args()

class Logger(object):
    def __init__(self, logpath, syspart=sys.stdout):
        self.terminal = syspart
        self.log = open(logpath, "a")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)
        self.log.flush()

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        # this handles the flush command by doing nothing.
        # you might want to specify some extra behavior here.
        pass


def log(*args):
    print(f'[{datetime.now()}]', *args)


# parameters
names = read_strings_from_txt(args.complex_names_path)

if os.path.exists(args.results_path) and not args.skip_existing:
    shutil.rmtree(args.results_path)
os.makedirs(args.results_path, exist_ok=True)
sys.stdout = Logger(logpath=f'{args.results_path}/gnina.log', syspart=sys.stdout)
sys.stderr = Logger(logpath=f'{args.results_path}/error.log', syspart=sys.stderr)

p2rank_cache_path = "results/.p2rank_cache"
if args.use_p2rank_pocket and not args.skip_p2rank:
    os.makedirs(p2rank_cache_path, exist_ok=True)
    pdb_files_cache = os.path.join(p2rank_cache_path,'pdb_files')
    os.makedirs(pdb_files_cache, exist_ok=True)
    with open(f"{p2rank_cache_path}/pdb_list_p2rank.txt", "w") as out:
        for name in names:
            if os.path.exists(os.path.join(args.data_dir, name, f'{name}{args.protein_file_suffix}.pdb')):
                shutil.copy(os.path.join(args.data_dir, name, f'{name}{args.protein_file_suffix}.pdb'), f'{pdb_files_cache}/{name}{args.protein_file_suffix}.pdb')
                out.write(os.path.join('pdb_files', f'{name}{args.protein_file_suffix}.pdb\n'))
    cmd = f"bash {args.prank_path} predict {p2rank_cache_path}/pdb_list_p2rank.txt -o {p2rank_cache_path}/p2rank_output -threads 4"
    os.system(cmd)


all_times = []
start_time = time.time()
for i, name in enumerate(names):
    os.makedirs(os.path.join(args.results_path, name), exist_ok=True)
    log('\n')
    log(f'complex {i} of {len(names)}')
    # call gnina to find binding pose
    rec_path = os.path.join(args.data_dir, name, f'{name}{args.protein_file_suffix}.pdb')
    prediction_output_name = os.path.join(args.results_path, name, f'{name}{args.file_suffix}.pdb')
    flexRes_output_name = os.path.join(args.results_path, name, f'{name}_flex_residues.pdb')
    log_path = os.path.join(args.results_path, name, f'{name}{args.file_suffix}.log')
    if args.seed_molecules_path is not None: seed_mol_path = os.path.join(args.seed_molecules_path, name, f'{name if args.seed_molecule_includes_name else ""}{args.seed_molecule_filename}')
    if args.skip_existing and os.path.exists(prediction_output_name): continue

    if args.pocket_mode:
        mol = read_mol(args.data_dir, name, remove_hs=False)
        rec = PandasPdb().read_pdb(rec_path)
        rec_df = rec.get(s='c-alpha')
        rec_pos = rec_df[['x_coord', 'y_coord', 'z_coord']].to_numpy().squeeze().astype(np.float32)
        lig_pos = mol.GetConformer().GetPositions()
        d = cdist(rec_pos, lig_pos)
        label = np.any(d < args.pocket_cutoff, axis=1)

        if np.any(label):
            center_pocket = rec_pos[label].mean(axis=0)
        else:
            print("No pocket residue below minimum distance ", args.pocket_cutoff, "taking closest at", np.min(d))
            center_pocket = rec_pos[np.argmin(np.min(d, axis=1)[0])]
        radius_pocket = np.max(np.linalg.norm(lig_pos - center_pocket[None, :], axis=1))
        diameter_pocket = radius_pocket * 2
        center_x = center_pocket[0]
        size_x = diameter_pocket + 8
        center_y = center_pocket[1]
        size_y = diameter_pocket + 8
        center_z = center_pocket[2]
        size_z = diameter_pocket + 8


    mol_rdkit = read_mol(args.data_dir, name, remove_hs=False)
    single_time = time.time()

    mol_rdkit.RemoveAllConformers()
    ps = AllChem.ETKDGv2()
    id = AllChem.EmbedMolecule(mol_rdkit, ps)
    if id == -1:
        print('rdkit pos could not be generated without using random pos. using random pos now.')
        ps.useRandomCoords = True
        AllChem.EmbedMolecule(mol_rdkit, ps)
        AllChem.MMFFOptimizeMolecule(mol_rdkit, confId=0)
    rdkit_mol_path = os.path.join(args.data_dir, name, f'{name}_rdkit_ligand.pdb')
    MolToPDBFile(mol_rdkit, rdkit_mol_path)

    fallback_without_p2rank = False
    if args.use_p2rank_pocket:
        if os.path.exists(f'{p2rank_cache_path}/p2rank_output/{name}{args.protein_file_suffix}.pdb_predictions.csv'):
            df = pd.read_csv(f'{p2rank_cache_path}/p2rank_output/{name}{args.protein_file_suffix}.pdb_predictions.csv')
            rdkit_lig_pos = mol_rdkit.GetConformer().GetPositions()
            diameter_pocket = np.max(cdist(rdkit_lig_pos, rdkit_lig_pos))
            size_x = diameter_pocket + args.autobox_add * 2
            size_y = diameter_pocket + args.autobox_add * 2
            size_z = diameter_pocket + args.autobox_add * 2
            if df.empty:
                fallback_without_p2rank = True
            else:
                center_x = df.iloc[0]['   center_x']
                center_y = df.iloc[0]['   center_y']
                center_z = df.iloc[0]['   center_z']
                if args.flexdist > 0:
                    mol = Chem.MolFromSmiles('C')
                    AllChem.EmbedMolecule(mol)
                    conf = mol.GetConformer()
                    conf.SetAtomPosition(0, Point3D(center_x, center_y, center_z))
                    pocket_center_cache_path = f'{p2rank_cache_path}/p2rank_output/{name}_pocketcenter_as_carbon.pdb'
                    MolToPDBFile(mol, pocket_center_cache_path)
        else:
            fallback_without_p2rank = True


    # create flexibility arguments
    if args.flexdist > 0 and args.seed_molecules_path is not None and not args.use_p2rank_pocket: # for EquiBind plus GNINA flexible and GNINA rigid plus GNINA flexible
        flex_args = f'--flexdist {args.flexdist} --flexdist_ligand {seed_mol_path} {"--flex_max " + str(args.flex_max) if args.flex_max > 0 else ""} --out_flex {flexRes_output_name}'
    elif args.flexdist > 0 and args.seed_molecules_path is None and not args.use_p2rank_pocket: # for GNINA flexible (all sidechains flexible with high enough flexdist)
        flex_args = f'--flexdist {args.flexdist} --flexdist_ligand {rdkit_mol_path} {"--flex_max " + str(args.flex_max) if args.flex_max > 0 else ""} --out_flex {flexRes_output_name}' 
    elif args.flexdist > 0 and args.seed_molecules_path is None and args.use_p2rank_pocket and not fallback_without_p2rank: # for GNINA flexible (all sidechains flexible with high enough flexdist)
        flex_args = f'--flexdist {args.flexdist + diameter_pocket/2} --flexdist_ligand {pocket_center_cache_path} {"--flex_max " + str(args.flex_max) if args.flex_max > 0 else ""} --out_flex {flexRes_output_name}'
    else:
        flex_args = ''

    log(f'processing {rec_path}')
    if not args.pocket_mode and not args.use_p2rank_pocket or fallback_without_p2rank:
        return_code = subprocess.run(
            f"{args.gnina_path} --receptor {rec_path} --ligand {rdkit_mol_path} --num_modes {args.num_modes} -o {prediction_output_name} {flex_args} {'--no_gpu' if args.no_gpu else ''} --autobox_ligand {rec_path if args.seed_molecules_path is None else seed_mol_path} --autobox_add {args.autobox_add} --log {log_path} --exhaustiveness {args.exhaustiveness} --cpu {args.num_cpu} {'--cnn_scoring none' if args.smina else ''}",
            shell=True)
    else:
        return_code = subprocess.run(
            f"{args.gnina_path} --receptor {rec_path} --ligand {rdkit_mol_path} --num_modes {args.num_modes} -o {prediction_output_name} {flex_args} {'--no_gpu' if args.no_gpu else ''} --log {log_path} --exhaustiveness {args.exhaustiveness} --cpu {args.num_cpu} {'--cnn_scoring none' if args.smina else ''} --center_x {center_x} --center_y {center_y} --center_z {center_z} --size_x {size_x} --size_y {size_y} --size_z {size_z}",
            shell=True)
    log(return_code)
    all_times.append(time.time() - single_time)

    log("single time: --- %s seconds ---" % (time.time() - single_time))
    log("time so far: --- %s seconds ---" % (time.time() - start_time))
    log('\n')
log(all_times)
log("--- %s seconds ---" % (time.time() - start_time))
