# Copyright (c) DP Technology.

# This source code is licensed under the MIT license found in the

# LICENSE file in the root directory of this source tree.

from IPython import embed as debug_embedded

import logging

import os

from collections.abc import Iterable

from sklearn.metrics import roc_auc_score

from xmlrpc.client import Boolean

import numpy as np

import torch

import pickle

from tqdm import tqdm

from unicore import checkpoint_utils

import unicore

from unicore.data import (AppendTokenDataset, Dictionary, EpochShuffleDataset,

                          FromNumpyDataset, NestedDictionaryDataset,

                          PrependTokenDataset, RawArrayDataset,LMDBDataset, RawLabelDataset,

                          RightPadDataset, RightPadDataset2D, TokenizeDataset,SortDataset,data_utils)

from unicore.tasks import UnicoreTask, register_task

from unimol.data import (AffinityDataset, CroppingPocketDataset,

                         CrossDistanceDataset, DistanceDataset,

                         EdgeTypeDataset, KeyDataset, LengthDataset,

                         NormalizeDataset, NormalizeDockingPoseDataset,

                         PrependAndAppend2DDataset, RemoveHydrogenDataset,

                         RemoveHydrogenPocketDataset, RightPadDatasetCoord,

                         RightPadDatasetCross2D, TTADockingPoseDataset, AffinityTestDataset, AffinityValidDataset, AffinityMolDataset, AffinityPocketDataset, ResamplingDataset)

#from skchem.metrics import bedroc_score

from rdkit.ML.Scoring.Scoring import CalcBEDROC, CalcAUC, CalcEnrichment

from sklearn.metrics import roc_curve

logger = logging.getLogger(__name__)

def re_new(y_true, y_score, ratio):

    fp = 0

    tp = 0

    p = sum(y_true)

    n = len(y_true) - p

    num = ratio*n

    sort_index = np.argsort(y_score)[::-1]

    for i in range(len(sort_index)):

        index = sort_index[i]

        if y_true[index] == 1:

            tp += 1

        else:

            fp += 1

            if fp>= num:

                break

    return (tp*n)/(p*fp)

def calc_re(y_true, y_score, ratio_list):

    fpr, tpr, thresholds = roc_curve(y_true, y_score, pos_label=1)

    #print(fpr, tpr)

    res = {}

    res2 = {}

    total_active_compounds = sum(y_true)

    total_compounds = len(y_true)

    # for ratio in ratio_list:

    #     for i, t in enumerate(fpr):

    #         if t > ratio:

    #             #print(fpr[i], tpr[i])

    #             if fpr[i-1]==0:

    #                 res[str(ratio)]=tpr[i]/fpr[i]

    #             else:

    #                 res[str(ratio)]=tpr[i-1]/fpr[i-1]

    #             break

    for ratio in ratio_list:

        res2[str(ratio)] = re_new(y_true, y_score, ratio)

    #print(res)

    #print(res2)

    return res2

def cal_metrics(y_true, y_score, alpha):

    """

    Calculate BEDROC score.

    Parameters:

    - y_true: true binary labels (0 or 1)

    - y_score: predicted scores or probabilities

    - alpha: parameter controlling the degree of early retrieval emphasis

    Returns:

    - BEDROC score

    """

        # concate res_single and labels

    scores = np.expand_dims(y_score, axis=1)

    y_true = np.expand_dims(y_true, axis=1)

    scores = np.concatenate((scores, y_true), axis=1)

    # inverse sort scores based on first column

    scores = scores[scores[:,0].argsort()[::-1]]

    bedroc = CalcBEDROC(scores, 1, 80.5)

    count = 0

    # sort y_score, return index

    index  = np.argsort(y_score)[::-1]

    for i in range(int(len(index)*0.005)):

        if y_true[index[i]] == 1:

            count += 1

    auc = CalcAUC(scores, 1)

    ef_list = CalcEnrichment(scores, 1, [0.005, 0.01, 0.02, 0.05])

    ef = {

        "0.005": ef_list[0],

        "0.01": ef_list[1],

        "0.02": ef_list[2],

        "0.05": ef_list[3]

    }

    re_list = calc_re(y_true, y_score, [0.005, 0.01, 0.02, 0.05])

    return auc, bedroc, ef, re_list

@register_task("drugclip")

