"""
RNN based NMT and CLF models
"""
import tensorflow as tf
import layers
import model_inputs
from exception import Error
from sampling_utils import SamplingUtils
from abc import ABC
from constants import TaskSetting
from model_components import WordLevelEncoder, SentLevelEncoder, AdvancedDecoder, VanillaDecoder


class RNNModel(ABC):
    def __init__(self, config, is_training):
        # initialize dropout
        self.dropout_source = None
        if config.rnn_use_dropout and config.rnn_dropout_source > 0.0:
            def dropout_source(x):
                return tf.layers.dropout(
                    x, noise_shape=(tf.shape(x)[0], tf.shape(x)[1], 1),
                    rate=config.rnn_dropout_source,
                    training=is_training)
            self.dropout_source = dropout_source

        self.dropout_embedding = None
        if config.rnn_use_dropout and config.rnn_dropout_embedding > 0.0:
            def dropout_embedding(e):
                return tf.layers.dropout(e, noise_shape=tf.shape(e),
                                            rate=config.rnn_dropout_embedding,
                                            training=is_training)
            self.dropout_embedding = dropout_embedding

        self.dropout_hidden = None
        if config.rnn_use_dropout and config.rnn_dropout_hidden > 0.0:
            def dropout_hidden(h):
                return tf.layers.dropout(h, noise_shape=tf.shape(h),
                                            rate=config.rnn_dropout_hidden,
                                            training=is_training)
            self.dropout_hidden = dropout_hidden

        self.dropout_target = None
        if config.rnn_use_dropout and config.rnn_dropout_target > 0.0:
            def dropout_target(y):
                return tf.layers.dropout(
                    y, noise_shape=(tf.shape(y)[0], tf.shape(y)[1], 1),
                    rate=config.rnn_dropout_target,
                    training=is_training)
            self.dropout_target = dropout_target
        
        # loss
        self._loss = None
        self.sampling_utils = SamplingUtils(config)

    @property
    def loss(self):
        if self._loss is None:
            raise Error('loss is not defined.')
        return self._loss


class VanillaNMT(RNNModel):
    def __init__(self, config):
        self.inputs = model_inputs.VanillaNMTModelInputs(config)
        super(VanillaNMT, self).__init__(config, self.inputs.training)

        batch_size = tf.shape(self.inputs.x)[-1]

        with tf.variable_scope("word-encoder"):
            self.encoder = WordLevelEncoder(config, batch_size, 
                                            self.dropout_source, self.dropout_embedding, self.dropout_hidden)
            ctx, embs = self.encoder.get_context(self.inputs.x, self.inputs.x_mask)
        with tf.variable_scope("decoder"):
            if config.tie_encoder_decoder_embeddings:
                tied_embeddings = self.encoder.emb_layer
            else:
                tied_embeddings = None
            self.decoder = VanillaDecoder(config, ctx, embs, self.inputs.x_mask,
                                          self.dropout_target, self.dropout_embedding, self.dropout_hidden, tied_embeddings)
            self.logits = self.decoder.score(self.inputs.y)

        with tf.variable_scope("loss"):
            self.loss_layer = layers.Masked_cross_entropy_loss(
                self.inputs.y, self.inputs.y_mask, config.label_smoothing,
                training=self.inputs.training)
            self._loss_per_sentence = self.loss_layer.forward(self.logits)
            self._loss = tf.reduce_mean(self._loss_per_sentence, keepdims=False)

    @property
    def loss_per_sentence(self):
        return self._loss_per_sentence


