# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# 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.
"""
neural network for word2vec
"""
from __future__ import print_function
import paddle
import paddle.nn.functional as F


def skip_gram_word2vec(dict_size,
                       embedding_size,
                       batch_size,
                       is_sparse=False,
                       neg_num=5):

    words = []
    input_word = paddle.static.data(
        name="input_word", shape=[None, 1], dtype='int64')
    true_word = paddle.static.data(
        name='true_label', shape=[None, 1], dtype='int64')
    neg_word = paddle.static.data(
        name="neg_label", shape=[None, neg_num], dtype='int64')

    words.append(input_word)
    words.append(true_word)
    words.append(neg_word)

    py_reader = paddle.io.DataLoader.from_generator(
        capacity=64, feed_list=words, use_double_buffer=True, iterable=False)

    words[0] = paddle.reshape(words[0], [-1])
    words[1] = paddle.reshape(words[1], [-1])
    init_width = 0.5 / embedding_size
    input_emb = paddle.static.nn.embedding(
        input=words[0],
        is_sparse=is_sparse,
        size=[dict_size, embedding_size],
        param_attr=paddle.ParamAttr(
            name='emb',
            initializer=paddle.nn.initializer.Uniform(-init_width, init_width)))

    true_emb_w = paddle.static.nn.embedding(
        input=words[1],
        is_sparse=is_sparse,
        size=[dict_size, embedding_size],
        param_attr=paddle.ParamAttr(
            name='emb_w',
            initializer=paddle.nn.initializer.Constant(value=0.0)))

    true_emb_b = paddle.static.nn.embedding(
        input=words[1],
        is_sparse=is_sparse,
        size=[dict_size, 1],
        param_attr=paddle.ParamAttr(
            name='emb_b',
            initializer=paddle.nn.initializer.Constant(value=0.0)))
    neg_word_reshape = paddle.reshape(words[2], shape=[-1])
    neg_word_reshape.stop_gradient = True

    neg_emb_w = paddle.static.nn.embedding(
        input=neg_word_reshape,
        is_sparse=is_sparse,
        size=[dict_size, embedding_size],
        param_attr=paddle.ParamAttr(name='emb_w', learning_rate=1.0))

    neg_emb_w_re = paddle.reshape(
        neg_emb_w, shape=[-1, neg_num, embedding_size])
    neg_emb_b = paddle.static.nn.embedding(
        input=neg_word_reshape,
        is_sparse=is_sparse,
        size=[dict_size, 1],
        param_attr=paddle.ParamAttr(name='emb_b', learning_rate=1.0))

    neg_emb_b_vec = paddle.reshape(neg_emb_b, shape=[-1, neg_num])
    true_logits = paddle.add(
        paddle.mean(paddle.multiply(input_emb, true_emb_w), keepdim=True),
        true_emb_b)
    input_emb_re = paddle.reshape(input_emb, shape=[-1, 1, embedding_size])
    neg_matmul = paddle.matmul(input_emb_re, neg_emb_w_re, transpose_y=True)
    neg_matmul_re = paddle.reshape(neg_matmul, shape=[-1, neg_num])
    neg_logits = paddle.add(neg_matmul_re, neg_emb_b_vec)
    #nce loss

    label_ones = paddle.full(
        shape=[batch_size, 1], fill_value=1.0, dtype='float32')
    label_zeros = paddle.full(
        shape=[batch_size, neg_num], fill_value=0.0, dtype='float32')

    true_xent = F.binary_cross_entropy_with_logits(
        true_logits, label_ones, reduction='none')
    neg_xent = F.binary_cross_entropy_with_logits(
        neg_logits, label_zeros, reduction='none')
    cost = paddle.add(
        paddle.sum(true_xent, axis=1), paddle.sum(neg_xent, axis=1))
    avg_cost = paddle.mean(cost)
    return avg_cost, py_reader


def infer_network(vocab_size, emb_size):
    analogy_a = paddle.static.data(
        name="analogy_a", shape=[None, 1], dtype='int64')
    analogy_b = paddle.static.data(
        name="analogy_b", shape=[None, 1], dtype='int64')
    analogy_c = paddle.static.data(
        name="analogy_c", shape=[None, 1], dtype='int64')
    all_label = paddle.static.data(
        name="all_label", shape=[vocab_size, 1], dtype='int64')
    all_label = paddle.reshape(all_label, [-1])
    emb_all_label = paddle.static.nn.embedding(
        input=all_label, size=[vocab_size, emb_size], param_attr="emb")

    analogy_a = paddle.reshape(analogy_a, [-1])
    emb_a = paddle.static.nn.embedding(
        input=analogy_a, size=[vocab_size, emb_size], param_attr="emb")
    analogy_b = paddle.reshape(analogy_b, [-1])
    emb_b = paddle.static.nn.embedding(
        input=analogy_b, size=[vocab_size, emb_size], param_attr="emb")
    analogy_c = paddle.reshape(analogy_c, [-1])
    emb_c = paddle.static.nn.embedding(
        input=analogy_c, size=[vocab_size, emb_size], param_attr="emb")
    target = paddle.add(paddle.add(emb_b, -emb_a), emb_c)
    emb_all_label_l2 = F.normalize(emb_all_label, p=2, axis=1)
    dist = paddle.matmul(x=target, y=emb_all_label_l2, transpose_y=True)
    values, pred_idx = paddle.topk(x=dist, k=4)
    return values, pred_idx