#                               Apache License
#                         Version 2.0, January 2004
#                     http://www.apache.org/licenses/

# TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

# 1. Definitions.

#   "License" shall mean the terms and conditions for use, reproduction,
#   and distribution as defined by Sections 1 through 9 of this document.

#   "Licensor" shall mean the copyright owner or entity authorized by
#   the copyright owner that is granting the License.

#   "Legal Entity" shall mean the union of the acting entity and all
#   other entities that control, are controlled by, or are under common
#   control with that entity. For the purposes of this definition,
#   "control" means (i) the power, direct or indirect, to cause the
#   direction or management of such entity, whether by contract or
#   otherwise, or (ii) ownership of fifty percent (50%) or more of the
#   outstanding shares, or (iii) beneficial ownership of such entity.

#   "You" (or "Your") shall mean an individual or Legal Entity
#   exercising permissions granted by this License.

#   "Source" form shall mean the preferred form for making modifications,
#   including but not limited to software source code, documentation
#   source, and configuration files.

#   "Object" form shall mean any form resulting from mechanical
#   transformation or translation of a Source form, including but
#   not limited to compiled object code, generated documentation,
#   and conversions to other media types.

#   "Work" shall mean the work of authorship, whether in Source or
#   Object form, made available under the License, as indicated by a
#   copyright notice that is included in or attached to the work
#   (an example is provided in the Appendix below).

#   "Derivative Works" shall mean any work, whether in Source or Object
#   form, that is based on (or derived from) the Work and for which the
#   editorial revisions, annotations, elaborations, or other modifications
#   represent, as a whole, an original work of authorship. For the purposes
#   of this License, Derivative Works shall not include works that remain
#   separable from, or merely link (or bind by name) to the interfaces of,
#   the Work and Derivative Works thereof.

#   "Contribution" shall mean any work of authorship, including
#   the original version of the Work and any modifications or additions
#   to that Work or Derivative Works thereof, that is intentionally
#   submitted to Licensor for inclusion in the Work by the copyright owner
#   or by an individual or Legal Entity authorized to submit on behalf of
#   the copyright owner. For the purposes of this definition, "submitted"
#   means any form of electronic, verbal, or written communication sent
#   to the Licensor or its representatives, including but not limited to
#   communication on electronic mailing lists, source code control systems,
#   and issue tracking systems that are managed by, or on behalf of, the
#   Licensor for the purpose of discussing and improving the Work, but
#   excluding communication that is conspicuously marked or otherwise
#   designated in writing by the copyright owner as "Not a Contribution."

#   "Contributor" shall mean Licensor and any individual or Legal Entity
#   on behalf of whom a Contribution has been received by Licensor and
#   subsequently incorporated within the Work.

# 2. Grant of Copyright License. Subject to the terms and conditions of
#   this License, each Contributor hereby grants to You a perpetual,
#   worldwide, non-exclusive, no-charge, royalty-free, irrevocable
#   copyright license to reproduce, prepare Derivative Works of,
#   publicly display, publicly perform, sublicense, and distribute the
#   Work and such Derivative Works in Source or Object form.

# 3. Grant of Patent License. Subject to the terms and conditions of
#   this License, each Contributor hereby grants to You a perpetual,
#   worldwide, non-exclusive, no-charge, royalty-free, irrevocable
#   (except as stated in this section) patent license to make, have made,
#   use, offer to sell, sell, import, and otherwise transfer the Work,
#   where such license applies only to those patent claims licensable
#   by such Contributor that are necessarily infringed by their
#   Contribution(s) alone or by combination of their Contribution(s)
#   with the Work to which such Contribution(s) was submitted. If You
#   institute patent litigation against any entity (including a
#   cross-claim or counterclaim in a lawsuit) alleging that the Work
#   or a Contribution incorporated within the Work constitutes direct
#   or contributory patent infringement, then any patent licenses
#   granted to You under this License for that Work shall terminate
#   as of the date such litigation is filed.

# 4. Redistribution. You may reproduce and distribute copies of the
#   Work or Derivative Works thereof in any medium, with or without
#   modifications, and in Source or Object form, provided that You
#   meet the following conditions:

#   (a) You must give any other recipients of the Work or
#       Derivative Works a copy of this License; and

#   (b) You must cause any modified files to carry prominent notices
#       stating that You changed the files; and

#   (c) You must retain, in the Source form of any Derivative Works
#       that You distribute, all copyright, patent, trademark, and
#       attribution notices from the Source form of the Work,
#       excluding those notices that do not pertain to any part of
#       the Derivative Works; and

#   (d) If the Work includes a "NOTICE" text file as part of its
#       distribution, then any Derivative Works that You distribute must
#       include a readable copy of the attribution notices contained
#       within such NOTICE file, excluding those notices that do not
#       pertain to any part of the Derivative Works, in at least one
#       of the following places: within a NOTICE text file distributed
#       as part of the Derivative Works; within the Source form or
#       documentation, if provided along with the Derivative Works; or,
#       within a display generated by the Derivative Works, if and
#       wherever such third-party notices normally appear. The contents
#       of the NOTICE file are for informational purposes only and
#       do not modify the License. You may add Your own attribution
#       notices within Derivative Works that You distribute, alongside
#       or as an addendum to the NOTICE text from the Work, provided
#       that such additional attribution notices cannot be construed
#       as modifying the License.

#   You may add Your own copyright statement to Your modifications and
#   may provide additional or different license terms and conditions
#   for use, reproduction, or distribution of Your modifications, or
#   for any such Derivative Works as a whole, provided Your use,
#   reproduction, and distribution of the Work otherwise complies with
#   the conditions stated in this License.

