# Copyright (c) DP Techonology, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import os
import numpy as np
import pandas as pd
from multiprocessing import Pool
from tqdm import tqdm
import glob
import argparse
from docking_utils import (
    docking_data_pre,
    ensemble_iterations,
    print_results,
    rmsd_func,
)
import warnings

warnings.filterwarnings(action="ignore")


def result_log(dir_path):
    ### result logging ###
    output_dir = os.path.join(dir_path, "cache")
    rmsd_results = []
    for path in glob.glob(os.path.join(output_dir, "*.docking.pkl")):
        (
            bst_predict_coords,
            holo_coords,
            bst_loss,
            smi,
            pocket,
            pocket_coords,
        ) = pd.read_pickle(path)
        rmsd = rmsd_func(holo_coords, bst_predict_coords)
        rmsd_results.append(rmsd)
    rmsd_results = np.array(rmsd_results)
    print_results(rmsd_results)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="docking")
    parser.add_argument(
        "--reference-file",
        type=str,
        default="./protein_ligand_binding_pose_prediction/test.lmdb",
        help="Location of the reference set",
    )
    parser.add_argument("--nthreads", type=int, default=40, help="num of threads")
    parser.add_argument(
        "--predict-file",
        type=str,
        default="./infer_pose/save_pose_test.out.pkl",
        help="Location of the prediction file",
    )
    parser.add_argument(
        "--output-path",
        type=str,
        default="./protein_ligand_binding_pose_prediction",
        help="Location of the docking output path",
    )
    args = parser.parse_args()

    raw_data_path, predict_path, dir_path, nthreads = (
        args.reference_file,
        args.predict_file,
        args.output_path,
        args.nthreads,
    )
    tta_times = 10
    (
        mol_list,
        smi_list,
        pocket_list,
        pocket_coords_list,
        distance_predict_list,
        holo_distance_predict_list,
        holo_coords_list,
        holo_center_coords_list,
    ) = docking_data_pre(raw_data_path, predict_path)
    iterations = ensemble_iterations(
        mol_list,
        smi_list,
        pocket_list,
        pocket_coords_list,
        distance_predict_list,
        holo_distance_predict_list,
        holo_coords_list,
        holo_center_coords_list,
        tta_times=tta_times,
    )
    sz = len(mol_list) // tta_times
    new_pocket_list = pocket_list[::tta_times]
    output_dir = os.path.join(dir_path, "cache")
    os.makedirs(output_dir, exist_ok=True)

    def dump(content):
        pocket = content[3]
        output_name = os.path.join(output_dir, "{}.pkl".format(pocket))
        try:
            os.remove(output_name)
        except:
            pass
        pd.to_pickle(content, output_name)
        return True

    with Pool(nthreads) as pool:
        for inner_output in tqdm(pool.imap(dump, iterations), total=sz):
            if not inner_output:
                print("fail to dump")

    def single_docking(pocket_name):
        input_name = os.path.join(output_dir, "{}.pkl".format(pocket_name))
        output_name = os.path.join(output_dir, "{}.docking.pkl".format(pocket_name))
        output_ligand_name = os.path.join(
            output_dir, "{}.ligand.sdf".format(pocket_name)
        )
        try:
            os.remove(output_name)
        except:
            pass
        try:
            os.remove(output_ligand_name)
        except:
            pass
        cmd = "python ./unimol/utils/coordinate_model.py --input {} --output {} --output-ligand {}".format(
            input_name, output_name, output_ligand_name
        )
        os.system(cmd)
        return True

    with Pool(nthreads) as pool:
        for inner_output in tqdm(
            pool.imap(single_docking, new_pocket_list), total=len(new_pocket_list)
        ):
            if not inner_output:
                print("fail to docking")

    result_log(args.output_path)