# Path: nmt_chainer/utilities/utils.py
# def compute_lexicon_matrix(src_batch, lexical_probability_dictionary, V_tgt):
#     real_mb_size = src_batch[0].data.shape[0]
#     max_source_size = len(src_batch)
#     lexicon_matrix = np.zeros((real_mb_size, max_source_size, V_tgt), dtype=np.float32)
#     for src_pos in six.moves.range(max_source_size):
#         # TODO: check if this is too slow
#         src_batch_cpu = cuda.to_cpu(src_batch[src_pos].data)
#         for num_mb in six.moves.range(real_mb_size):
#             src_idx = int(src_batch_cpu[num_mb])
#             if src_idx in lexical_probability_dictionary:
#                 for tgt_idx, lex_prob in six.iteritems(lexical_probability_dictionary[src_idx]):
#                     lexicon_matrix[num_mb][src_pos][tgt_idx] = lex_prob
#     return lexicon_matrix


import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import six
import logging
from chainer import cuda, Variable
from chainer import Link, Chain, ChainList
from . import rnn_cells
from nmt_chainer.utilities.utils import compute_lexicon_matrix
from . import attention
from . import decoder_cells
from . import encoders
        src_batch is as in the sequence parameter of Encoder
        tgt_batch is as in the targets parameter of Decoder
        src_mask is as in the mask parameter of Encoder

        return loss and attention values
    """

    def __init__(self, Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl, attn_cls=attention.AttentionModule, init_orth=False, use_bn_length=0,
                 encoder_cell_type=rnn_cells.LSTMCell,
                 decoder_cell_type=rnn_cells.LSTMCell,
                 lexical_probability_dictionary=None, lex_epsilon=1e-3,
                 use_goto_attention=False
                 ):
        log.info("constructing encoder decoder with Vi:%i Ei:%i Hi:%i Vo:%i Eo:%i Ho:%i Ha:%i Hl:%i" %
                 (Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl))
        super(EncoderDecoder, self).__init__(
            enc=encoders.make_encoder(Vi, Ei, Hi, init_orth=init_orth, use_bn_length=use_bn_length,
                                      cell_type=encoder_cell_type),
            dec=decoder_cells.Decoder(Vo, Eo, Ho, Ha, 2 * Hi, Hl, attn_cls=attn_cls, init_orth=init_orth,
                                      cell_type=decoder_cell_type, use_goto_attention=use_goto_attention)
        )
        self.Vo = Vo
        self.lexical_probability_dictionary = lexical_probability_dictionary
        self.lex_epsilon = lex_epsilon

    def encdec_type(self):
        return "classic"

    def compute_lexicon_probability_matrix(self, src_batch):
        if self.lexical_probability_dictionary is not None:

# Path: nmt_chainer/utilities/utils.py
# def compute_lexicon_matrix(src_batch, lexical_probability_dictionary, V_tgt):
#     real_mb_size = src_batch[0].data.shape[0]
#     max_source_size = len(src_batch)
#     lexicon_matrix = np.zeros((real_mb_size, max_source_size, V_tgt), dtype=np.float32)
#     for src_pos in six.moves.range(max_source_size):
#         # TODO: check if this is too slow
#         src_batch_cpu = cuda.to_cpu(src_batch[src_pos].data)
#         for num_mb in six.moves.range(real_mb_size):
#             src_idx = int(src_batch_cpu[num_mb])
#             if src_idx in lexical_probability_dictionary:
#                 for tgt_idx, lex_prob in six.iteritems(lexical_probability_dictionary[src_idx]):
#                     lexicon_matrix[num_mb][src_pos][tgt_idx] = lex_prob
#     return lexicon_matrix


# Path: nmt_chainer/models/encoder_decoder.py
import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import six
import logging
from chainer import cuda, Variable
from chainer import Link, Chain, ChainList
from . import rnn_cells
from nmt_chainer.utilities.utils import compute_lexicon_matrix
from . import attention
from . import decoder_cells
from . import encoders

        src_batch is as in the sequence parameter of Encoder
        tgt_batch is as in the targets parameter of Decoder
        src_mask is as in the mask parameter of Encoder

        return loss and attention values
    """

    def __init__(self, Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl, attn_cls=attention.AttentionModule, init_orth=False, use_bn_length=0,
                 encoder_cell_type=rnn_cells.LSTMCell,
                 decoder_cell_type=rnn_cells.LSTMCell,
                 lexical_probability_dictionary=None, lex_epsilon=1e-3,
                 use_goto_attention=False
                 ):
        log.info("constructing encoder decoder with Vi:%i Ei:%i Hi:%i Vo:%i Eo:%i Ho:%i Ha:%i Hl:%i" %
                 (Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl))
        super(EncoderDecoder, self).__init__(
            enc=encoders.make_encoder(Vi, Ei, Hi, init_orth=init_orth, use_bn_length=use_bn_length,
                                      cell_type=encoder_cell_type),
            dec=decoder_cells.Decoder(Vo, Eo, Ho, Ha, 2 * Hi, Hl, attn_cls=attn_cls, init_orth=init_orth,
                                      cell_type=decoder_cell_type, use_goto_attention=use_goto_attention)
        )
        self.Vo = Vo
        self.lexical_probability_dictionary = lexical_probability_dictionary
        self.lex_epsilon = lex_epsilon

    def encdec_type(self):
        return "classic"

    def compute_lexicon_probability_matrix(self, src_batch):
        if self.lexical_probability_dictionary is not None:

            lexicon_probability_matrix = compute_lexicon_matrix(