# coding=utf-8
import os
import pdb
import copy
import torch
import torch.nn.functional as F
from torch import nn

from torch.nn import CrossEntropyLoss
from transformers import (
        BertConfig,
        BertModel,
        RobertaModel,
        BertForTokenClassification,
        BertTokenizer,
        RobertaConfig,
        RobertaForTokenClassification,
        RobertaTokenizer, 
        AutoTokenizer,
)

class BERTMultiNER2(BertForTokenClassification):
    def __init__(self, config, num_labels=3):
        super(BERTMultiNER2, self).__init__(config)
        self.num_labels = num_labels
        self.bert = BertModel(config)
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        
        self.dise_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # disease
        self.chem_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # chemical
        self.gene_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # gene/protein
        self.spec_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # species
        self.cellline_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # cell line
        self.dna_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # dna
        self.rna_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # rna
        self.celltype_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # cell type
        
        # self.biological_structure_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # biological structure
        # self.diagnostic_procedure_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # diagnostic procedure
        # self.duration_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # duration
        # self.date_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # date
        # self.therapeutic_procedure_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # therapeutic procedure
        # self.sign_symptom_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # sign/symptom
        # self.lab_value_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # lab value
        

        self.dise_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.chem_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.gene_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.spec_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.cellline_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.dna_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.rna_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.celltype_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        
        # self.biological_structure_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.diagnostic_procedure_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)    
        # self.duration_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)    
        # self.date_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.therapeutic_procedure_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.sign_symptom_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.lab_value_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)

        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, entity_type_ids=None):
        sequence_output = self.bert(input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, head_mask=None)[0]
        batch_size,max_len,feat_dim = sequence_output.shape
        sequence_output = self.dropout(sequence_output)

        if entity_type_ids[0][0].item() == 0:
            '''
            Raw text data with trained parameters
            '''
            dise_sequence_output = F.relu(self.dise_classifier_2(sequence_output)) # disease logit value
            chem_sequence_output = F.relu(self.chem_classifier_2(sequence_output)) # chemical logit value
            gene_sequence_output = F.relu(self.gene_classifier_2(sequence_output)) # gene/protein logit value
            spec_sequence_output = F.relu(self.spec_classifier_2(sequence_output)) # species logit value
            cellline_sequence_output = F.relu(self.cellline_classifier_2(sequence_output)) # cell line logit value
            dna_sequence_output = F.relu(self.dna_classifier_2(sequence_output)) # dna logit value
            rna_sequence_output = F.relu(self.rna_classifier_2(sequence_output)) # rna logit value
            celltype_sequence_output = F.relu(self.celltype_classifier_2(sequence_output)) # cell type logit value
            
            # biological_structure_sequence_output = F.relu(self.biological_structure_classifier_2(sequence_output)) # biological structure logit value
            # diagnostic_procedure_sequence_output = F.relu(self.diagnostic_procedure_classifier_2(sequence_output)) # diagnostic procedure logit value
            # duration_sequence_output = F.relu(self.duration_classifier_2(sequence_output)) # duration logit value
            # date_sequence_output = F.relu(self.date_classifier_2(sequence_output)) # date logit value
            # therapeutic_procedure_sequence_output = F.relu(self.therapeutic_procedure_classifier_2(sequence_output)) # therapeutic procedure logit value
            # sign_symptom_sequence_output = F.relu(self.sign_symptom_classifier_2(sequence_output)) # sign/symptom logit value
            # lab_value_sequence_output = F.relu(self.lab_value_classifier_2(sequence_output)) # lab value logit value


            dise_logits = self.dise_classifier(dise_sequence_output) # disease logit value
            chem_logits = self.chem_classifier(chem_sequence_output) # chemical logit value
            gene_logits = self.gene_classifier(gene_sequence_output) # gene/protein logit value
            spec_logits = self.spec_classifier(spec_sequence_output) # species logit value
            cellline_logits = self.cellline_classifier(cellline_sequence_output) # cell line logit value
            dna_logits = self.dna_classifier(dna_sequence_output) # dna logit value
            rna_logits = self.rna_classifier(rna_sequence_output) # rna logit value
            celltype_logits = self.celltype_classifier(celltype_sequence_output) # cell type logit value
            
            # biological_logits = self.biological_structure_classifier(biological_structure_sequence_output) # biological structure logit value
            # diagnostic_logits = self.diagnostic_procedure_classifier(diagnostic_procedure_sequence_output) # diagnostic procedure logit value
            # duration_logits = self.duration_classifier(duration_sequence_output) # duration logit value
            # date_logits = self.date_classifier(date_sequence_output) # date logit value
            # therapeutic_logits = self.therapeutic_procedure_classifier(therapeutic_procedure_sequence_output) # therapeutic procedure logit value
            # sign_symptom_logits = self.sign_symptom_classifier(sign_symptom_sequence_output) # sign/symptom logit value
            # lab_value_logits = self.lab_value_classifier(lab_value_sequence_output) # lab value logit value

            

            # update logit and sequence_output
            sequence_output = dise_sequence_output + chem_sequence_output + gene_sequence_output + spec_sequence_output + cellline_sequence_output + dna_sequence_output + rna_sequence_output + celltype_sequence_output 
            # + \
            #     biological_structure_sequence_output + diagnostic_procedure_sequence_output + duration_sequence_output + date_sequence_output + \
            #     therapeutic_procedure_sequence_output + sign_symptom_sequence_output + lab_value_sequence_output 
                
            logits = (dise_logits, chem_logits, gene_logits, spec_logits, cellline_logits, 
                      dna_logits, rna_logits, celltype_logits)
                    #   biological_logits, diagnostic_logits,
                    #   duration_logits, date_logits, therapeutic_logits,
                    #   sign_symptom_logits, lab_value_logits)
        else:
            ''' 
            Train, Eval, Test with pre-defined entity type tags
            '''
            # make 1*1 conv to adopt entity type
            dise_idx = copy.deepcopy(entity_type_ids)
            chem_idx = copy.deepcopy(entity_type_ids)
            gene_idx = copy.deepcopy(entity_type_ids)
            spec_idx = copy.deepcopy(entity_type_ids)
            cellline_idx = copy.deepcopy(entity_type_ids)
            dna_idx = copy.deepcopy(entity_type_ids)
            rna_idx = copy.deepcopy(entity_type_ids)
            celltype_idx = copy.deepcopy(entity_type_ids)
            
            # biological_idx = copy.deepcopy(entity_type_ids)
            # diagnostic_idx = copy.deepcopy(entity_type_ids)
            # duration_idx = copy.deepcopy(entity_type_ids)
            # date_idx = copy.deepcopy(entity_type_ids)
            # therapeutic_idx = copy.deepcopy(entity_type_ids)
            # sign_symptom_idx = copy.deepcopy(entity_type_ids)
            # lab_value_idx = copy.deepcopy(entity_type_ids)

            

            dise_idx[dise_idx != 1] = 0
            chem_idx[chem_idx != 2] = 0
            gene_idx[gene_idx != 3] = 0
            spec_idx[spec_idx != 4] = 0
            cellline_idx[cellline_idx != 5] = 0
            dna_idx[dna_idx != 6] = 0
            rna_idx[rna_idx != 7] = 0
            celltype_idx[celltype_idx != 8] = 0
            # biological_idx[biological_idx != 9] = 0
            # diagnostic_idx[diagnostic_idx != 10] = 0
            # duration_idx[duration_idx != 11] = 0
            # date_idx[date_idx != 12] = 0
            # therapeutic_idx[therapeutic_idx != 13] = 0
            # sign_symptom_idx[sign_symptom_idx != 14] = 0
            # lab_value_idx[lab_value_idx != 15] = 0

            dise_sequence_output = dise_idx.unsqueeze(-1) * sequence_output        
            chem_sequence_output = chem_idx.unsqueeze(-1) * sequence_output
            gene_sequence_output = gene_idx.unsqueeze(-1) * sequence_output
            spec_sequence_output = spec_idx.unsqueeze(-1) * sequence_output
            cellline_sequence_output = cellline_idx.unsqueeze(-1) * sequence_output
            dna_sequence_output = dna_idx.unsqueeze(-1) * sequence_output
            rna_sequence_output = rna_idx.unsqueeze(-1) * sequence_output
            celltype_sequence_output = celltype_idx.unsqueeze(-1) * sequence_output
            # biological_structure_sequence_output = biological_idx.unsqueeze(-1) * sequence_output
            # diagnostic_procedure_sequence_output = diagnostic_idx.unsqueeze(-1) * sequence_output
            # duration_sequence_output = duration_idx.unsqueeze(-1) * sequence_output
            # date_sequence_output = date_idx.unsqueeze(-1) * sequence_output
            # therapeutic_procedure_sequence_output = therapeutic_idx.unsqueeze(-1) * sequence_output
            # sign_symptom_sequence_output = sign_symptom_idx.unsqueeze(-1) * sequence_output
            # lab_value_sequence_output = lab_value_idx.unsqueeze(-1) * sequence_output

            
            # F.tanh or F.relu
            dise_sequence_output = F.relu(self.dise_classifier_2(dise_sequence_output)) # disease logit value
            chem_sequence_output = F.relu(self.chem_classifier_2(chem_sequence_output)) # chemical logit value
            gene_sequence_output = F.relu(self.gene_classifier_2(gene_sequence_output)) # gene/protein logit value
            spec_sequence_output = F.relu(self.spec_classifier_2(spec_sequence_output)) # species logit value
            cellline_sequence_output = F.relu(self.cellline_classifier_2(cellline_sequence_output)) # cell line logit value
            dna_sequence_output = F.relu(self.dna_classifier_2(dna_sequence_output)) # dna logit value
            rna_sequence_output = F.relu(self.rna_classifier_2(rna_sequence_output)) # rna logit value
            celltype_sequence_output = F.relu(self.celltype_classifier_2(celltype_sequence_output)) # cell type logit value

            # biological_structure_sequence_output = F.relu(self.biological_structure_classifier_2(biological_structure_sequence_output)) # biological structure logit value
            # diagnostic_procedure_sequence_output = F.relu(self.diagnostic_procedure_classifier_2(diagnostic_procedure_sequence_output)) # diagnostic procedure logit value
            # duration_sequence_output = F.relu(self.duration_classifier_2(duration_sequence_output)) # duration logit value
            # date_sequence_output = F.relu(self.date_classifier_2(date_sequence_output)) # date logit value
            # therapeutic_procedure_sequence_output = F.relu(self.therapeutic_procedure_classifier_2(therapeutic_procedure_sequence_output)) # therapeutic procedure logit value
            # sign_symptom_sequence_output = F.relu(self.sign_symptom_classifier_2(sign_symptom_sequence_output)) # sign/symptom logit value
            # lab_value_sequence_output = F.relu(self.lab_value_classifier_2(lab_value_sequence_output)) # lab value logit value

            

            dise_logits = self.dise_classifier(dise_sequence_output) # disease logit value
            chem_logits = self.chem_classifier(chem_sequence_output) # chemical logit value
            gene_logits = self.gene_classifier(gene_sequence_output) # gene/protein logit value
            spec_logits = self.spec_classifier(spec_sequence_output) # species logit value
            cellline_logits = self.cellline_classifier(cellline_sequence_output) # cell line logit value
            dna_logits = self.dna_classifier(dna_sequence_output) # dna logit value
            rna_logits = self.rna_classifier(rna_sequence_output) # rna logit value
            celltype_logits = self.celltype_classifier(celltype_sequence_output) # cell type logit value
            # biological_logits = self.biological_structure_classifier(biological_structure_sequence_output) # biological structure logit value
            # diagnostic_logits = self.diagnostic_procedure_classifier(diagnostic_procedure_sequence_output) # diagnostic procedure logit value
            # duration_logits = self.duration_classifier(duration_sequence_output) # duration logit value
            # date_logits = self.date_classifier(date_sequence_output)
            # therapeutic_logits = self.therapeutic_procedure_classifier(therapeutic_procedure_sequence_output) # therapeutic procedure logit value
            # sign_symptom_logits = self.sign_symptom_classifier(sign_symptom_sequence_output) # sign/symptom logit value
            # lab_value_logits = self.lab_value_classifier(lab_value_sequence_output) # lab value logit value

            

            # update logit and sequence_output
            sequence_output =dise_sequence_output + chem_sequence_output + gene_sequence_output + spec_sequence_output + cellline_sequence_output + dna_sequence_output + rna_sequence_output + celltype_sequence_output 
            # \
            #     + biological_structure_sequence_output + diagnostic_procedure_sequence_output + duration_sequence_output + date_sequence_output \
            #     + therapeutic_procedure_sequence_output + sign_symptom_sequence_output + lab_value_sequence_output
                
            logits = dise_logits + chem_logits + gene_logits + spec_logits + cellline_logits \
                + dna_logits + rna_logits + celltype_logits 
                # + biological_logits \
                # + diagnostic_logits + duration_logits + date_logits + therapeutic_logits \
                # + sign_symptom_logits + lab_value_logits 
                

        outputs = (logits, sequence_output)
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            # Only keep active parts of the loss
            if attention_mask is not None:
                if entity_type_ids[0][0].item() == 0:
                    active_loss = attention_mask.view(-1) == 1
                    dise_logits, chem_logits, gene_logits, spec_logits, cellline_logits, \
                    dna_logits, rna_logits, celltype_logits = logits
                    
                    #  biological_logits, diagnostic_logits, \
                    # duration_logits, date_logits, therapeutic_logits, \
                    # sign_symptom_logits, lab_value_logits


                    active_dise_logits = dise_logits.view(-1, self.num_labels)
                    active_chem_logits = chem_logits.view(-1, self.num_labels)
                    active_gene_logits = gene_logits.view(-1, self.num_labels)
                    active_spec_logits = spec_logits.view(-1, self.num_labels)
                    active_cellline_logits = cellline_logits.view(-1, self.num_labels)
                    active_dna_logits = dna_logits.view(-1, self.num_labels)
                    active_rna_logits = rna_logits.view(-1, self.num_labels)
                    active_celltype_logits = celltype_logits.view(-1, self.num_labels)
                    # active_biological_logits = biological_logits.view(-1, self.num_labels)
                    # active_diagnostic_logits = diagnostic_logits.view(-1, self.num_labels)
                    # active_duration_logits = duration_logits.view(-1, self.num_labels)
                    # active_date_logits = date_logits.view(-1, self.num_labels)
                    # active_therapeutic_logits = therapeutic_logits.view(-1, self.num_labels)
                    # active_sign_symptom_logits = sign_symptom_logits.view(-1, self.num_labels)
                    # active_lab_value_logits = lab_value_logits.view(-1, self.num_labels)
                    
                    
                    active_labels = torch.where(
                        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                    )
                    dise_loss = loss_fct(active_dise_logits, active_labels)
                    chem_loss = loss_fct(active_chem_logits, active_labels)
                    gene_loss = loss_fct(active_gene_logits, active_labels)
                    spec_loss = loss_fct(active_spec_logits, active_labels)
                    cellline_loss = loss_fct(active_cellline_logits, active_labels)
                    dna_loss = loss_fct(active_dna_logits, active_labels)
                    rna_loss = loss_fct(active_rna_logits, active_labels)
                    celltype_loss = loss_fct(active_celltype_logits, active_labels)
                    # biological_loss = loss_fct(active_biological_logits, active_labels)
                    # diagnostic_loss = loss_fct(active_diagnostic_logits, active_labels)
                    # duration_loss = loss_fct(active_duration_logits, active_labels)
                    # date_loss = loss_fct(active_date_logits, active_labels)
                    # therapeutic_loss = loss_fct(active_therapeutic_logits, active_labels)
                    # sign_symptom_loss = loss_fct(active_sign_symptom_logits, active_labels)
                    # lab_value_loss = loss_fct(active_lab_value_logits, active_labels)
                    
                    loss = dise_loss + chem_loss + gene_loss + spec_loss + cellline_loss + dna_loss + rna_loss + celltype_loss 
                    # \
                    #      + biological_loss + diagnostic_loss \
                    #     + duration_loss + date_loss + therapeutic_loss + sign_symptom_loss + lab_value_loss 
                        
                    return ((loss,) + outputs)
                else:
                    active_loss = attention_mask.view(-1) == 1
                    active_logits = logits.view(-1, self.num_labels)
                    active_labels = torch.where(
                        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                    )
                    loss = loss_fct(active_logits, active_labels)
                    return ((loss,) + outputs)
            else:
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
                return loss
        else:
            return logits