class HANModel(RNNModel):
    def __init__(self, config):
        self.inputs = model_inputs.CLFModelInputs(config)
        super(HANModel, self).__init__(config, self.inputs.training)
        
        self._build_hierarchical_rnn(config)
        *_, batch_size = tf.unstack(tf.shape(self.inputs.x))

        with tf.variable_scope("fully-connect-1"):
            self.fc_layer_1 = layers.FeedForwardLayer(in_size=2 * config.sent_rnn_state_size, out_size=config.sec_repr_size, batch_size=batch_size)
            self.sec_repr = self.fc_layer_1.forward(self.sent_attn_context)

        with tf.variable_scope("fully-connect-2"):
            self.fc_layer_2 = layers.FeedForwardLayer(in_size=config.sec_repr_size, out_size=TaskSetting.NUM_CLASS, batch_size=batch_size)
            # logits: [1, batch_size, class_num]
            self.logits = tf.expand_dims(self.fc_layer_2.forward(self.sec_repr), axis=0, name='logits')
            self._preds = tf.argmax(tf.squeeze(self.logits, axis=0), axis=1, name='preds')
            self._dist = tf.nn.softmax(tf.squeeze(self.logits, axis=0), axis=1, name='dist')

        with tf.variable_scope("clf-loss"):
            # we can view the output label as a one-token sentence
            self.loss_layer = layers.Masked_cross_entropy_loss(
                self.inputs.y, self.inputs.y_mask, config.label_smoothing,
                training=self.inputs.training)
            self._loss_per_sec = self.loss_layer.forward(self.logits)
            self._loss = tf.reduce_mean(self._loss_per_sec, keepdims=False)

    def __init_word_encoder_n_attn(self, config, total_sent_num):
        # word-level encoder which is shared with NMT
        with tf.variable_scope("word-encoder", reuse=tf.AUTO_REUSE):
            # the input of Encoder put batch_size as its last dimension
            self.word_encoder = WordLevelEncoder(config, total_sent_num,
                                                 self.dropout_source, self.dropout_embedding, self.dropout_hidden)

            self.word_level_ctx, self.word_embeds = self.word_encoder.get_context(self.word_input, self.word_input_mask)
            # word_level_ctx - a simple concate of bi-direction GRU states
            # [time_steps, batch_size, 2 * state_size]
            word_level_query = tf.get_variable(name='word-level-query', 
                                                    shape=[2 * config.word_rnn_state_size], 
                                                    dtype=tf.float32, 
                                                    initializer=tf.random_normal_initializer,
                                                    trainable=True)

        with tf.variable_scope("word-attention", reuse=tf.AUTO_REUSE):
            self.word_attn_layer = layers.AttentionStep(
                                        context=self.word_level_ctx,
                                        context_state_size=2 * config.word_rnn_state_size,          # the size of context
                                        context_mask=self.word_input_mask,
                                        state_size=2 * config.word_rnn_state_size,                  # the size of query vector
                                        hidden_size=2 * config.word_rnn_state_size,                
                                        use_layer_norm=config.rnn_layer_normalization,   
                                        dropout_context=self.dropout_hidden,
                                        dropout_state=self.dropout_hidden)
            # word_attn [batch_size * sent_len, state_size * 2]
            self.word_attn_context, self.word_attn_dist = self.word_attn_layer.forward(
                        tf.broadcast_to(word_level_query, [total_sent_num, 2 * config.word_rnn_state_size])
                        )

    def __init_sent_encoder_n_attn(self, config, total_sec_num):
        # sent-level encoder
        with tf.variable_scope("sent-encoder", reuse=tf.AUTO_REUSE):
            self.sent_encoder = SentLevelEncoder(config, total_sec_num, self.dropout_hidden)

            sent_level_ctx = self.sent_encoder.get_context(self.sent_input, self.sent_input_mask)

            sent_level_query = tf.get_variable(name='sent-level-query',
                                                shape=[2 * config.sent_rnn_state_size],
                                                dtype=tf.float32,
                                                initializer=tf.random_normal_initializer,
                                                trainable=True)

        with tf.variable_scope("sent-attention", reuse=tf.AUTO_REUSE):
            # sent_attn [batch_size, state_size * 2]
            self.sent_attn_layer = layers.AttentionStep(
                context=sent_level_ctx,
                context_state_size=2 * config.sent_rnn_state_size,
                context_mask=self.sent_input_mask,
                state_size=2 * config.sent_rnn_state_size,
                hidden_size=2 * config.sent_rnn_state_size,
                use_layer_norm=config.rnn_layer_normalization,
                dropout_context=self.dropout_hidden,
                dropout_state=self.dropout_hidden)
            
            # now batch channel first
            self.sent_attn_context, self.sent_attn_dist = self.sent_attn_layer.forward(
                tf.broadcast_to(sent_level_query, (total_sec_num, 2 * config.sent_rnn_state_size))
                )

    def _build_hierarchical_rnn(self, config):
        # batch_size is the alias of the number of sections
        _, sent_len, word_len, batch_size = tf.unstack(tf.shape(self.inputs.x))
        
        # reshape layer for word encoder
        self.word_input = tf.reshape(
            tf.transpose(self.inputs.x, (0, 2, 1, 3)),   # n_factor * sent_len * word_len * n_sec -> n_factor * word_len * sent_len * n_sec
            (config.factors, word_len, sent_len * batch_size))
        self.word_input_mask = tf.reshape(
            tf.transpose(self.inputs.x_mask, (1, 0, 2)),   # sent_len * word_len * n_sec ->  word_len * sent_len * n_sec
            (word_len, sent_len * batch_size))

        # init word-level encoder and attention
        self.__init_word_encoder_n_attn(config, sent_len * batch_size)

        # reshape layer for sent encoder
        # reshape context vector, always batch channel last
        self.sent_input = tf.reshape(self.word_attn_context, (sent_len, batch_size, self.word_attn_context.shape[-1]))
        self.sent_input_mask = tf.sign(tf.reduce_sum(self.inputs.x_mask, axis=1))

        # init sentence-level encoder and attention
        self.__init_sent_encoder_n_attn(config, batch_size)

    @property
    def loss_per_sec(self):
        return self._loss_per_sec
    
    @property
    def preds(self):
        return self._preds

    @property
    def output_dist(self):
        return self._dist


