# Copyright (c) 2020 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.

import math
import numpy as np

import paddle.fluid as fluid
import paddle.fluid.layers as layers

from paddlerec.core.utils import envs
from paddlerec.core.model import Model as ModelBase


class Model(ModelBase):
    def __init__(self, config):
        ModelBase.__init__(self, config)
        self.init_config()

    def init_config(self):
        self._fetch_interval = 1
        self.items_num, self.ins_num = self.config_read(
            envs.get_global_env("hyper_parameters.config_path", None,
                                self._namespace))
        self.train_batch_size = envs.get_global_env("batch_size", None,
                                                    "train.reader")
        self.evaluate_batch_size = envs.get_global_env("batch_size", None,
                                                       "evaluate.reader")
        self.hidden_size = envs.get_global_env(
            "hyper_parameters.sparse_feature_dim", None, self._namespace)
        self.step = envs.get_global_env(
            "hyper_parameters.gnn_propogation_steps", None, self._namespace)

    def config_read(self, config_path=None):
        if config_path is None:
            raise ValueError(
                "please set train.model.hyper_parameters.config_path at first")
        with open(config_path, "r") as fin:
            item_nums = int(fin.readline().strip())
            ins_nums = int(fin.readline().strip())
        return item_nums, ins_nums

    def input(self, bs):
        self.items = fluid.data(
            name="items", shape=[bs, -1],
            dtype="int64")  # [batch_size, uniq_max]
        self.seq_index = fluid.data(
            name="seq_index", shape=[bs, -1, 2],
            dtype="int32")  # [batch_size, seq_max, 2]
        self.last_index = fluid.data(
            name="last_index", shape=[bs, 2], dtype="int32")  # [batch_size, 2]
        self.adj_in = fluid.data(
            name="adj_in", shape=[bs, -1, -1],
            dtype="float32")  # [batch_size, seq_max, seq_max]
        self.adj_out = fluid.data(
            name="adj_out", shape=[bs, -1, -1],
            dtype="float32")  # [batch_size, seq_max, seq_max]
        self.mask = fluid.data(
            name="mask", shape=[bs, -1, 1],
            dtype="float32")  # [batch_size, seq_max, 1]
        self.label = fluid.data(
            name="label", shape=[bs, 1], dtype="int64")  # [batch_size, 1]

        res = [
            self.items, self.seq_index, self.last_index, self.adj_in,
            self.adj_out, self.mask, self.label
        ]
        return res

    def train_input(self):
        res = self.input(self.train_batch_size)
        self._data_var = res

        use_dataloader = envs.get_global_env("hyper_parameters.use_DataLoader",
                                             False, self._namespace)

        if self._platform != "LINUX" or use_dataloader:
            self._data_loader = fluid.io.DataLoader.from_generator(
                feed_list=self._data_var,
                capacity=256,
                use_double_buffer=False,
                iterable=False)

    def net(self, items_num, hidden_size, step, bs):
        stdv = 1.0 / math.sqrt(hidden_size)

        def embedding_layer(input,
                            table_name,
                            emb_dim,
                            initializer_instance=None):
            emb = fluid.embedding(
                input=input,
                size=[items_num, emb_dim],
                param_attr=fluid.ParamAttr(
                    name=table_name, initializer=initializer_instance), )
            return emb

        sparse_initializer = fluid.initializer.Uniform(low=-stdv, high=stdv)
        items_emb = embedding_layer(self.items, "emb", hidden_size,
                                    sparse_initializer)
        pre_state = items_emb
        for i in range(step):
            pre_state = layers.reshape(
                x=pre_state, shape=[bs, -1, hidden_size])
            state_in = layers.fc(
                input=pre_state,
                name="state_in",
                size=hidden_size,
                act=None,
                num_flatten_dims=2,
                param_attr=fluid.ParamAttr(
                    initializer=fluid.initializer.Uniform(
                        low=-stdv, high=stdv)),
                bias_attr=fluid.ParamAttr(
                    initializer=fluid.initializer.Uniform(
                        low=-stdv, high=stdv)))  # [batch_size, uniq_max, h]
            state_out = layers.fc(
                input=pre_state,
                name="state_out",
                size=hidden_size,
                act=None,
                num_flatten_dims=2,
                param_attr=fluid.ParamAttr(
                    initializer=fluid.initializer.Uniform(
                        low=-stdv, high=stdv)),
                bias_attr=fluid.ParamAttr(
                    initializer=fluid.initializer.Uniform(
                        low=-stdv, high=stdv)))  # [batch_size, uniq_max, h]

            state_adj_in = layers.matmul(self.adj_in,
                                         state_in)  # [batch_size, uniq_max, h]
            state_adj_out = layers.matmul(
                self.adj_out, state_out)  # [batch_size, uniq_max, h]

            gru_input = layers.concat([state_adj_in, state_adj_out], axis=2)

            gru_input = layers.reshape(
                x=gru_input, shape=[-1, hidden_size * 2])
            gru_fc = layers.fc(input=gru_input,
                               name="gru_fc",
                               size=3 * hidden_size,
                               bias_attr=False)
            pre_state, _, _ = fluid.layers.gru_unit(
                input=gru_fc,
                hidden=layers.reshape(
                    x=pre_state, shape=[-1, hidden_size]),
                size=3 * hidden_size)

        final_state = layers.reshape(pre_state, shape=[bs, -1, hidden_size])
        seq = layers.gather_nd(final_state, self.seq_index)
        last = layers.gather_nd(final_state, self.last_index)

        seq_fc = layers.fc(
            input=seq,
            name="seq_fc",
            size=hidden_size,
            bias_attr=False,
            act=None,
            num_flatten_dims=2,
            param_attr=fluid.ParamAttr(initializer=fluid.initializer.Uniform(
                low=-stdv, high=stdv)))  # [batch_size, seq_max, h]
        last_fc = layers.fc(input=last,
                            name="last_fc",
                            size=hidden_size,
                            bias_attr=False,
                            act=None,
                            num_flatten_dims=1,
                            param_attr=fluid.ParamAttr(
                                initializer=fluid.initializer.Uniform(
                                    low=-stdv, high=stdv)))  # [bathc_size, h]

        seq_fc_t = layers.transpose(
            seq_fc, perm=[1, 0, 2])  # [seq_max, batch_size, h]
        add = layers.elementwise_add(seq_fc_t,
                                     last_fc)  # [seq_max, batch_size, h]
        b = layers.create_parameter(
            shape=[hidden_size],
            dtype='float32',
            default_initializer=fluid.initializer.Constant(value=0.0))  # [h]
        add = layers.elementwise_add(add, b)  # [seq_max, batch_size, h]

        add_sigmoid = layers.sigmoid(add)  # [seq_max, batch_size, h]
        add_sigmoid = layers.transpose(
            add_sigmoid, perm=[1, 0, 2])  # [batch_size, seq_max, h]

        weight = layers.fc(
            input=add_sigmoid,
            name="weight_fc",
            size=1,
            act=None,
            num_flatten_dims=2,
            bias_attr=False,
            param_attr=fluid.ParamAttr(initializer=fluid.initializer.Uniform(
                low=-stdv, high=stdv)))  # [batch_size, seq_max, 1]
        weight *= self.mask
        weight_mask = layers.elementwise_mul(
            seq, weight, axis=0)  # [batch_size, seq_max, h]
        global_attention = layers.reduce_sum(
            weight_mask, dim=1)  # [batch_size, h]

        final_attention = layers.concat(
            [global_attention, last], axis=1)  # [batch_size, 2*h]
        final_attention_fc = layers.fc(
            input=final_attention,
            name="final_attention_fc",
            size=hidden_size,
            bias_attr=False,
            act=None,
            param_attr=fluid.ParamAttr(initializer=fluid.initializer.Uniform(
                low=-stdv, high=stdv)))  # [batch_size, h]

        # all_vocab = layers.create_global_var(
        #     shape=[items_num - 1],
        #     value=0,
        #     dtype="int64",
        #     persistable=True,
        #     name="all_vocab")
        all_vocab = np.arange(1, items_num).reshape((-1)).astype('int32')
        all_vocab = fluid.layers.cast(
            x=fluid.layers.assign(all_vocab), dtype='int64')

        all_emb = fluid.embedding(
            input=all_vocab,
            param_attr=fluid.ParamAttr(
                name="emb",
                initializer=fluid.initializer.Uniform(
                    low=-stdv, high=stdv)),
            size=[items_num, hidden_size])  # [all_vocab, h]

        logits = layers.matmul(
            x=final_attention_fc, y=all_emb,
            transpose_y=True)  # [batch_size, all_vocab]
        softmax = layers.softmax_with_cross_entropy(
            logits=logits, label=self.label)  # [batch_size, 1]
        self.loss = layers.reduce_mean(softmax)  # [1]
        self.acc = layers.accuracy(input=logits, label=self.label, k=20)

    def avg_loss(self):
        self._cost = self.loss

    def metrics(self):
        self._metrics["LOSS"] = self.loss
        self._metrics["train_acc"] = self.acc

    def train_net(self):
        self.train_input()
        self.net(self.items_num, self.hidden_size, self.step,
                 self.train_batch_size)
        self.avg_loss()
        self.metrics()

    def optimizer(self):
        learning_rate = envs.get_global_env("hyper_parameters.learning_rate",
                                            None, self._namespace)
        step_per_epoch = self.ins_num // self.train_batch_size
        decay_steps = envs.get_global_env("hyper_parameters.decay_steps", None,
                                          self._namespace)
        decay_rate = envs.get_global_env("hyper_parameters.decay_rate", None,
                                         self._namespace)
        l2 = envs.get_global_env("hyper_parameters.l2", None, self._namespace)
        optimizer = fluid.optimizer.Adam(
            learning_rate=fluid.layers.exponential_decay(
                learning_rate=learning_rate,
                decay_steps=decay_steps * step_per_epoch,
                decay_rate=decay_rate),
            regularization=fluid.regularizer.L2DecayRegularizer(
                regularization_coeff=l2))

        return optimizer

    def infer_input(self):
        self._reader_namespace = "evaluate.reader"
        res = self.input(self.evaluate_batch_size)
        self._infer_data_var = res

        self._infer_data_loader = fluid.io.DataLoader.from_generator(
            feed_list=self._infer_data_var,
            capacity=64,
            use_double_buffer=False,
            iterable=False)

    def infer_net(self):
        self.infer_input()
        self.net(self.items_num, self.hidden_size, self.step,
                 self.evaluate_batch_size)
        self._infer_results['acc'] = self.acc
        self._infer_results['loss'] = self.loss