# 5. Submission of Contributions. Unless You explicitly state otherwise,
#   any Contribution intentionally submitted for inclusion in the Work
#   by You to the Licensor shall be under the terms and conditions of
#   this License, without any additional terms or conditions.
#   Notwithstanding the above, nothing herein shall supersede or modify
#   the terms of any separate license agreement you may have executed
#   with Licensor regarding such Contributions.

# 6. Trademarks. This License does not grant permission to use the trade
#   names, trademarks, service marks, or product names of the Licensor,
#   except as required for reasonable and customary use in describing the
#   origin of the Work and reproducing the content of the NOTICE file.

# 7. Disclaimer of Warranty. Unless required by applicable law or
#   agreed to in writing, Licensor provides the Work (and each
#   Contributor provides its Contributions) on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
#   implied, including, without limitation, any warranties or conditions
#   of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
#   PARTICULAR PURPOSE. You are solely responsible for determining the
#   appropriateness of using or redistributing the Work and assume any
#   risks associated with Your exercise of permissions under this License.

# 8. Limitation of Liability. In no event and under no legal theory,
#   whether in tort (including negligence), contract, or otherwise,
#   unless required by applicable law (such as deliberate and grossly
#   negligent acts) or agreed to in writing, shall any Contributor be
#   liable to You for damages, including any direct, indirect, special,
#   incidental, or consequential damages of any character arising as a
#   result of this License or out of the use or inability to use the
#   Work (including but not limited to damages for loss of goodwill,
#   work stoppage, computer failure or malfunction, or any and all
#   other commercial damages or losses), even if such Contributor
#   has been advised of the possibility of such damages.

# 9. Accepting Warranty or Additional Liability. While redistributing
#   the Work or Derivative Works thereof, You may choose to offer,
#   and charge a fee for, acceptance of support, warranty, indemnity,
#   or other liability obligations and/or rights consistent with this
#   License. However, in accepting such obligations, You may act only
#   on Your own behalf and on Your sole responsibility, not on behalf
#   of any other Contributor, and only if You agree to indemnify,
#   defend, and hold each Contributor harmless for any liability
#   incurred by, or claims asserted against, such Contributor by reason
#   of your accepting any such warranty or additional liability.

# END OF TERMS AND CONDITIONS

# Copyright 2017- IBM Corporation

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn


class Attention(nn.Module):
    """ Applies attention mechanism on the `context` using the `query`.

    Args:
        dimensions (int): Dimensionality of the query and context.
        attention_type (str, optional): How to compute the attention score:

            * dot: :math:`score(H_j,q) = H_j^T q`
            * general: :math:`score(H_j, q) = H_j^T W_a q`


    Examples:

         >>> attention = Attention(256)
         >>> query = Variable(torch.randn(5, 1, 256))
         >>> context = Variable(torch.randn(5, 5, 256))
         >>> output, weights = attention(query, context)
         >>> output.size()
         torch.Size([5, 1, 256])
         >>> weights.size()
         torch.Size([5, 1, 5])
    """

    def __init__(self, dimensions, attention_type='general'):
        super(Attention, self).__init__()

        if attention_type not in ['dot', 'general']:
            raise ValueError('Invalid attention type selected.')

        self.attention_type = attention_type
        if self.attention_type == 'general':
            self.linear_in = nn.Linear(dimensions, dimensions, bias=False)

        self.linear_out = nn.Linear(dimensions * 2, dimensions, bias=False)
        self.softmax = nn.Softmax(dim=0)
        self.tanh = nn.Tanh()

    def forward(self, query, context):
        """
        Args:
            query (:class:`torch.FloatTensor` [batch size, output length, dimensions]): Sequence of
                queries to query the context.
            context (:class:`torch.FloatTensor` [batch size, query length, dimensions]): Data
                overwhich to apply the attention mechanism.
        Returns:
            :class:`tuple` with `output` and `weights`:
            * **output** (:class:`torch.LongTensor` [batch size, output length, dimensions]):
              Tensor containing the attended features.
            * **weights** (:class:`torch.FloatTensor` [batch size, output length, query length]):
              Tensor containing attention weights.
        """
        batch_size, output_len, dimensions = query.size()
        query_len = context.size(1)

        if self.attention_type == "general":
            query = query.view(batch_size * output_len, dimensions)
            query = self.linear_in(query)
            query = query.view(batch_size, output_len, dimensions)

        # TODO: Include mask on PADDING_INDEX?

        # (batch_size, output_len, dimensions) * (batch_size, query_len, dimensions) ->
        # (batch_size, output_len, query_len)
        attention_scores = torch.bmm(query, context.transpose(1, 2).contiguous())

        # Compute weights across every context sequence
        attention_scores = attention_scores.view(batch_size * output_len, query_len)
        attention_weights = self.softmax(attention_scores)
        attention_weights = attention_weights.view(batch_size, output_len, query_len)

        # (batch_size, output_len, query_len) * (batch_size, query_len, dimensions) ->
        # (batch_size, output_len, dimensions)
        mix = torch.bmm(attention_weights, context)

        # concat -> (batch_size * output_len, 2*dimensions)
        combined = torch.cat((mix, query), dim=2)
        combined = combined.view(batch_size * output_len, 2 * dimensions)

        # Apply linear_out on every 2nd dimension of concat
        # output -> (batch_size, output_len, dimensions)
        output = self.linear_out(combined).view(batch_size, output_len, dimensions)
        output = self.tanh(output)

        return output, attention_weights