class StructureAwareNMT(HANModel):
    def __init__(self, config):
        self.inputs = model_inputs.NMTModelInputs(config)
        super(HANModel, self).__init__(config, self.inputs.training)

        _, sent_len, _, batch_size = tf.unstack(tf.shape(self.inputs.x))
        _, word_len_y, _ = tf.unstack(tf.shape(self.inputs.y))

        self._build_hierarchical_rnn(config)

        # the second FC layer for text-Clf is abandoned
        with tf.variable_scope("fully-connect-1", reuse=tf.AUTO_REUSE):
            self.fc_layer_1 = layers.FeedForwardLayer(in_size=2*config.sent_rnn_state_size, 
                                                      out_size=config.sec_repr_size, 
                                                      batch_size=batch_size)
            self.sec_repr = self.fc_layer_1.forward(self.sent_attn_context)

        self.input_y_2d = tf.reshape(
            tf.transpose(self.inputs.y, (1, 0, 2)),
            (word_len_y, sent_len * batch_size))
        self.input_y_mask_2d = tf.reshape(
            tf.transpose(self.inputs.y_mask, (1, 0, 2)),
            (word_len_y, sent_len * batch_size))
        
        with tf.variable_scope("decoder"):
            if config.tie_encoder_decoder_embeddings:
                tied_embeddings = self.word_encoder.emb_layer
            else:
                tied_embeddings = None
            self.decoder = AdvancedDecoder(config, self.word_level_ctx, self.word_embeds, self.word_input_mask,
                                           self.dropout_target, self.dropout_embedding, self.dropout_hidden, tied_embeddings)

            # sec_repr is produced by HAN, which is not a part of decoder
            self.sec_repr_broadcast = tf.reshape(
                tf.broadcast_to(self.sec_repr, (sent_len, batch_size, config.sec_repr_size)),
                (sent_len * batch_size, config.sec_repr_size))
            self.logits = self.decoder.score(self.input_y_2d, self.sec_repr_broadcast)

        with tf.variable_scope("nmt-loss"):
            self.loss_layer = layers.Masked_cross_entropy_loss(
                self.input_y_2d, self.input_y_mask_2d, config.label_smoothing,
                training=self.inputs.training)
            self._loss_per_sentence = self.loss_layer.forward(self.logits)    # sum word cross entropy over sentences
            self._loss = tf.reduce_mean(self._loss_per_sentence, keepdims=False)

    @property
    def loss_per_sentence(self):
        return self._loss_per_sentence


class JointModel(RNNModel):
    def __init__(self, config):
        self.clf_model = HANModel(config)
        self.nmt_model = StructureAwareNMT(config)


def model_builder(config):
    MODELS = {'clf': HANModel, 'nmt': StructureAwareNMT, 'joint': JointModel, 'baseline-nmt': VanillaNMT}
    return MODELS[config.mode](config)

"""
RNN based NMT and CLF models
"""
import tensorflow as tf
import layers
import model_inputs
from exception import Error
from sampling_utils import SamplingUtils
from abc import ABC
from constants import TaskSetting
from model_components import WordLevelEncoder, SentLevelEncoder, AdvancedDecoder, VanillaDecoder


