from transformers.models.bert.modeling_bert import (BertEmbeddings, BertModel,
BertForSequenceClassification,BertPreTrainedModel,BertEncoder,BertPooler)
from transformers.models.roberta.modeling_roberta import RobertaModel,RobertaEncoder
import torch.nn as nn
import torch
import torch.nn.functional as F
from fastNLP.core.metrics import MetricBase,seq_len_to_mask
import GNN_global_representation
import GNN_edge_regression
from utils import *
import ipdb
#增加：糖基化结构用dgl的图结构表示（node_index,masses.glyco_process)--写GNN作为糖的全局表示
Ablation="GIN"
loc=True
GIN_hidden_dim=64
class Acid_BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""
    def __init__(self, config):
        super().__init__()
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.charge_embedding=nn.Embedding(10,config.hidden_size,padding_idx=0)
        self.a_embedding=nn.Embedding(30,config.hidden_size,padding_idx=0)
        self.phos_embedding=nn.Embedding(10,config.hidden_size)#修饰三种加上padding###全部调大了，其他的修饰也在这里
        if Ablation=="GIN":
            self.gly_embedding=GNN_global_representation.GIN(20, 16, config.hidden_size,init_eps=0)
        if Ablation=="Nogly":
            pass
        if Ablation=="GCN":
            #self.gly_embedding=GNN_global_representation.GCN(20, 16, config.hidden_size)
            print("input GCN model")
        if Ablation not in ["GIN","GCN","Nogly"]:
            raise NameError
        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))


    def forward(
        self, peptide_tokens=None, position_ids=None,decoration=None,charge=None,batched_graph=None, inputs_embeds=None, past_key_values_length=0
    ):  
        N=peptide_tokens.size(0)
        L=peptide_tokens.size(1)
        sequence=peptide_tokens
        charge_embed=self.charge_embedding(charge.unsqueeze(1).expand(N,L))
        feat=batched_graph.ndata["attr"]
        if Ablation=="GIN" or Ablation =="GCN":
            gly_embedding=self.gly_embedding(batched_graph,feat)
            gly_embedding=gly_embedding.unsqueeze(1).expand(N,L,gly_embedding.size(1))*((decoration==5).unsqueeze(-1).expand(
    N,L,gly_embedding.size(1)))
        if Ablation=="Nogly":
            pass

        assert sequence.size(0) == decoration.size(0)
        # ipdb.set_trace()
        #decoration传入了decoration_ids，包括了decoration_ACE的信息，见preprocess.py
        phos_embed = self.phos_embedding(decoration)
        
        if loc==False:
            phos_embed = self.phos_embedding(decoration-5*(decoration==5).to(int))

        if peptide_tokens is not None:
            input_shape = peptide_tokens.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]

        if position_ids is None:
            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

        inputs_embeds = self.a_embedding(peptide_tokens)
        if Ablation !="Nogly":
            embeddings = inputs_embeds+phos_embed+charge_embed+gly_embedding
        if Ablation =="Nogly":
            embeddings = inputs_embeds+phos_embed+charge_embed

        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

# ----------------------- BY ------------------------------#
class _2deepchargeModelBY_bert(BertModel):#input:sequence:N*L(N:batch)##########bertmodel
    #####输入是cls A sep B sep C ....所以长度为2L，最后取所有的sep出来预测因为pretrain的原因只能写死在这里面numcol
    def __init__(self,config):
        super().__init__(config)
        self.config = config
        self.embeddings = Acid_BertEmbeddings(config)
        self.encoder = BertEncoder(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.activation=nn.GELU()
        self.num_col=num_col
        self.pooler = BertPooler(config)
        self.init_weights()
        print(f"GNN_ablation {GNN_edge_ablation}!!!")
        if GNN_edge_ablation=="GIN":
            self.BY_pred=GNN_edge_regression.GIN(20, GIN_hidden_dim, config.hidden_size,init_eps=0)
        if GNN_edge_ablation=="GCN":
            self.BY_pred=GNN_edge_regression.GCN(20, GIN_hidden_dim, config.hidden_size)
        if GNN_edge_ablation=="GAT":
            self.BY_pred=GNN_edge_regression.GAT(20, GIN_hidden_dim, config.hidden_size,num_heads=4)
        self.peptide_rep_linear=nn.Linear(config.hidden_size,GIN_hidden_dim)

    def forward(self,input_ids,peptide_tokens,peptide_length,charge,decoration,decoration_ids,_id,decoration_ACE,PlausibleStruct,peptide,return_dict=None,head_mask=None):
        #input:input_ids:N*2L(N:batch)
        # length:N*1(N:batch)
        # phos:N*2L(N:batch)
        #charge:N*1
        
        # ninput=self.dropout(ninput)
        key_padding_mask=seq_len_to_mask(peptide_length*2)#对肽段两倍长度以外的部分出现了mask
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        input_shape = input_ids.size()
        batch_size, seq_length = input_shape #现在的batchsize为128
        device = input_ids.device 
        # past_key_values_length = 0
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(key_padding_mask, input_shape, device)
        batched_graph=PlausibleStruct
        # import sys
        # sys.path.append('..')
        from masses import glyco_process
        import dgl
        graph_list=[]
        for i in batched_graph:
            _,_,Struct_tokens,nodef=glyco_process(i)
            node2idx={"P":0,"H" : 2, "N" : 1, "A" : 4, "G" : 5 , "F" :3}
            nodef="P"+nodef
            Struct_feature=[node2idx[i] for i in nodef]
            grapbatched_graphh=dgl.graph(Struct_tokens)
            grapbatched_graphh.ndata["attr"]=torch.Tensor(Struct_feature).to(int)
            graph_list.append(grapbatched_graphh)
        batched_graph = dgl.batch(graph_list).to(device)
        embedding_output = self.embeddings(peptide_tokens=input_ids, batched_graph=batched_graph,
        position_ids=None,
        decoration=decoration_ids,charge=charge, inputs_embeds=None, past_key_values_length=0
        ) #导出的结果是batchsize*peptide_length*embedding_size
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            use_cache=False,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]#####N*2L*E
        E=sequence_output.size(-1)
        output=sequence_output[:,2:-1:2,:] #取[SEP]表示
        assert output.size(1)==int(seq_length/2-1)
        batched_graph = dgl.batch(graph_list).to(device)
        feat = batched_graph.ndata["attr"]
        peptide_ind=torch.nonzero(feat==0).squeeze()
        peptide_rep=self.peptide_rep_linear(sequence_output[:,0,:])
        BY_pred=self.BY_pred(batched_graph,feat,peptide_rep,peptide_ind) #利用global_edge_regression对于图以及节点坐标得到边的碎裂可能性
        return {'pred':BY_pred,'sequence':peptide_tokens,'charge':charge,
        "decoration":decoration,"seq_len":peptide_length,"_id":_id,
        "PlausibleStruct":PlausibleStruct,"peptide":peptide}