class DrugCLIP(UnicoreTask):

    """Task for training transformer auto-encoder models."""

    @staticmethod

    def add_args(parser):

        """Add task-specific arguments to the parser."""

        parser.add_argument(

            "data",

            help="downstream data path",

        )

        parser.add_argument(

            "--finetune-mol-model",

            default=None,

            type=str,

            help="pretrained molecular model path",

        )

        parser.add_argument(

            "--finetune-pocket-model",

            default=None,

            type=str,

            help="pretrained pocket model path",

        )

        parser.add_argument(

            "--dist-threshold",

            type=float,

            default=6.0,

            help="threshold for the distance between the molecule and the pocket",

        )

        parser.add_argument(

            "--max-pocket-atoms",

            type=int,

            default=256,

            help="selected maximum number of atoms in a pocket",

        )

        parser.add_argument(

            "--test-model",

            default=False,

            type=Boolean,

            help="whether test model",

        )

        parser.add_argument("--reg", action="store_true", help="regression task")

    def __init__(self, args, dictionary, pocket_dictionary):

        super().__init__(args)

        self.dictionary = dictionary

        self.pocket_dictionary = pocket_dictionary

        self.seed = args.seed

        # add mask token

        self.mask_idx = dictionary.add_symbol("[MASK]", is_special=True)

        self.pocket_mask_idx = pocket_dictionary.add_symbol("[MASK]", is_special=True)

        self.mol_reps = None

        self.keys = None

    @classmethod

    def setup_task(cls, args, **kwargs):

        mol_dictionary = Dictionary.load(os.path.join(args.data, "dict_mol.txt"))

        pocket_dictionary = Dictionary.load(os.path.join(args.data, "dict_pkt.txt"))

        logger.info("ligand dictionary: {} types".format(len(mol_dictionary)))

        logger.info("pocket dictionary: {} types".format(len(pocket_dictionary)))

        return cls(args, mol_dictionary, pocket_dictionary)

    def load_dataset(self, split, **kwargs):

        """Load a given dataset split.

        'smi','pocket','atoms','coordinates','pocket_atoms','pocket_coordinates'

        Args:

            split (str): name of the data scoure (e.g., bppp)

        """

        data_path = os.path.join(self.args.data, split + ".lmdb")

        dataset = LMDBDataset(data_path)

        if split.startswith("train"):

            smi_dataset = KeyDataset(dataset, "smi")

            poc_dataset = KeyDataset(dataset, "pocket")

            dataset = AffinityDataset(

                dataset,

                self.args.seed,

                "atoms",

                "coordinates",

                "pocket_atoms",

                "pocket_coordinates",

                "label",

                True,

            )

            tgt_dataset = KeyDataset(dataset, "affinity")

        else:

            dataset = AffinityDataset(

                dataset,

                self.args.seed,

                "atoms",

                "coordinates",

                "pocket_atoms",

                "pocket_coordinates",

                "label",

            )

            tgt_dataset = KeyDataset(dataset, "affinity")

            smi_dataset = KeyDataset(dataset, "smi")

            poc_dataset = KeyDataset(dataset, "pocket")

        def PrependAndAppend(dataset, pre_token, app_token):

            dataset = PrependTokenDataset(dataset, pre_token)

            return AppendTokenDataset(dataset, app_token)

        dataset = RemoveHydrogenPocketDataset(

            dataset,

            "pocket_atoms",

            "pocket_coordinates",

            True,

            True,

        )

        dataset = CroppingPocketDataset(

            dataset,

            self.seed,

            "pocket_atoms",

            "pocket_coordinates",

            self.args.max_pocket_atoms,

        )

        dataset = RemoveHydrogenDataset(dataset, "atoms", "coordinates", True, True)

        apo_dataset = NormalizeDataset(dataset, "coordinates")

        apo_dataset = NormalizeDataset(apo_dataset, "pocket_coordinates")

        src_dataset = KeyDataset(apo_dataset, "atoms")

        mol_len_dataset = LengthDataset(src_dataset)

        src_dataset = TokenizeDataset(

            src_dataset, self.dictionary, max_seq_len=self.args.max_seq_len

        )

        coord_dataset = KeyDataset(apo_dataset, "coordinates")

        src_dataset = PrependAndAppend(

            src_dataset, self.dictionary.bos(), self.dictionary.eos()

        )

        edge_type = EdgeTypeDataset(src_dataset, len(self.dictionary))

        coord_dataset = FromNumpyDataset(coord_dataset)

        distance_dataset = DistanceDataset(coord_dataset)

        coord_dataset = PrependAndAppend(coord_dataset, 0.0, 0.0)

        distance_dataset = PrependAndAppend2DDataset(distance_dataset, 0.0)

        src_pocket_dataset = KeyDataset(apo_dataset, "pocket_atoms")

        pocket_len_dataset = LengthDataset(src_pocket_dataset)

        src_pocket_dataset = TokenizeDataset(

            src_pocket_dataset,

            self.pocket_dictionary,

            max_seq_len=self.args.max_seq_len,

        )

        coord_pocket_dataset = KeyDataset(apo_dataset, "pocket_coordinates")

        src_pocket_dataset = PrependAndAppend(

            src_pocket_dataset,

            self.pocket_dictionary.bos(),

            self.pocket_dictionary.eos(),

        )

        pocket_edge_type = EdgeTypeDataset(

            src_pocket_dataset, len(self.pocket_dictionary)

        )

        coord_pocket_dataset = FromNumpyDataset(coord_pocket_dataset)

        distance_pocket_dataset = DistanceDataset(coord_pocket_dataset)

        coord_pocket_dataset = PrependAndAppend(coord_pocket_dataset, 0.0, 0.0)

        distance_pocket_dataset = PrependAndAppend2DDataset(

            distance_pocket_dataset, 0.0

        )

        nest_dataset = NestedDictionaryDataset(

            {

                "net_input": {

                    "mol_src_tokens": RightPadDataset(

                        src_dataset,

                        pad_idx=self.dictionary.pad(),

                    ),

                    "mol_src_distance": RightPadDataset2D(

                        distance_dataset,

                        pad_idx=0,

                    ),

                    "mol_src_edge_type": RightPadDataset2D(

                        edge_type,

                        pad_idx=0,

                    ),

                    "pocket_src_tokens": RightPadDataset(

                        src_pocket_dataset,

                        pad_idx=self.pocket_dictionary.pad(),

                    ),

                    "pocket_src_distance": RightPadDataset2D(

                        distance_pocket_dataset,

                        pad_idx=0,

                    ),

                    "pocket_src_edge_type": RightPadDataset2D(

                        pocket_edge_type,

                        pad_idx=0,

                    ),

                    "pocket_src_coord": RightPadDatasetCoord(

                        coord_pocket_dataset,

                        pad_idx=0,

                    ),

                    "mol_len": RawArrayDataset(mol_len_dataset),

                    "pocket_len": RawArrayDataset(pocket_len_dataset)

                },

                "target": {

                    "finetune_target": RawLabelDataset(tgt_dataset),

                },

                "smi_name": RawArrayDataset(smi_dataset),

                "pocket_name": RawArrayDataset(poc_dataset),

            },

        )

        if split == "train":

            with data_utils.numpy_seed(self.args.seed):

                shuffle = np.random.permutation(len(src_dataset))

            self.datasets[split] = SortDataset(

                nest_dataset,

                sort_order=[shuffle],

            )

            self.datasets[split] = ResamplingDataset(

                self.datasets[split]

            )

        else:

            self.datasets[split] = nest_dataset

    def load_mols_dataset(self, data_path,atoms,coords, **kwargs):

        dataset = LMDBDataset(data_path)

        label_dataset = KeyDataset(dataset, "label")

        dataset = AffinityMolDataset(

            dataset,

            self.args.seed,

            atoms,

            coords,

            False,

        )

        smi_dataset = KeyDataset(dataset, "smi")

        def PrependAndAppend(dataset, pre_token, app_token):

            dataset = PrependTokenDataset(dataset, pre_token)

            return AppendTokenDataset(dataset, app_token)

        dataset = RemoveHydrogenDataset(dataset, "atoms", "coordinates", True, True)

        apo_dataset = NormalizeDataset(dataset, "coordinates")

        src_dataset = KeyDataset(apo_dataset, "atoms")

        len_dataset = LengthDataset(src_dataset)

        src_dataset = TokenizeDataset(

            src_dataset, self.dictionary, max_seq_len=self.args.max_seq_len

        )

        coord_dataset = KeyDataset(apo_dataset, "coordinates")

        src_dataset = PrependAndAppend(

            src_dataset, self.dictionary.bos(), self.dictionary.eos()

        )

        edge_type = EdgeTypeDataset(src_dataset, len(self.dictionary))

        coord_dataset = FromNumpyDataset(coord_dataset)

        distance_dataset = DistanceDataset(coord_dataset)

        coord_dataset = PrependAndAppend(coord_dataset, 0.0, 0.0)

        distance_dataset = PrependAndAppend2DDataset(distance_dataset, 0.0)

        nest_dataset = NestedDictionaryDataset(

            {

                "net_input": {

                    "mol_src_tokens": RightPadDataset(

                        src_dataset,

                        pad_idx=self.dictionary.pad(),

                    ),

                    "mol_src_distance": RightPadDataset2D(

                        distance_dataset,

                        pad_idx=0,

                    ),

                    "mol_src_edge_type": RightPadDataset2D(

                        edge_type,

                        pad_idx=0,

                    ),

                },

                "smi_name": RawArrayDataset(smi_dataset),

                "target":  RawArrayDataset(label_dataset),

                "mol_len": RawArrayDataset(len_dataset),

            },

        )

        return nest_dataset

    def load_retrieval_mols_dataset(self, data_path,atoms,coords, **kwargs):

        dataset = LMDBDataset(data_path)

        dataset = AffinityMolDataset(

            dataset,

            self.args.seed,

            atoms,

            coords,

            False,

        )

        smi_dataset = KeyDataset(dataset, "smi")

        def PrependAndAppend(dataset, pre_token, app_token):

            dataset = PrependTokenDataset(dataset, pre_token)

            return AppendTokenDataset(dataset, app_token)

        dataset = RemoveHydrogenDataset(dataset, "atoms", "coordinates", True, True)

        apo_dataset = NormalizeDataset(dataset, "coordinates")

        src_dataset = KeyDataset(apo_dataset, "atoms")

        len_dataset = LengthDataset(src_dataset)

        src_dataset = TokenizeDataset(

            src_dataset, self.dictionary, max_seq_len=self.args.max_seq_len

        )

        coord_dataset = KeyDataset(apo_dataset, "coordinates")

        src_dataset = PrependAndAppend(

            src_dataset, self.dictionary.bos(), self.dictionary.eos()

        )

        edge_type = EdgeTypeDataset(src_dataset, len(self.dictionary))

        coord_dataset = FromNumpyDataset(coord_dataset)

        distance_dataset = DistanceDataset(coord_dataset)

        coord_dataset = PrependAndAppend(coord_dataset, 0.0, 0.0)

        distance_dataset = PrependAndAppend2DDataset(distance_dataset, 0.0)

        nest_dataset = NestedDictionaryDataset(

            {

                "net_input": {

                    "mol_src_tokens": RightPadDataset(

                        src_dataset,

                        pad_idx=self.dictionary.pad(),

                    ),

                    "mol_src_distance": RightPadDataset2D(

                        distance_dataset,

                        pad_idx=0,

                    ),

                    "mol_src_edge_type": RightPadDataset2D(

                        edge_type,

                        pad_idx=0,

                    ),

                },

                "smi_name": RawArrayDataset(smi_dataset),

                "mol_len": RawArrayDataset(len_dataset),

            },

        )

        return nest_dataset

    def load_pockets_dataset(self, data_path, **kwargs):

        dataset = LMDBDataset(data_path)

        dataset = AffinityPocketDataset(

            dataset,

            self.args.seed,

            "pocket_atoms",

            "pocket_coordinates",

            False,

            "pocket"

        )

        poc_dataset = KeyDataset(dataset, "pocket")

        def PrependAndAppend(dataset, pre_token, app_token):

            dataset = PrependTokenDataset(dataset, pre_token)

            return AppendTokenDataset(dataset, app_token)

        dataset = RemoveHydrogenPocketDataset(

            dataset,

            "pocket_atoms",

            "pocket_coordinates",

            True,

            True,

        )

        dataset = CroppingPocketDataset(

            dataset,

            self.seed,

            "pocket_atoms",

            "pocket_coordinates",

            self.args.max_pocket_atoms,

        )

        apo_dataset = NormalizeDataset(dataset, "pocket_coordinates")

        src_pocket_dataset = KeyDataset(apo_dataset, "pocket_atoms")

        len_dataset = LengthDataset(src_pocket_dataset)

        src_pocket_dataset = TokenizeDataset(

            src_pocket_dataset,

            self.pocket_dictionary,

            max_seq_len=self.args.max_seq_len,

        )

        coord_pocket_dataset = KeyDataset(apo_dataset, "pocket_coordinates")

        src_pocket_dataset = PrependAndAppend(

            src_pocket_dataset,

            self.pocket_dictionary.bos(),

            self.pocket_dictionary.eos(),

        )

        pocket_edge_type = EdgeTypeDataset(

            src_pocket_dataset, len(self.pocket_dictionary)

        )

        coord_pocket_dataset = FromNumpyDataset(coord_pocket_dataset)

        distance_pocket_dataset = DistanceDataset(coord_pocket_dataset)

        coord_pocket_dataset = PrependAndAppend(coord_pocket_dataset, 0.0, 0.0)

        distance_pocket_dataset = PrependAndAppend2DDataset(

            distance_pocket_dataset, 0.0

        )

        nest_dataset = NestedDictionaryDataset(

            {

                "net_input": {

                    "pocket_src_tokens": RightPadDataset(

                        src_pocket_dataset,

                        pad_idx=self.pocket_dictionary.pad(),

                    ),

                    "pocket_src_distance": RightPadDataset2D(

                        distance_pocket_dataset,

                        pad_idx=0,

                    ),

                    "pocket_src_edge_type": RightPadDataset2D(

                        pocket_edge_type,

                        pad_idx=0,

                    ),

                    "pocket_src_coord": RightPadDatasetCoord(

                        coord_pocket_dataset,

                        pad_idx=0,

                    ),

                },

                "pocket_name": RawArrayDataset(poc_dataset),

                "pocket_len": RawArrayDataset(len_dataset),

            },

        )

        return nest_dataset

    def build_model(self, args):

        from unicore import models

        model = models.build_model(args, self)

        if args.finetune_mol_model is not None:

            print("load pretrain model weight from...", args.finetune_mol_model)

            state = checkpoint_utils.load_checkpoint_to_cpu(

                args.finetune_mol_model,

            )

            model.mol_model.load_state_dict(state["model"], strict=False)

        if args.finetune_pocket_model is not None:

            print("load pretrain model weight from...", args.finetune_pocket_model)

            state = checkpoint_utils.load_checkpoint_to_cpu(

                args.finetune_pocket_model,

            )

            model.pocket_model.load_state_dict(state["model"], strict=False)

        return model

    def train_step(

        self, sample, model, loss, optimizer, update_num, ignore_grad=False

    ):

        """

        Do forward and backward, and return the loss as computed by *loss*

        for the given *model* and *sample*.

        Args:

            sample (dict): the mini-batch. The format is defined by the

                :class:`~unicore.data.UnicoreDataset`.

            model (~unicore.models.BaseUnicoreModel): the model

            loss (~unicore.losses.UnicoreLoss): the loss

            optimizer (~unicore.optim.UnicoreOptimizer): the optimizer

            update_num (int): the current update

            ignore_grad (bool): multiply loss by 0 if this is set to True

        Returns:

            tuple:

                - the loss

                - the sample size, which is used as the denominator for the

                  gradient

                - logging outputs to display while training

        """

        model.train()

        model.set_num_updates(update_num)

        with torch.autograd.profiler.record_function("forward"):

            loss, sample_size, logging_output = loss(model, sample)

        if ignore_grad:

            loss *= 0

        with torch.autograd.profiler.record_function("backward"):

            optimizer.backward(loss)

        return loss, sample_size, logging_output

    def valid_step(self, sample, model, loss, test=False):

        model.eval()

        with torch.no_grad():

            loss, sample_size, logging_output = loss(model, sample)

        return loss, sample_size, logging_output

    def test_pcba_target(self, name, model, **kwargs):

        """Encode a dataset with the molecule encoder."""

        #names = "PPARG"

        data_path = "./data/lit_pcba/" + name + "/mols.lmdb"

        mol_dataset = self.load_mols_dataset(data_path, "atoms", "coordinates")

        num_data = len(mol_dataset)

        bsz=64

        #print(num_data//bsz)

        mol_reps = []

        mol_names = []

        labels = []

        # generate mol data

        mol_data = torch.utils.data.DataLoader(mol_dataset, batch_size=bsz, collate_fn=mol_dataset.collater)

        for _, sample in enumerate(tqdm(mol_data)):

            sample = unicore.utils.move_to_cuda(sample)

            dist = sample["net_input"]["mol_src_distance"]

            et = sample["net_input"]["mol_src_edge_type"]

            st = sample["net_input"]["mol_src_tokens"]

            mol_padding_mask = st.eq(model.mol_model.padding_idx)

            mol_x = model.mol_model.embed_tokens(st)

            n_node = dist.size(-1)

            gbf_feature = model.mol_model.gbf(dist, et)

            gbf_result = model.mol_model.gbf_proj(gbf_feature)

            graph_attn_bias = gbf_result

            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()

            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)

            mol_outputs = model.mol_model.encoder(

                mol_x, padding_mask=mol_padding_mask, attn_mask=graph_attn_bias

            )

            mol_encoder_rep = mol_outputs[0][:,0,:]

            mol_emb = model.mol_project(mol_encoder_rep)

            mol_emb = mol_emb / mol_emb.norm(dim=1, keepdim=True)

            mol_emb = mol_emb.detach().cpu().numpy()

            mol_reps.append(mol_emb)

            mol_names.extend(sample["smi_name"])

            labels.extend(sample["target"].detach().cpu().numpy())

        mol_reps = np.concatenate(mol_reps, axis=0)

        labels = np.array(labels, dtype=np.int32)

        # generate pocket data

        data_path = "./data/lit_pcba/" + name + "/pockets.lmdb"

        pocket_dataset = self.load_pockets_dataset(data_path)

        pocket_data = torch.utils.data.DataLoader(pocket_dataset, batch_size=bsz, collate_fn=pocket_dataset.collater)

        pocket_reps = []

        for _, sample in enumerate(tqdm(pocket_data)):

            sample = unicore.utils.move_to_cuda(sample)

            dist = sample["net_input"]["pocket_src_distance"]

            et = sample["net_input"]["pocket_src_edge_type"]

            st = sample["net_input"]["pocket_src_tokens"]

            pocket_padding_mask = st.eq(model.pocket_model.padding_idx)

            pocket_x = model.pocket_model.embed_tokens(st)

            n_node = dist.size(-1)

            gbf_feature = model.pocket_model.gbf(dist, et)

            gbf_result = model.pocket_model.gbf_proj(gbf_feature)

            graph_attn_bias = gbf_result

            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()

            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)

            pocket_outputs = model.pocket_model.encoder(

                pocket_x, padding_mask=pocket_padding_mask, attn_mask=graph_attn_bias

            )

            pocket_encoder_rep = pocket_outputs[0][:,0,:]

            pocket_emb = model.pocket_project(pocket_encoder_rep)

            pocket_emb = pocket_emb / pocket_emb.norm(dim=1, keepdim=True)

            pocket_emb = pocket_emb.detach().cpu().numpy()

            pocket_names = sample["pocket_name"]

            pocket_reps.append(pocket_emb)

        pocket_reps = np.concatenate(pocket_reps, axis=0)

        res = pocket_reps @ mol_reps.T

        res_single = res.max(axis=0)

        auc, bedroc, ef_list, re_list = cal_metrics(labels, res_single, 80.5)

        return auc, bedroc, ef_list, re_list

    def test_pcba(self, model, **kwargs):

        #ckpt_date = self.args.finetune_from_model.split("/")[-2]

        #save_name = "/home/gaobowen/DrugClip/test_results/pcba/" + ckpt_date + ".txt"

        save_name = ""

        targets = os.listdir("./data/lit_pcba/")

        #print(targets)

        auc_list = []

        ef_list = []

        bedroc_list = []

        re_list = {

            "0.005": [],

            "0.01": [],

            "0.02": [],

            "0.05": []

        }

        ef_list = {

            "0.005": [],

            "0.01": [],

            "0.02": [],

            "0.05": []

        }

        for target in targets:

            auc, bedroc, ef, re = self.test_pcba_target(target, model)

            auc_list.append(auc)

            bedroc_list.append(bedroc)

            for key in ef:

                ef_list[key].append(ef[key])

            # print("re", re)

            # print("ef", ef)

            for key in re:

                re_list[key].append(re[key])

        print(auc_list)

        print(ef_list)

        print("auc 25%", np.percentile(auc_list, 25))

        print("auc 50%", np.percentile(auc_list, 50))

        print("auc 75%", np.percentile(auc_list, 75))

        print("auc mean", np.mean(auc_list))

        print("bedroc 25%", np.percentile(bedroc_list, 25))

        print("bedroc 50%", np.percentile(bedroc_list, 50))

        print("bedroc 75%", np.percentile(bedroc_list, 75))

        print("bedroc mean", np.mean(bedroc_list))

        #print(np.median(auc_list))

        #print(np.median(ef_list))

        for key in ef_list:

            print("ef", key, "25%", np.percentile(ef_list[key], 25))

            print("ef",key, "50%", np.percentile(ef_list[key], 50))

            print("ef",key, "75%", np.percentile(ef_list[key], 75))

            print("ef",key, "mean", np.mean(ef_list[key]))

        for key in re_list:

            print("re",key, "25%", np.percentile(re_list[key], 25))

            print("re",key, "50%", np.percentile(re_list[key], 50))

            print("re",key, "75%", np.percentile(re_list[key], 75))

            print("re",key, "mean", np.mean(re_list[key]))

        return 

    def test_dude_target(self, target, model, **kwargs):

        data_path = "./data/DUD-E/raw/all/" + target + "/mols.lmdb"

        mol_dataset = self.load_mols_dataset(data_path, "atoms", "coordinates")

        num_data = len(mol_dataset)

        bsz=64

        print(num_data//bsz)

        mol_reps = []

        mol_names = []

        labels = []

        # generate mol data

        mol_data = torch.utils.data.DataLoader(mol_dataset, batch_size=bsz, collate_fn=mol_dataset.collater)

        for _, sample in enumerate(tqdm(mol_data)):

            sample = unicore.utils.move_to_cuda(sample)

            dist = sample["net_input"]["mol_src_distance"]

            et = sample["net_input"]["mol_src_edge_type"]

            st = sample["net_input"]["mol_src_tokens"]

            mol_padding_mask = st.eq(model.mol_model.padding_idx)

            mol_x = model.mol_model.embed_tokens(st)

            n_node = dist.size(-1)

            gbf_feature = model.mol_model.gbf(dist, et)

            gbf_result = model.mol_model.gbf_proj(gbf_feature)

            graph_attn_bias = gbf_result

            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()

            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)

            mol_outputs = model.mol_model.encoder(

                mol_x, padding_mask=mol_padding_mask, attn_mask=graph_attn_bias

            )

            mol_encoder_rep = mol_outputs[0][:,0,:]

            mol_emb = mol_encoder_rep

            mol_emb = model.mol_project(mol_encoder_rep)

            mol_emb = mol_emb / mol_emb.norm(dim=-1, keepdim=True)

            #print(mol_emb.dtype)

            mol_emb = mol_emb.detach().cpu().numpy()

            #print(mol_emb.dtype)

            mol_reps.append(mol_emb)

            mol_names.extend(sample["smi_name"])

            labels.extend(sample["target"].detach().cpu().numpy())

        mol_reps = np.concatenate(mol_reps, axis=0)

        labels = np.array(labels, dtype=np.int32)

        # generate pocket data

        data_path = "./data/DUD-E/raw/all/" + target + "/pocket.lmdb"

        pocket_dataset = self.load_pockets_dataset(data_path)

        pocket_data = torch.utils.data.DataLoader(pocket_dataset, batch_size=bsz, collate_fn=pocket_dataset.collater)

        pocket_reps = []

        for _, sample in enumerate(tqdm(pocket_data)):

            sample = unicore.utils.move_to_cuda(sample)

            dist = sample["net_input"]["pocket_src_distance"]

            et = sample["net_input"]["pocket_src_edge_type"]

            st = sample["net_input"]["pocket_src_tokens"]

            pocket_padding_mask = st.eq(model.pocket_model.padding_idx)

            pocket_x = model.pocket_model.embed_tokens(st)

            n_node = dist.size(-1)

            gbf_feature = model.pocket_model.gbf(dist, et)

            gbf_result = model.pocket_model.gbf_proj(gbf_feature)

            graph_attn_bias = gbf_result

            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()

            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)

            pocket_outputs = model.pocket_model.encoder(

                pocket_x, padding_mask=pocket_padding_mask, attn_mask=graph_attn_bias

            )

            pocket_encoder_rep = pocket_outputs[0][:,0,:]

            #pocket_emb = pocket_encoder_rep

            pocket_emb = model.pocket_project(pocket_encoder_rep)

            pocket_emb = pocket_emb / pocket_emb.norm(dim=-1, keepdim=True)

            pocket_emb = pocket_emb.detach().cpu().numpy()

            pocket_reps.append(pocket_emb)

        pocket_reps = np.concatenate(pocket_reps, axis=0)

        print(pocket_reps.shape)

        res = pocket_reps @ mol_reps.T

        res_single = res.max(axis=0)

        auc, bedroc, ef_list, re_list = cal_metrics(labels, res_single, 80.5)

        print(target)

        print(np.sum(labels), len(labels)-np.sum(labels))

        return auc, bedroc, ef_list, re_list, res_single, labels

    def test_dude(self, model, **kwargs):

        targets = os.listdir("./data/DUD-E/raw/all/")

        auc_list = []

        bedroc_list = []

        ef_list = []

        res_list= []

        labels_list = []

        re_list = {

            "0.005": [],

            "0.01": [],

            "0.02": [],

            "0.05": [],

        }

        ef_list = {

            "0.005": [],

            "0.01": [],

            "0.02": [],

            "0.05": [],

        }

        for i,target in enumerate(targets):

            auc, bedroc, ef, re, res_single, labels = self.test_dude_target(target, model)

            auc_list.append(auc)

            bedroc_list.append(bedroc)

            for key in ef:

                ef_list[key].append(ef[key])

            for key in re_list:

                re_list[key].append(re[key])

            res_list.append(res_single)

            labels_list.append(labels)

        res = np.concatenate(res_list, axis=0)

        labels = np.concatenate(labels_list, axis=0)

        print("auc mean", np.mean(auc_list))

        print("bedroc mean", np.mean(bedroc_list))

        for key in ef_list:

            print("ef", key, "mean", np.mean(ef_list[key]))

        for key in re_list:

            print("re", key, "mean",  np.mean(re_list[key]))

        # save printed results 

        return

    def encode_mols_once(self, model, data_path, emb_dir, atoms, coords, **kwargs):

        # cache path is embdir/data_path.pkl

        cache_path = os.path.join(emb_dir, data_path.split("/")[-1] + ".pkl")

        if os.path.exists(cache_path):

            with open(cache_path, "rb") as f:

                mol_reps, mol_names = pickle.load(f)

            return mol_reps, mol_names

        mol_dataset = self.load_retrieval_mols_dataset(data_path,atoms,coords)

        mol_reps = []

        mol_names = []

        bsz=32

        mol_data = torch.utils.data.DataLoader(mol_dataset, batch_size=bsz, collate_fn=mol_dataset.collater)

        for _, sample in enumerate(tqdm(mol_data)):

            sample = unicore.utils.move_to_cuda(sample)

            dist = sample["net_input"]["mol_src_distance"]

            et = sample["net_input"]["mol_src_edge_type"]

            st = sample["net_input"]["mol_src_tokens"]

            mol_padding_mask = st.eq(model.mol_model.padding_idx)

            mol_x = model.mol_model.embed_tokens(st)

            n_node = dist.size(-1)

            gbf_feature = model.mol_model.gbf(dist, et)

            gbf_result = model.mol_model.gbf_proj(gbf_feature)

            graph_attn_bias = gbf_result

            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()

            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)

            mol_outputs = model.mol_model.encoder(

                mol_x, padding_mask=mol_padding_mask, attn_mask=graph_attn_bias

            )

            mol_encoder_rep = mol_outputs[0][:,0,:]

            mol_emb = model.mol_project(mol_encoder_rep)

            mol_emb = mol_emb / mol_emb.norm(dim=-1, keepdim=True)

            mol_emb = mol_emb.detach().cpu().numpy()

            mol_reps.append(mol_emb)

            mol_names.extend(sample["smi_name"])

        mol_reps = np.concatenate(mol_reps, axis=0)

        # save the results

        with open(cache_path, "wb") as f:

            pickle.dump([mol_reps, mol_names], f)

        return mol_reps, mol_names

    def retrieve_mols(self, model, mol_path, pocket_path, emb_dir, k, **kwargs):

        os.makedirs(emb_dir, exist_ok=True)        

        mol_reps, mol_names = self.encode_mols_once(model, mol_path, emb_dir,  "atoms", "coordinates")

        pocket_dataset = self.load_pockets_dataset(pocket_path)

        pocket_data = torch.utils.data.DataLoader(pocket_dataset, batch_size=16, collate_fn=pocket_dataset.collater)

        pocket_reps = []

        pocket_names = []

        for _, sample in enumerate(tqdm(pocket_data)):

            sample = unicore.utils.move_to_cuda(sample)

            dist = sample["net_input"]["pocket_src_distance"]

            et = sample["net_input"]["pocket_src_edge_type"]

            st = sample["net_input"]["pocket_src_tokens"]

            pocket_padding_mask = st.eq(model.pocket_model.padding_idx)

            pocket_x = model.pocket_model.embed_tokens(st)

            n_node = dist.size(-1)

            gbf_feature = model.pocket_model.gbf(dist, et)

            gbf_result = model.pocket_model.gbf_proj(gbf_feature)

            graph_attn_bias = gbf_result

            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()

            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)

            pocket_outputs = model.pocket_model.encoder(

                pocket_x, padding_mask=pocket_padding_mask, attn_mask=graph_attn_bias

            )

            pocket_encoder_rep = pocket_outputs[0][:,0,:]

            pocket_emb = model.pocket_project(pocket_encoder_rep)

            pocket_emb = pocket_emb / pocket_emb.norm(dim=-1, keepdim=True)

            pocket_emb = pocket_emb.detach().cpu().numpy()

            pocket_reps.append(pocket_emb)

            pocket_names.extend(sample["pocket_name"])

        pocket_reps = np.concatenate(pocket_reps, axis=0)

        res = pocket_reps @ mol_reps.T

        res = res.max(axis=0)

        # get top k results

        top_k = np.argsort(res)[::-1][:k]

        # return names and scores

        return [mol_names[i] for i in top_k], res[top_k]