class RNNModel(ABC):
    def __init__(self, config, is_training):
        # initialize dropout
        self.dropout_source = None
        if config.rnn_use_dropout and config.rnn_dropout_source > 0.0:
            def dropout_source(x):
                return tf.layers.dropout(
                    x, noise_shape=(tf.shape(x)[0], tf.shape(x)[1], 1),
                    rate=config.rnn_dropout_source,
                    training=is_training)
            self.dropout_source = dropout_source

        self.dropout_embedding = None
        if config.rnn_use_dropout and config.rnn_dropout_embedding > 0.0:
            def dropout_embedding(e):
                return tf.layers.dropout(e, noise_shape=tf.shape(e),
                                            rate=config.rnn_dropout_embedding,
                                            training=is_training)
            self.dropout_embedding = dropout_embedding

        self.dropout_hidden = None
        if config.rnn_use_dropout and config.rnn_dropout_hidden > 0.0:
            def dropout_hidden(h):
                return tf.layers.dropout(h, noise_shape=tf.shape(h),
                                            rate=config.rnn_dropout_hidden,
                                            training=is_training)
            self.dropout_hidden = dropout_hidden

        self.dropout_target = None
        if config.rnn_use_dropout and config.rnn_dropout_target > 0.0:
            def dropout_target(y):
                return tf.layers.dropout(
                    y, noise_shape=(tf.shape(y)[0], tf.shape(y)[1], 1),
                    rate=config.rnn_dropout_target,
                    training=is_training)
            self.dropout_target = dropout_target
        
        # loss
        self._loss = None
        self.sampling_utils = SamplingUtils(config)

    @property
    def loss(self):
        if self._loss is None:
            raise Error('loss is not defined.')
        return self._loss


class VanillaNMT(RNNModel):
    def __init__(self, config):
        self.inputs = model_inputs.VanillaNMTModelInputs(config)
        super(VanillaNMT, self).__init__(config, self.inputs.training)

        batch_size = tf.shape(self.inputs.x)[-1]

        with tf.variable_scope("word-encoder"):
            self.encoder = WordLevelEncoder(config, batch_size, 
                                            self.dropout_source, self.dropout_embedding, self.dropout_hidden)
            ctx, embs = self.encoder.get_context(self.inputs.x, self.inputs.x_mask)
        with tf.variable_scope("decoder"):
            if config.tie_encoder_decoder_embeddings:
                tied_embeddings = self.encoder.emb_layer
            else:
                tied_embeddings = None
            self.decoder = VanillaDecoder(config, ctx, embs, self.inputs.x_mask,
                                          self.dropout_target, self.dropout_embedding, self.dropout_hidden, tied_embeddings)
            self.logits = self.decoder.score(self.inputs.y)

        with tf.variable_scope("loss"):
            self.loss_layer = layers.Masked_cross_entropy_loss(
                self.inputs.y, self.inputs.y_mask, config.label_smoothing,
                training=self.inputs.training)
            self._loss_per_sentence = self.loss_layer.forward(self.logits)
            self._loss = tf.reduce_mean(self._loss_per_sentence, keepdims=False)

    @property
    def loss_per_sentence(self):
        return self._loss_per_sentence