class RoBERTaMultiNER2(RobertaForTokenClassification):
    def __init__(self, config, num_labels=3):
        super(RoBERTaMultiNER2, self).__init__(config)
        self.num_labels = num_labels
        self.roberta = RobertaModel(config)
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        
        self.dise_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # disease
        self.chem_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # chemical
        self.gene_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # gene/protein
        self.spec_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # species
        self.cellline_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # cell line
        self.dna_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # dna
        self.rna_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # rna
        self.celltype_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # cell type
        
        # self.biological_structure_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # biological structure
        # self.diagnostic_procedure_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # diagnostic procedure
        # self.duration_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # duration
        # self.date_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # date
        # self.therapeutic_procedure_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # therapeutic procedure
        # self.sign_symptom_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # sign/symptom
        # self.lab_value_classifier = torch.nn.Linear(config.hidden_size, self.num_labels) # lab value
        

        self.dise_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.chem_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.gene_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.spec_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.cellline_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.dna_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.rna_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        self.celltype_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        
        # self.biological_structure_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.diagnostic_procedure_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)    
        # self.duration_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)    
        # self.date_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.therapeutic_procedure_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.sign_symptom_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)
        # self.lab_value_classifier_2 = torch.nn.Linear(config.hidden_size, config.hidden_size)

        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, entity_type_ids=None):
        sequence_output = self.roberta(input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, head_mask=None)[0]
        batch_size,max_len,feat_dim = sequence_output.shape
        sequence_output = self.dropout(sequence_output)

        if entity_type_ids[0][0].item() == 0:
            '''
            Raw text data with trained parameters
            '''
            dise_sequence_output = F.relu(self.dise_classifier_2(sequence_output)) # disease logit value
            chem_sequence_output = F.relu(self.chem_classifier_2(sequence_output)) # chemical logit value
            gene_sequence_output = F.relu(self.gene_classifier_2(sequence_output)) # gene/protein logit value
            spec_sequence_output = F.relu(self.spec_classifier_2(sequence_output)) # species logit value
            cellline_sequence_output = F.relu(self.cellline_classifier_2(sequence_output)) # cell line logit value
            dna_sequence_output = F.relu(self.dna_classifier_2(sequence_output)) # dna logit value
            rna_sequence_output = F.relu(self.rna_classifier_2(sequence_output)) # rna logit value
            celltype_sequence_output = F.relu(self.celltype_classifier_2(sequence_output)) # cell type logit value
            
            # biological_structure_sequence_output = F.relu(self.biological_structure_classifier_2(sequence_output)) # biological structure logit value
            # diagnostic_procedure_sequence_output = F.relu(self.diagnostic_procedure_classifier_2(sequence_output)) # diagnostic procedure logit value
            # duration_sequence_output = F.relu(self.duration_classifier_2(sequence_output)) # duration logit value
            # date_sequence_output = F.relu(self.date_classifier_2(sequence_output)) # date logit value
            # therapeutic_procedure_sequence_output = F.relu(self.therapeutic_procedure_classifier_2(sequence_output)) # therapeutic procedure logit value
            # sign_symptom_sequence_output = F.relu(self.sign_symptom_classifier_2(sequence_output)) # sign/symptom logit value
            # lab_value_sequence_output = F.relu(self.lab_value_classifier_2(sequence_output)) # lab value logit value


            dise_logits = self.dise_classifier(dise_sequence_output) # disease logit value
            chem_logits = self.chem_classifier(chem_sequence_output) # chemical logit value
            gene_logits = self.gene_classifier(gene_sequence_output) # gene/protein logit value
            spec_logits = self.spec_classifier(spec_sequence_output) # species logit value
            cellline_logits = self.cellline_classifier(cellline_sequence_output) # cell line logit value
            dna_logits = self.dna_classifier(dna_sequence_output) # dna logit value
            rna_logits = self.rna_classifier(rna_sequence_output) # rna logit value
            celltype_logits = self.celltype_classifier(celltype_sequence_output) # cell type logit value
            
            # biological_logits = self.biological_structure_classifier(biological_structure_sequence_output) # biological structure logit value
            # diagnostic_logits = self.diagnostic_procedure_classifier(diagnostic_procedure_sequence_output) # diagnostic procedure logit value
            # duration_logits = self.duration_classifier(duration_sequence_output) # duration logit value
            # date_logits = self.date_classifier(date_sequence_output) # date logit value
            # therapeutic_logits = self.therapeutic_procedure_classifier(therapeutic_procedure_sequence_output) # therapeutic procedure logit value
            # sign_symptom_logits = self.sign_symptom_classifier(sign_symptom_sequence_output) # sign/symptom logit value
            # lab_value_logits = self.lab_value_classifier(lab_value_sequence_output) # lab value logit value

            

            # update logit and sequence_output
            sequence_output = dise_sequence_output + chem_sequence_output + gene_sequence_output + spec_sequence_output + cellline_sequence_output + dna_sequence_output + rna_sequence_output + celltype_sequence_output 
            # + \
            #     biological_structure_sequence_output + diagnostic_procedure_sequence_output + duration_sequence_output + date_sequence_output + \
            #     therapeutic_procedure_sequence_output + sign_symptom_sequence_output + lab_value_sequence_output 
                
            logits = (dise_logits, chem_logits, gene_logits, spec_logits, cellline_logits, 
                      dna_logits, rna_logits, celltype_logits)
                    #   biological_logits, diagnostic_logits,
                    #   duration_logits, date_logits, therapeutic_logits,
                    #   sign_symptom_logits, lab_value_logits)
        else:
            ''' 
            Train, Eval, Test with pre-defined entity type tags
            '''
            # make 1*1 conv to adopt entity type
            dise_idx = copy.deepcopy(entity_type_ids)
            chem_idx = copy.deepcopy(entity_type_ids)
            gene_idx = copy.deepcopy(entity_type_ids)
            spec_idx = copy.deepcopy(entity_type_ids)
            cellline_idx = copy.deepcopy(entity_type_ids)
            dna_idx = copy.deepcopy(entity_type_ids)
            rna_idx = copy.deepcopy(entity_type_ids)
            celltype_idx = copy.deepcopy(entity_type_ids)
            
            # biological_idx = copy.deepcopy(entity_type_ids)
            # diagnostic_idx = copy.deepcopy(entity_type_ids)
            # duration_idx = copy.deepcopy(entity_type_ids)
            # date_idx = copy.deepcopy(entity_type_ids)
            # therapeutic_idx = copy.deepcopy(entity_type_ids)
            # sign_symptom_idx = copy.deepcopy(entity_type_ids)
            # lab_value_idx = copy.deepcopy(entity_type_ids)

            

            dise_idx[dise_idx != 1] = 0
            chem_idx[chem_idx != 2] = 0
            gene_idx[gene_idx != 3] = 0
            spec_idx[spec_idx != 4] = 0
            cellline_idx[cellline_idx != 5] = 0
            dna_idx[dna_idx != 6] = 0
            rna_idx[rna_idx != 7] = 0
            celltype_idx[celltype_idx != 8] = 0
            # biological_idx[biological_idx != 9] = 0
            # diagnostic_idx[diagnostic_idx != 10] = 0
            # duration_idx[duration_idx != 11] = 0
            # date_idx[date_idx != 12] = 0
            # therapeutic_idx[therapeutic_idx != 13] = 0
            # sign_symptom_idx[sign_symptom_idx != 14] = 0
            # lab_value_idx[lab_value_idx != 15] = 0


            dise_sequence_output = dise_idx.unsqueeze(-1) * sequence_output        
            chem_sequence_output = chem_idx.unsqueeze(-1) * sequence_output
            gene_sequence_output = gene_idx.unsqueeze(-1) * sequence_output
            spec_sequence_output = spec_idx.unsqueeze(-1) * sequence_output
            cellline_sequence_output = cellline_idx.unsqueeze(-1) * sequence_output
            dna_sequence_output = dna_idx.unsqueeze(-1) * sequence_output
            rna_sequence_output = rna_idx.unsqueeze(-1) * sequence_output
            celltype_sequence_output = celltype_idx.unsqueeze(-1) * sequence_output
            # biological_structure_sequence_output = biological_idx.unsqueeze(-1) * sequence_output
            # diagnostic_procedure_sequence_output = diagnostic_idx.unsqueeze(-1) * sequence_output
            # duration_sequence_output = duration_idx.unsqueeze(-1) * sequence_output
            # date_sequence_output = date_idx.unsqueeze(-1) * sequence_output
            # therapeutic_procedure_sequence_output = therapeutic_idx.unsqueeze(-1) * sequence_output
            # sign_symptom_sequence_output = sign_symptom_idx.unsqueeze(-1) * sequence_output
            # lab_value_sequence_output = lab_value_idx.unsqueeze(-1) * sequence_output

            
            # F.tanh or F.relu
            dise_sequence_output = F.relu(self.dise_classifier_2(dise_sequence_output)) # disease logit value
            chem_sequence_output = F.relu(self.chem_classifier_2(chem_sequence_output)) # chemical logit value
            gene_sequence_output = F.relu(self.gene_classifier_2(gene_sequence_output)) # gene/protein logit value
            spec_sequence_output = F.relu(self.spec_classifier_2(spec_sequence_output)) # species logit value
            cellline_sequence_output = F.relu(self.cellline_classifier_2(cellline_sequence_output)) # cell line logit value
            dna_sequence_output = F.relu(self.dna_classifier_2(dna_sequence_output)) # dna logit value
            rna_sequence_output = F.relu(self.rna_classifier_2(rna_sequence_output)) # rna logit value
            celltype_sequence_output = F.relu(self.celltype_classifier_2(celltype_sequence_output)) # cell type logit value

            # biological_structure_sequence_output = F.relu(self.biological_structure_classifier_2(biological_structure_sequence_output)) # biological structure logit value
            # diagnostic_procedure_sequence_output = F.relu(self.diagnostic_procedure_classifier_2(diagnostic_procedure_sequence_output)) # diagnostic procedure logit value
            # duration_sequence_output = F.relu(self.duration_classifier_2(duration_sequence_output)) # duration logit value
            # date_sequence_output = F.relu(self.date_classifier_2(date_sequence_output)) # date logit value
            # therapeutic_procedure_sequence_output = F.relu(self.therapeutic_procedure_classifier_2(therapeutic_procedure_sequence_output)) # therapeutic procedure logit value
            # sign_symptom_sequence_output = F.relu(self.sign_symptom_classifier_2(sign_symptom_sequence_output)) # sign/symptom logit value
            # lab_value_sequence_output = F.relu(self.lab_value_classifier_2(lab_value_sequence_output)) # lab value logit value

            

            dise_logits = self.dise_classifier(dise_sequence_output) # disease logit value
            chem_logits = self.chem_classifier(chem_sequence_output) # chemical logit value
            gene_logits = self.gene_classifier(gene_sequence_output) # gene/protein logit value
            spec_logits = self.spec_classifier(spec_sequence_output) # species logit value
            cellline_logits = self.cellline_classifier(cellline_sequence_output) # cell line logit value
            dna_logits = self.dna_classifier(dna_sequence_output) # dna logit value
            rna_logits = self.rna_classifier(rna_sequence_output) # rna logit value
            celltype_logits = self.celltype_classifier(celltype_sequence_output) # cell type logit value
            # biological_logits = self.biological_structure_classifier(biological_structure_sequence_output) # biological structure logit value
            # diagnostic_logits = self.diagnostic_procedure_classifier(diagnostic_procedure_sequence_output) # diagnostic procedure logit value
            # duration_logits = self.duration_classifier(duration_sequence_output) # duration logit value
            # date_logits = self.date_classifier(date_sequence_output)
            # therapeutic_logits = self.therapeutic_procedure_classifier(therapeutic_procedure_sequence_output) # therapeutic procedure logit value
            # sign_symptom_logits = self.sign_symptom_classifier(sign_symptom_sequence_output) # sign/symptom logit value
            # lab_value_logits = self.lab_value_classifier(lab_value_sequence_output) # lab value logit value

            

            # update logit and sequence_output
            sequence_output =dise_sequence_output + chem_sequence_output + gene_sequence_output + spec_sequence_output + cellline_sequence_output + dna_sequence_output + rna_sequence_output + celltype_sequence_output 
                # + biological_structure_sequence_output + diagnostic_procedure_sequence_output + duration_sequence_output + date_sequence_output \
                # + therapeutic_procedure_sequence_output + sign_symptom_sequence_output + lab_value_sequence_output
                
            logits = dise_logits + chem_logits + gene_logits + spec_logits + cellline_logits \
                + dna_logits + rna_logits + celltype_logits 
                # + biological_logits \
                # + diagnostic_logits + duration_logits + date_logits + therapeutic_logits \
                # + sign_symptom_logits + lab_value_logits 
                

        outputs = (logits, sequence_output)
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            # Only keep active parts of the loss
            if attention_mask is not None:
                if entity_type_ids[0][0].item() == 0:
                    active_loss = attention_mask.view(-1) == 1
                    dise_logits, chem_logits, gene_logits, spec_logits, cellline_logits, \
                    dna_logits, rna_logits, celltype_logits = logits
                    
                    #  biological_logits, diagnostic_logits, \
                    # duration_logits, date_logits, therapeutic_logits, \
                    # sign_symptom_logits, lab_value_logits


                    active_dise_logits = dise_logits.view(-1, self.num_labels)
                    active_chem_logits = chem_logits.view(-1, self.num_labels)
                    active_gene_logits = gene_logits.view(-1, self.num_labels)
                    active_spec_logits = spec_logits.view(-1, self.num_labels)
                    active_cellline_logits = cellline_logits.view(-1, self.num_labels)
                    active_dna_logits = dna_logits.view(-1, self.num_labels)
                    active_rna_logits = rna_logits.view(-1, self.num_labels)
                    active_celltype_logits = celltype_logits.view(-1, self.num_labels)
                    # active_biological_logits = biological_logits.view(-1, self.num_labels)
                    # active_diagnostic_logits = diagnostic_logits.view(-1, self.num_labels)
                    # active_duration_logits = duration_logits.view(-1, self.num_labels)
                    # active_date_logits = date_logits.view(-1, self.num_labels)
                    # active_therapeutic_logits = therapeutic_logits.view(-1, self.num_labels)
                    # active_sign_symptom_logits = sign_symptom_logits.view(-1, self.num_labels)
                    # active_lab_value_logits = lab_value_logits.view(-1, self.num_labels)
                    
                    
                    active_labels = torch.where(
                        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                    )
                    dise_loss = loss_fct(active_dise_logits, active_labels)
                    chem_loss = loss_fct(active_chem_logits, active_labels)
                    gene_loss = loss_fct(active_gene_logits, active_labels)
                    spec_loss = loss_fct(active_spec_logits, active_labels)
                    cellline_loss = loss_fct(active_cellline_logits, active_labels)
                    dna_loss = loss_fct(active_dna_logits, active_labels)
                    rna_loss = loss_fct(active_rna_logits, active_labels)
                    celltype_loss = loss_fct(active_celltype_logits, active_labels)
                    # biological_loss = loss_fct(active_biological_logits, active_labels)
                    # diagnostic_loss = loss_fct(active_diagnostic_logits, active_labels)
                    # duration_loss = loss_fct(active_duration_logits, active_labels)
                    # date_loss = loss_fct(active_date_logits, active_labels)
                    # therapeutic_loss = loss_fct(active_therapeutic_logits, active_labels)
                    # sign_symptom_loss = loss_fct(active_sign_symptom_logits, active_labels)
                    # lab_value_loss = loss_fct(active_lab_value_logits, active_labels)
                    
                    loss = dise_loss + chem_loss + gene_loss + spec_loss + cellline_loss + dna_loss + rna_loss + celltype_loss 
                        #  + biological_loss + diagnostic_loss \
                        # + duration_loss + date_loss + therapeutic_loss + sign_symptom_loss + lab_value_loss 
                        
                    return ((loss,) + outputs)
                else:
                    active_loss = attention_mask.view(-1) == 1
                    active_logits = logits.view(-1, self.num_labels)
                    active_labels = torch.where(
                        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                    )
                    loss = loss_fct(active_logits, active_labels)
                    return ((loss,) + outputs)
            else:
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
                return loss
        else:
            return logits


class NER(BertForTokenClassification):
    def __init__(self, config, num_labels=3):
        super(NER, self).__init__(config)
        self.num_labels = num_labels
        self.bert = BertModel(config)
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        self.classifier = torch.nn.Linear(config.hidden_size, self.num_labels)

        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
        sequence_output = self.bert(input_ids, token_type_ids, attention_mask, head_mask=None)[0]
        batch_size,max_len,feat_dim = sequence_output.shape
        sequence_output = self.dropout(sequence_output)

        logits = self.classifier(sequence_output)

        outputs = (logits, sequence_output)
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            # Only keep active parts of the loss
            if attention_mask is not None:
                active_loss = attention_mask.view(-1) == 1
                active_logits = logits.view(-1, self.num_labels)
                active_labels = torch.where(
                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                )
                loss = loss_fct(active_logits, active_labels)
                return ((loss,) + outputs)
            else:
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
                return loss
        else:
            return logits