# -*- coding: utf-8 -*-
# @Author: <NAME>
# @Date:   2017-10-17 16:47:32
# @Last Modified by:   <NAME>,     Contact: <EMAIL>
# @Last Modified time: 2017-12-14 12:04:48

import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from covebilstm import CoveBiLSTM
from crf import CRF


class Cove_BiLSTM_CRF(nn.Module):
    def __init__(self, data):
        super(Cove_BiLSTM_CRF, self).__init__()
        print "build batched covelstmcrf..."
        self.gpu = data.HP_gpu
        self.average_batch = data.HP_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.lstm = CoveBiLSTM(data)
        self.crf = CRF(label_size, self.gpu)

    def neg_log_likelihood_loss(self, word_inputs, word_seq_lengths, char_inputs, char_seq_lengths, char_seq_recover,
                                batch_label, mask):
        outs = self.lstm.get_output_score(word_inputs, word_seq_lengths, char_inputs, char_seq_lengths,
                                          char_seq_recover)
        batch_size = word_inputs.size(0)
        seq_len = word_inputs.size(1)
        total_loss = self.crf.neg_log_likelihood_loss(outs, mask, batch_label)
        scores, tag_seq = self.crf._viterbi_decode(outs, mask)
        if self.average_batch:
            total_loss = total_loss / batch_size
        return total_loss, tag_seq

    def forward(self, word_inputs, word_seq_lengths, char_inputs, char_seq_lengths, char_seq_recover, mask):
        outs = self.lstm.get_output_score(word_inputs, word_seq_lengths, char_inputs, char_seq_lengths,
                                          char_seq_recover)
        batch_size = word_inputs.size(0)
        seq_len = word_inputs.size(1)
        scores, tag_seq = self.crf._viterbi_decode(outs, mask)
        return tag_seq

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