class HANModel(RNNModel):
    def __init__(self, config):
        self.inputs = model_inputs.CLFModelInputs(config)
        super(HANModel, self).__init__(config, self.inputs.training)
        
        self._build_hierarchical_rnn(config)
        *_, batch_size = tf.unstack(tf.shape(self.inputs.x))

        with tf.variable_scope("fully-connect-1"):
            self.fc_layer_1 = layers.FeedForwardLayer(in_size=2 * config.sent_rnn_state_size, out_size=config.sec_repr_size, batch_size=batch_size)
            self.sec_repr = self.fc_layer_1.forward(self.sent_attn_context)

        with tf.variable_scope("fully-connect-2"):
            self.fc_layer_2 = layers.FeedForwardLayer(in_size=config.sec_repr_size, out_size=TaskSetting.NUM_CLASS, batch_size=batch_size)
            # logits: [1, batch_size, class_num]
            self.logits = tf.expand_dims(self.fc_layer_2.forward(self.sec_repr), axis=0, name='logits')
            self._preds = tf.argmax(tf.squeeze(self.logits, axis=0), axis=1, name='preds')
            self._dist = tf.nn.softmax(tf.squeeze(self.logits, axis=0), axis=1, name='dist')

        with tf.variable_scope("clf-loss"):
            # we can view the output label as a one-token sentence
            self.loss_layer = layers.Masked_cross_entropy_loss(
                self.inputs.y, self.inputs.y_mask, config.label_smoothing,
                training=self.inputs.training)
            self._loss_per_sec = self.loss_layer.forward(self.logits)
            self._loss = tf.reduce_mean(self._loss_per_sec, keepdims=False)

    def __init_word_encoder_n_attn(self, config, total_sent_num):
        # word-level encoder which is shared with NMT
        with tf.variable_scope("word-encoder", reuse=tf.AUTO_REUSE):
            # the input of Encoder put batch_size as its last dimension
            self.word_encoder = WordLevelEncoder(config, total_sent_num,
                                                 self.dropout_source, self.dropout_embedding, self.dropout_hidden)

            self.word_level_ctx, self.word_embeds = self.word_encoder.get_context(self.word_input, self.word_input_mask)
            # word_level_ctx - a simple concate of bi-direction GRU states
            # [time_steps, batch_size, 2 * state_size]
            word_level_query = tf.get_variable(name='word-level-query', 
                                                    shape=[2 * config.word_rnn_state_size], 
                                                    dtype=tf.float32, 
                                                    initializer=tf.random_normal_initializer,
                                                    trainable=True)

        with tf.variable_scope("word-attention", reuse=tf.AUTO_REUSE):
            self.word_attn_layer = layers.AttentionStep(
                                        context=self.word_level_ctx,
                                        context_state_size=2 * config.word_rnn_state_size,          # the size of context
                                        context_mask=self.word_input_mask,
                                        state_size=2 * config.word_rnn_state_size,                  # the size of query vector
                                        hidden_size=2 * config.word_rnn_state_size,                
                                        use_layer_norm=config.rnn_layer_normalization,   
                                        dropout_context=self.dropout_hidden,
                                        dropout_state=self.dropout_hidden)
            # word_attn [batch_size * sent_len, state_size * 2]
            self.word_attn_context, self.word_attn_dist = self.word_attn_layer.forward(
                        tf.broadcast_to(word_level_query, [total_sent_num, 2 * config.word_rnn_state_size])
                        )

    def __init_sent_encoder_n_attn(self, config, total_sec_num):
        # sent-level encoder
        with tf.variable_scope("sent-encoder", reuse=tf.AUTO_REUSE):
            self.sent_encoder = SentLevelEncoder(config, total_sec_num, self.dropout_hidden)

            sent_level_ctx = self.sent_encoder.get_context(self.sent_input, self.sent_input_mask)

            sent_level_query = tf.get_variable(name='sent-level-query',
                                                shape=[2 * config.sent_rnn_state_size],
                                                dtype=tf.float32,
                                                initializer=tf.random_normal_initializer,
                                                trainable=True)

        with tf.variable_scope("sent-attention", reuse=tf.AUTO_REUSE):
            # sent_attn [batch_size, state_size * 2]
            self.sent_attn_layer = layers.AttentionStep(
                context=sent_level_ctx,
                context_state_size=2 * config.sent_rnn_state_size,
                context_mask=self.sent_input_mask,
                state_size=2 * config.sent_rnn_state_size,
                hidden_size=2 * config.sent_rnn_state_size,
                use_layer_norm=config.rnn_layer_normalization,
                dropout_context=self.dropout_hidden,
                dropout_state=self.dropout_hidden)
            
            # now batch channel first
            self.sent_attn_context, self.sent_attn_dist = self.sent_attn_layer.forward(
                tf.broadcast_to(sent_level_query, (total_sec_num, 2 * config.sent_rnn_state_size))
                )

    def _build_hierarchical_rnn(self, config):
        # batch_size is the alias of the number of sections
        _, sent_len, word_len, batch_size = tf.unstack(tf.shape(self.inputs.x))
        
        # reshape layer for word encoder
        self.word_input = tf.reshape(
            tf.transpose(self.inputs.x, (0, 2, 1, 3)),   # n_factor * sent_len * word_len * n_sec -> n_factor * word_len * sent_len * n_sec
            (config.factors, word_len, sent_len * batch_size))
        self.word_input_mask = tf.reshape(
            tf.transpose(self.inputs.x_mask, (1, 0, 2)),   # sent_len * word_len * n_sec ->  word_len * sent_len * n_sec
            (word_len, sent_len * batch_size))

        # init word-level encoder and attention
        self.__init_word_encoder_n_attn(config, sent_len * batch_size)

        # reshape layer for sent encoder
        # reshape context vector, always batch channel last
        self.sent_input = tf.reshape(self.word_attn_context, (sent_len, batch_size, self.word_attn_context.shape[-1]))
        self.sent_input_mask = tf.sign(tf.reduce_sum(self.inputs.x_mask, axis=1))

        # init sentence-level encoder and attention
        self.__init_sent_encoder_n_attn(config, batch_size)

    @property
    def loss_per_sec(self):
        return self._loss_per_sec
    
    @property
    def preds(self):
        return self._preds

    @property
    def output_dist(self):
        return self._dist


