# 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 copy
import torch
import pandas as pd
from rdkit import Chem
import pickle
import argparse
from docking_utils import rmsd_func
import warnings

warnings.filterwarnings(action="ignore")


def single_SF_loss(
    predict_coords,
    pocket_coords,
    distance_predict,
    holo_distance_predict,
    dist_threshold=4.5,
):
    dist = torch.norm(predict_coords.unsqueeze(1) - pocket_coords.unsqueeze(0), dim=-1)
    holo_dist = torch.norm(
        predict_coords.unsqueeze(1) - predict_coords.unsqueeze(0), dim=-1
    )
    distance_mask = distance_predict < dist_threshold
    cross_dist_score = (
        (dist[distance_mask] - distance_predict[distance_mask]) ** 2
    ).mean()
    dist_score = ((holo_dist - holo_distance_predict) ** 2).mean()
    loss = cross_dist_score * 1.0 + dist_score * 5.0
    return loss


def scoring(
    predict_coords,
    pocket_coords,
    distance_predict,
    holo_distance_predict,
    dist_threshold=4.5,
):
    predict_coords = predict_coords.detach()
    dist = torch.norm(predict_coords.unsqueeze(1) - pocket_coords.unsqueeze(0), dim=-1)
    holo_dist = torch.norm(
        predict_coords.unsqueeze(1) - predict_coords.unsqueeze(0), dim=-1
    )
    distance_mask = distance_predict < dist_threshold
    cross_dist_score = (
        (dist[distance_mask] - distance_predict[distance_mask]) ** 2
    ).mean()
    dist_score = ((holo_dist - holo_distance_predict) ** 2).mean()
    return cross_dist_score.numpy(), dist_score.numpy()


def dock_with_gradient(
    coords,
    pocket_coords,
    distance_predict_tta,
    holo_distance_predict_tta,
    loss_func=single_SF_loss,
    holo_coords=None,
    iterations=20000,
    early_stoping=5,
):
    bst_loss, bst_coords, bst_meta_info = 10000.0, coords, None
    for i, (distance_predict, holo_distance_predict) in enumerate(
        zip(distance_predict_tta, holo_distance_predict_tta)
    ):
        new_coords = copy.deepcopy(coords)
        _coords, _loss, _meta_info = single_dock_with_gradient(
            new_coords,
            pocket_coords,
            distance_predict,
            holo_distance_predict,
            loss_func=loss_func,
            holo_coords=holo_coords,
            iterations=iterations,
            early_stoping=early_stoping,
        )
        if bst_loss > _loss:
            bst_coords = _coords
            bst_loss = _loss
            bst_meta_info = _meta_info
    return bst_coords, bst_loss, bst_meta_info


def single_dock_with_gradient(
    coords,
    pocket_coords,
    distance_predict,
    holo_distance_predict,
    loss_func=single_SF_loss,
    holo_coords=None,
    iterations=20000,
    early_stoping=5,
):
    coords = torch.from_numpy(coords).float()
    pocket_coords = torch.from_numpy(pocket_coords).float()
    distance_predict = torch.from_numpy(distance_predict).float()
    holo_distance_predict = torch.from_numpy(holo_distance_predict).float()

    if holo_coords is not None:
        holo_coords = torch.from_numpy(holo_coords).float()

    coords.requires_grad = True
    optimizer = torch.optim.LBFGS([coords], lr=1.0)
    bst_loss, times = 10000.0, 0
    for i in range(iterations):

        def closure():
            optimizer.zero_grad()
            loss = loss_func(
                coords, pocket_coords, distance_predict, holo_distance_predict
            )
            loss.backward()
            return loss

        loss = optimizer.step(closure)
        if loss.item() < bst_loss:
            bst_loss = loss.item()
            times = 0
        else:
            times += 1
            if times > early_stoping:
                break

    meta_info = scoring(coords, pocket_coords, distance_predict, holo_distance_predict)
    return coords.detach().numpy(), loss.detach().numpy(), meta_info


def set_coord(mol, coords):
    for i in range(coords.shape[0]):
        mol.GetConformer(0).SetAtomPosition(i, coords[i].tolist())
    return mol


def add_coord(mol, xyz):
    x, y, z = xyz
    conf = mol.GetConformer(0)
    pos = conf.GetPositions()
    pos[:, 0] += x
    pos[:, 1] += y
    pos[:, 2] += z
    for i in range(pos.shape[0]):
        conf.SetAtomPosition(
            i, Chem.rdGeometry.Point3D(pos[i][0], pos[i][1], pos[i][2])
        )
    return mol


def single_docking(input_path, output_path, output_ligand_path):
    content = pd.read_pickle(input_path)
    (
        init_coords_tta,
        mol,
        smi,
        pocket,
        pocket_coords,
        distance_predict_tta,
        holo_distance_predict_tta,
        holo_coords,
        holo_cener_coords,
    ) = content
    sample_times = len(init_coords_tta)
    bst_predict_coords, bst_loss, bst_meta_info = None, 1000.0, None
    for i in range(sample_times):
        init_coords = init_coords_tta[i]
        predict_coords, loss, meta_info = dock_with_gradient(
            init_coords,
            pocket_coords,
            distance_predict_tta,
            holo_distance_predict_tta,
            holo_coords=holo_coords,
            loss_func=single_SF_loss,
        )
        if loss < bst_loss:
            bst_loss = loss
            bst_predict_coords = predict_coords
            bst_meta_info = meta_info

    _rmsd = round(rmsd_func(holo_coords, bst_predict_coords), 4)
    _cross_score = round(float(bst_meta_info[0]), 4)
    _self_score = round(float(bst_meta_info[1]), 4)
    print(f"{pocket}-{smi}-RMSD:{_rmsd}-{_cross_score}-{_self_score}")
    mol = Chem.RemoveHs(mol)
    mol = set_coord(mol, bst_predict_coords)

    if output_path is not None:
        with open(output_path, "wb") as f:
            pickle.dump(
                [bst_predict_coords, holo_coords, bst_loss, smi, pocket, pocket_coords],
                f,
            )
    if output_ligand_path is not None:
        mol = add_coord(mol, holo_cener_coords.numpy())
        Chem.MolToMolFile(mol, output_ligand_path)

    return True


if __name__ == "__main__":
    torch.set_num_threads(1)
    torch.manual_seed(0)
    parser = argparse.ArgumentParser(description="Docking with gradient")
    parser.add_argument("--input", type=str, help="input file.")
    parser.add_argument("--output", type=str, default=None, help="output path.")
    parser.add_argument(
        "--output-ligand", type=str, default=None, help="output ligand sdf path."
    )
    args = parser.parse_args()

    single_docking(args.input, args.output, args.output_ligand)