# -*- coding: utf-8 -*-
# @Author: Jie Yang
# @Date:   2017-10-17 16:47:32
# @Last Modified by:   Jie Yang,     Contact: jieynlp@gmail.com
# @Last Modified time: 2019-02-13 11:49:38

from __future__ import print_function
from __future__ import absolute_import
import torch
import torch.nn as nn
from .wordsequence import WordSequence
from .crf import CRF
import sys
class SeqLabel(nn.Module):
    def __init__(self, data):
        """

        :param data:
        """
        super(SeqLabel, self).__init__()
        self.use_crf = data.use_crf
        if not data.silence:
            print("build sequence labeling network...")
            print("use_char: ", data.use_char)
            if data.use_word_seq:
                print("word feature extractor: ", data.word_feature_extractor)
                if data.use_char:
                    print("char feature extractor: ", data.char_feature_extractor)
            
            print("use crf: ", self.use_crf)

        self.gpu = data.HP_gpu
        self.device = data.device
        self.device = data.device
        self.average_batch = data.average_batch_loss
        ## add two more label for downlayer lstm, use original label size for CRF
        label_size = data.label_alphabet_size
        # data.label_alphabet_size += 2
        self.word_hidden = WordSequence(data)
        if self.use_crf:
            self.crf = CRF(label_size, self.gpu, self.device ,data.silence )


    def calculate_loss(self, *input):
        """

        :param input:
        :return:
        """
        word_inputs = input[0]
        batch_label = input[7]
        mask = input[8]
        outs = self.word_hidden(*input)
        if self.gpu:
            outs = outs.to(self.device)
        batch_size = word_inputs.size(0)
        seq_len = word_inputs.size(1)
        if self.use_crf:
            total_loss = self.crf.neg_log_likelihood_loss(outs, mask, batch_label)
            scores, tag_seq = self.crf._viterbi_decode(outs, mask)
        else:
            loss_function = nn.CrossEntropyLoss(ignore_index=0)#, reduction='sum')
            outs = outs.view(batch_size * seq_len, -1)
            total_loss = loss_function(outs, batch_label.view(batch_size * seq_len))
            _, tag_seq  = torch.max(outs, 1)
            tag_seq = tag_seq.view(batch_size, seq_len)
        if self.average_batch:
            total_loss = total_loss / batch_size
        return total_loss, tag_seq


    def forward(self, *input):
        """

        :param input:
        :return:
        """
        word_inputs = input[0]
        mask = input[8]
        outs = self.word_hidden(*input)
        batch_size = word_inputs.size(0)
        seq_len = word_inputs.size(1)
        if self.use_crf:
            scores, tag_seq = self.crf._viterbi_decode(outs, mask)
        else:
            outs = outs.view(batch_size * seq_len, -1)
            _, tag_seq  = torch.max(outs, 1)
            tag_seq = tag_seq.view(batch_size, seq_len)
            ## filter padded position with zero
            tag_seq = mask.long() * tag_seq
        return tag_seq


    # def get_lstm_features(self, word_inputs, word_seq_lengths, char_inputs, char_seq_lengths, char_seq_recover):
    #     return self.word_hidden(word_inputs, word_seq_lengths, char_inputs, char_seq_lengths, char_seq_recover)


    def decode_nbest(self, *input):
        """

        :param input:
        :return:
        """
        word_inputs, feature_inputs, word_seq_lengths, char_inputs, char_seq_lengths, char_seq_recover, batch_word_text, training, mask,  nbest = input[:10]
        if not self.use_crf:
            print("Nbest output is currently supported only for CRF! Exit...")
            sys.exit()
        outs = self.word_hidden(*input)
        scores, tag_seq = self.crf._viterbi_decode_nbest(outs, mask, nbest)
        return scores, tag_seq