class StructureAwareNMT(HANModel):
    def __init__(self, config):
        self.inputs = model_inputs.NMTModelInputs(config)
        super(HANModel, self).__init__(config, self.inputs.training)

        _, sent_len, _, batch_size = tf.unstack(tf.shape(self.inputs.x))
        _, word_len_y, _ = tf.unstack(tf.shape(self.inputs.y))

        self._build_hierarchical_rnn(config)

        # the second FC layer for text-Clf is abandoned
        with tf.variable_scope("fully-connect-1", reuse=tf.AUTO_REUSE):
            self.fc_layer_1 = layers.FeedForwardLayer(in_size=2*config.sent_rnn_state_size, 
                                                      out_size=config.sec_repr_size, 
                                                      batch_size=batch_size)
            self.sec_repr = self.fc_layer_1.forward(self.sent_attn_context)

        self.input_y_2d = tf.reshape(
            tf.transpose(self.inputs.y, (1, 0, 2)),
            (word_len_y, sent_len * batch_size))
        self.input_y_mask_2d = tf.reshape(
            tf.transpose(self.inputs.y_mask, (1, 0, 2)),
            (word_len_y, sent_len * batch_size))
        
        with tf.variable_scope("decoder"):
            if config.tie_encoder_decoder_embeddings:
                tied_embeddings = self.word_encoder.emb_layer
            else:
                tied_embeddings = None
            self.decoder = AdvancedDecoder(config, self.word_level_ctx, self.word_embeds, self.word_input_mask,
                                           self.dropout_target, self.dropout_embedding, self.dropout_hidden, tied_embeddings)

            # sec_repr is produced by HAN, which is not a part of decoder
            self.sec_repr_broadcast = tf.reshape(
                tf.broadcast_to(self.sec_repr, (sent_len, batch_size, config.sec_repr_size)),
                (sent_len * batch_size, config.sec_repr_size))
            self.logits = self.decoder.score(self.input_y_2d, self.sec_repr_broadcast)

        with tf.variable_scope("nmt-loss"):
            self.loss_layer = layers.Masked_cross_entropy_loss(
                self.input_y_2d, self.input_y_mask_2d, config.label_smoothing,
                training=self.inputs.training)
            self._loss_per_sentence = self.loss_layer.forward(self.logits)    # sum word cross entropy over sentences
            self._loss = tf.reduce_mean(self._loss_per_sentence, keepdims=False)

    @property
    def loss_per_sentence(self):
        return self._loss_per_sentence


class JointModel(RNNModel):
    def __init__(self, config):
        self.clf_model = HANModel(config)
        self.nmt_model = StructureAwareNMT(config)


def model_builder(config):
    MODELS = {'clf': HANModel, 'nmt': StructureAwareNMT, 'joint': JointModel, 'baseline-nmt': VanillaNMT}
    return MODELS[config.mode](config)

RNN based NMT and CLF models # initialize dropout # loss # logits: [1, batch_size, class_num] # we can view the output label as a one-token sentence # word-level encoder which is shared with NMT # the input of Encoder put batch_size as its last dimension # word_level_ctx - a simple concate of bi-direction GRU states # [time_steps, batch_size, 2 * state_size] # the size of context # the size of query vector # word_attn [batch_size * sent_len, state_size * 2] # sent-level encoder # sent_attn [batch_size, state_size * 2] # now batch channel first # batch_size is the alias of the number of sections # reshape layer for word encoder # n_factor * sent_len * word_len * n_sec -> n_factor * word_len * sent_len * n_sec # sent_len * word_len * n_sec ->  word_len * sent_len * n_sec # init word-level encoder and attention # reshape layer for sent encoder # reshape context vector, always batch channel last # init sentence-level encoder and attention # the second FC layer for text-Clf is abandoned # sec_repr is produced by HAN, which is not a part of decoder # sum word cross entropy over sentences