# Path: nmt_chainer/utilities/utils.py
# def ortho_init(link):
#     if isinstance(link, chainer.links.Linear):
#         print("init ortho", link)
#         link.W.data[...] = gen_ortho(link.W.data.shape)
#     elif isinstance(link, chainer.links.GRU):
#         print("init ortho", link)
#         for name_lin in "W_r U_r W_z U_z W U".split(" "):
#             print("case", name_lin, getattr(link, name_lin))
#             ortho_init(getattr(link, name_lin))
#     elif isinstance(link, chainer.links.Maxout):
#         print("init ortho", link)
#         ortho_init(link.linear)
#     else:
#         raise NotImplemented


import logging
import chainer.functions as F
import chainer.links as L
import chainerx
from chainer import Chain, ChainList, Link, Variable, cuda
from nmt_chainer.utilities.utils import ortho_init

#try:
#    batch_matmul = F.batch_matmul
#except AttributeError: # for chainer>=3
def batch_matmul(a, b, transa=False, transb=False):
    return F.matmul(a[:, :, None], b, transa=transa, transb=transb)

class AttentionModule(Chain):
    """ Attention Module for computing the current context during decoding.
        The __call_ takes 2 parameters: fb_concat and mask.

        fb_concat should be the result of a call to Encoder.
        mask is as in the description of Encoder

        Return a chainer variable of shape (mb_size, Hi) and type float32
    """

    def __init__(self, Hi, Ha, Ho, init_orth=False, prev_word_embedding_size = None):
        super(AttentionModule, self).__init__(
            al_lin_h=L.Linear(Hi, Ha, nobias=False),
            al_lin_s=L.Linear(Ho, Ha, nobias=True),
            al_lin_o=L.Linear(Ha, 1, nobias=True)
        )
        self.Hi = Hi
        self.Ha = Ha

        if prev_word_embedding_size is not None:
            self.add_link("al_lin_y", L.Linear(prev_word_embedding_size, Ha))

        if init_orth:

# Path: nmt_chainer/utilities/utils.py
# def ortho_init(link):
#     if isinstance(link, chainer.links.Linear):
#         print("init ortho", link)
#         link.W.data[...] = gen_ortho(link.W.data.shape)
#     elif isinstance(link, chainer.links.GRU):
#         print("init ortho", link)
#         for name_lin in "W_r U_r W_z U_z W U".split(" "):
#             print("case", name_lin, getattr(link, name_lin))
#             ortho_init(getattr(link, name_lin))
#     elif isinstance(link, chainer.links.Maxout):
#         print("init ortho", link)
#         ortho_init(link.linear)
#     else:
#         raise NotImplemented


# Path: nmt_chainer/models/attention.py
import logging
import chainer.functions as F
import chainer.links as L
import chainerx
from chainer import Chain, ChainList, Link, Variable, cuda
from nmt_chainer.utilities.utils import ortho_init


#try:
#    batch_matmul = F.batch_matmul
#except AttributeError: # for chainer>=3
def batch_matmul(a, b, transa=False, transb=False):
    return F.matmul(a[:, :, None], b, transa=transa, transb=transb)

class AttentionModule(Chain):
    """ Attention Module for computing the current context during decoding.
        The __call_ takes 2 parameters: fb_concat and mask.

        fb_concat should be the result of a call to Encoder.
        mask is as in the description of Encoder

        Return a chainer variable of shape (mb_size, Hi) and type float32
    """

    def __init__(self, Hi, Ha, Ho, init_orth=False, prev_word_embedding_size = None):
        super(AttentionModule, self).__init__(
            al_lin_h=L.Linear(Hi, Ha, nobias=False),
            al_lin_s=L.Linear(Ho, Ha, nobias=True),
            al_lin_o=L.Linear(Ha, 1, nobias=True)
        )
        self.Hi = Hi
        self.Ha = Ha

        if prev_word_embedding_size is not None:
            self.add_link("al_lin_y", L.Linear(prev_word_embedding_size, Ha))

        if init_orth:

            ortho_init(self.al_lin_h)