import torch
import torch.nn as nn
from transformers import PreTrainedModel
from collections import OrderedDict 
from transformers.modeling_outputs import SequenceClassifierOutput
from typing import List, Optional, Tuple, Union
from .configuration import MultiLabelClassifierConfig

class MultiLabelClassifierModel(PreTrainedModel):
    config_class = MultiLabelClassifierConfig    

    def __init__(self, config):
        super().__init__(config)
        
        self.nlp_model = torch.hub.load('huggingface/pytorch-transformers', 'model', config.transformer_name)
        self.rnn = nn.GRU(config.embedding_dim, 
                      config.hidden_dim,
                      num_layers = config.num_layers,
                      bidirectional = config.bidirectional,
                      batch_first = True,
                      dropout = 0 if config.num_layers < 2 else config.dropout)
        self.dropout = nn.Dropout(config.dropout)
        self.out = nn.Linear(config.hidden_dim * 2 if config.bidirectional else config.hidden_dim, config.num_classes)

    def forward(self, 
                input_ids: Optional[torch.Tensor] = None,
                attention_mask: Optional[torch.Tensor] = None,
                token_type_ids: Optional[torch.Tensor] = None,
                position_ids: Optional[torch.Tensor] = None,
                head_mask: Optional[torch.Tensor] = None,
                inputs_embeds: Optional[torch.Tensor] = None,
                output_attentions: Optional[bool] = None,
                output_hidden_states: Optional[bool] = None,
                return_dict: Optional[bool] = None,
        )-> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        output = self.nlp_model(
          input_ids=input_ids,
          attention_mask=attention_mask,
          token_type_ids=token_type_ids,
          position_ids=position_ids,
          head_mask=head_mask,
          inputs_embeds=inputs_embeds,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
        )
        _, hidden = self.rnn(output['last_hidden_state'])
        if self.rnn.bidirectional:
            hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1))
        else:
            hidden = self.dropout(hidden[-1,:,:])

        logits = self.out(hidden)
        return SequenceClassifierOutput(
            logits=logits,
            hidden_states=output.hidden_states,
            attentions=output.attentions,
        )