import torch
from torch import nn
from torch.nn import init

from fastNLP.modules.encoder.bert import BertModel


class Classifier(nn.Module):
    def __init__(self, hidden_size):
        super(Classifier, self).__init__()
        self.linear = nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, inputs, mask_cls):
        h = self.linear(inputs).squeeze(-1) # [batch_size, seq_len]
        sent_scores = self.sigmoid(h) * mask_cls.float()
        return sent_scores


class BertSum(nn.Module):
    
    def __init__(self, hidden_size=768):
        super(BertSum, self).__init__()
        
        self.hidden_size = hidden_size

        self.encoder = BertModel.from_pretrained('/path/to/uncased_L-12_H-768_A-12')
        self.decoder = Classifier(self.hidden_size)

    def forward(self, article, segment_id, cls_id):
         
        # print(article.device)
        # print(segment_id.device)
        # print(cls_id.device)

        input_mask = 1 - (article == 0).long()
        mask_cls = 1 - (cls_id == -1).long()
        assert input_mask.size() == article.size()
        assert mask_cls.size() == cls_id.size()

        bert_out = self.encoder(article, token_type_ids=segment_id, attention_mask=input_mask)
        bert_out = bert_out[0][-1] # last layer

        sent_emb = bert_out[torch.arange(bert_out.size(0)).unsqueeze(1), cls_id]
        sent_emb = sent_emb * mask_cls.unsqueeze(-1).float()
        assert sent_emb.size() == (article.size(0), cls_id.size(1), self.hidden_size) # [batch_size, seq_len, hidden_size]

        sent_scores = self.decoder(sent_emb, mask_cls) # [batch_size, seq_len]
        assert sent_scores.size() == (article.size(0), cls_id.size(1))

        return {'pred': sent_scores, 'mask': mask_cls}