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

from __future__ import print_function

import unittest
import numpy

import paddle.fluid as fluid
import paddle.fluid.layers as layers
import paddle.fluid.core as core
from paddle.fluid.framework import program_guard, Program

from paddle.fluid.executor import Executor
from paddle.fluid import framework

from paddle.fluid.layers.rnn import LSTMCell, GRUCell, RNNCell
from paddle.fluid.layers import rnn as dynamic_rnn
from paddle.fluid import contrib
from paddle.fluid.contrib.layers import basic_lstm
import paddle.fluid.layers.utils as utils

import numpy as np


class TestLSTMCellError(unittest.TestCase):

    def test_errors(self):
        with program_guard(Program(), Program()):
            batch_size, input_size, hidden_size = 4, 16, 16
            inputs = fluid.data(name='inputs',
                                shape=[None, input_size],
                                dtype='float32')
            pre_hidden = fluid.data(name='pre_hidden',
                                    shape=[None, hidden_size],
                                    dtype='float32')
            pre_cell = fluid.data(name='pre_cell',
                                  shape=[None, hidden_size],
                                  dtype='float32')
            cell = LSTMCell(hidden_size)

            def test_input_Variable():
                np_input = np.random.random(
                    (batch_size, input_size)).astype("float32")
                cell(np_input, [pre_hidden, pre_cell])

            self.assertRaises(TypeError, test_input_Variable)

            def test_pre_hidden_Variable():
                np_pre_hidden = np.random.random(
                    (batch_size, hidden_size)).astype("float32")
                cell(inputs, [np_pre_hidden, pre_cell])

            self.assertRaises(TypeError, test_pre_hidden_Variable)

            def test_pre_cell_Variable():
                np_pre_cell = np.random.random(
                    (batch_size, input_size)).astype("float32")
                cell(inputs, [pre_hidden, np_pre_cell])

            self.assertRaises(TypeError, test_pre_cell_Variable)

            def test_input_type():
                error_inputs = fluid.data(name='error_inputs',
                                          shape=[None, input_size],
                                          dtype='int32')
                cell(error_inputs, [pre_hidden, pre_cell])

            self.assertRaises(TypeError, test_input_type)

            def test_pre_hidden_type():
                error_pre_hidden = fluid.data(name='error_pre_hidden',
                                              shape=[None, hidden_size],
                                              dtype='int32')
                cell(inputs, [error_pre_hidden, pre_cell])

            self.assertRaises(TypeError, test_pre_hidden_type)

            def test_pre_cell_type():
                error_pre_cell = fluid.data(name='error_pre_cell',
                                            shape=[None, hidden_size],
                                            dtype='int32')
                cell(inputs, [pre_hidden, error_pre_cell])

            self.assertRaises(TypeError, test_pre_cell_type)

            def test_dtype():
                # the input type must be Variable
                LSTMCell(hidden_size, dtype="int32")

            self.assertRaises(TypeError, test_dtype)


class TestLSTMCell(unittest.TestCase):

    def setUp(self):
        self.batch_size = 4
        self.input_size = 16
        self.hidden_size = 16

    def test_run(self):
        inputs = fluid.data(name='inputs',
                            shape=[None, self.input_size],
                            dtype='float32')
        pre_hidden = fluid.data(name='pre_hidden',
                                shape=[None, self.hidden_size],
                                dtype='float32')
        pre_cell = fluid.data(name='pre_cell',
                              shape=[None, self.hidden_size],
                              dtype='float32')

        cell = LSTMCell(self.hidden_size)
        lstm_hidden_new, lstm_states_new = cell(inputs, [pre_hidden, pre_cell])

        lstm_unit = contrib.layers.rnn_impl.BasicLSTMUnit(
            "basicLSTM", self.hidden_size, None, None, None, None, 1.0,
            "float32")
        lstm_hidden, lstm_cell = lstm_unit(inputs, pre_hidden, pre_cell)

        if core.is_compiled_with_cuda():
            place = core.CUDAPlace(0)
        else:
            place = core.CPUPlace()
        exe = Executor(place)
        exe.run(framework.default_startup_program())

        inputs_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.input_size)).astype('float32')
        pre_hidden_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.hidden_size)).astype('float32')
        pre_cell_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.hidden_size)).astype('float32')

        param_names = [[
            "LSTMCell/BasicLSTMUnit_0.w_0", "basicLSTM/BasicLSTMUnit_0.w_0"
        ], ["LSTMCell/BasicLSTMUnit_0.b_0", "basicLSTM/BasicLSTMUnit_0.b_0"]]

        for names in param_names:
            param = np.array(fluid.global_scope().find_var(
                names[0]).get_tensor())
            param = np.random.uniform(-0.1, 0.1,
                                      size=param.shape).astype('float32')
            fluid.global_scope().find_var(names[0]).get_tensor().set(
                param, place)
            fluid.global_scope().find_var(names[1]).get_tensor().set(
                param, place)

        out = exe.run(feed={
            'inputs': inputs_np,
            'pre_hidden': pre_hidden_np,
            'pre_cell': pre_cell_np
        },
                      fetch_list=[lstm_hidden_new, lstm_hidden])

        self.assertTrue(np.allclose(out[0], out[1], rtol=1e-4, atol=0))


class TestGRUCellError(unittest.TestCase):

    def test_errors(self):
        with program_guard(Program(), Program()):
            batch_size, input_size, hidden_size = 4, 16, 16
            inputs = fluid.data(name='inputs',
                                shape=[None, input_size],
                                dtype='float32')
            pre_hidden = layers.data(name='pre_hidden',
                                     shape=[None, hidden_size],
                                     append_batch_size=False,
                                     dtype='float32')
            cell = GRUCell(hidden_size)

            def test_input_Variable():
                np_input = np.random.random(
                    (batch_size, input_size)).astype("float32")
                cell(np_input, pre_hidden)

            self.assertRaises(TypeError, test_input_Variable)

            def test_pre_hidden_Variable():
                np_pre_hidden = np.random.random(
                    (batch_size, hidden_size)).astype("float32")
                cell(inputs, np_pre_hidden)

            self.assertRaises(TypeError, test_pre_hidden_Variable)

            def test_input_type():
                error_inputs = fluid.data(name='error_inputs',
                                          shape=[None, input_size],
                                          dtype='int32')
                cell(error_inputs, pre_hidden)

            self.assertRaises(TypeError, test_input_type)

            def test_pre_hidden_type():
                error_pre_hidden = fluid.data(name='error_pre_hidden',
                                              shape=[None, hidden_size],
                                              dtype='int32')
                cell(inputs, error_pre_hidden)

            self.assertRaises(TypeError, test_pre_hidden_type)

            def test_dtype():
                # the input type must be Variable
                GRUCell(hidden_size, dtype="int32")

            self.assertRaises(TypeError, test_dtype)


class TestGRUCell(unittest.TestCase):

    def setUp(self):
        self.batch_size = 4
        self.input_size = 16
        self.hidden_size = 16

    def test_run(self):
        inputs = fluid.data(name='inputs',
                            shape=[None, self.input_size],
                            dtype='float32')
        pre_hidden = layers.data(name='pre_hidden',
                                 shape=[None, self.hidden_size],
                                 append_batch_size=False,
                                 dtype='float32')

        cell = GRUCell(self.hidden_size)
        gru_hidden_new, _ = cell(inputs, pre_hidden)

        gru_unit = contrib.layers.rnn_impl.BasicGRUUnit("basicGRU",
                                                        self.hidden_size, None,
                                                        None, None, None,
                                                        "float32")
        gru_hidden = gru_unit(inputs, pre_hidden)

        if core.is_compiled_with_cuda():
            place = core.CUDAPlace(0)
        else:
            place = core.CPUPlace()
        exe = Executor(place)
        exe.run(framework.default_startup_program())

        inputs_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.input_size)).astype('float32')
        pre_hidden_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.hidden_size)).astype('float32')

        param_names = [
            ["GRUCell/BasicGRUUnit_0.w_0", "basicGRU/BasicGRUUnit_0.w_0"],
            ["GRUCell/BasicGRUUnit_0.w_1", "basicGRU/BasicGRUUnit_0.w_1"],
            ["GRUCell/BasicGRUUnit_0.b_0", "basicGRU/BasicGRUUnit_0.b_0"],
            ["GRUCell/BasicGRUUnit_0.b_1", "basicGRU/BasicGRUUnit_0.b_1"]
        ]

        for names in param_names:
            param = np.array(fluid.global_scope().find_var(
                names[0]).get_tensor())
            param = np.random.uniform(-0.1, 0.1,
                                      size=param.shape).astype('float32')
            fluid.global_scope().find_var(names[0]).get_tensor().set(
                param, place)
            fluid.global_scope().find_var(names[1]).get_tensor().set(
                param, place)

        out = exe.run(feed={
            'inputs': inputs_np,
            'pre_hidden': pre_hidden_np
        },
                      fetch_list=[gru_hidden_new, gru_hidden])

        self.assertTrue(np.allclose(out[0], out[1], rtol=1e-4, atol=0))


class TestRnnError(unittest.TestCase):

    def test_errors(self):
        with program_guard(Program(), Program()):
            batch_size = 4
            input_size = 16
            hidden_size = 16
            seq_len = 4
            inputs = fluid.data(name='inputs',
                                shape=[None, input_size],
                                dtype='float32')
            pre_hidden = layers.data(name='pre_hidden',
                                     shape=[None, hidden_size],
                                     append_batch_size=False,
                                     dtype='float32')
            inputs_basic_lstm = fluid.data(name='inputs_basic_lstm',
                                           shape=[None, None, input_size],
                                           dtype='float32')
            sequence_length = fluid.data(name="sequence_length",
                                         shape=[None],
                                         dtype='int64')

            inputs_dynamic_rnn = layers.transpose(inputs_basic_lstm,
                                                  perm=[1, 0, 2])
            cell = LSTMCell(hidden_size, name="LSTMCell_for_rnn")
            np_inputs_dynamic_rnn = np.random.random(
                (seq_len, batch_size, input_size)).astype("float32")

            def test_input_Variable():
                dynamic_rnn(cell=cell,
                            inputs=np_inputs_dynamic_rnn,
                            sequence_length=sequence_length,
                            is_reverse=False)

            self.assertRaises(TypeError, test_input_Variable)

            def test_input_list():
                dynamic_rnn(cell=cell,
                            inputs=[np_inputs_dynamic_rnn],
                            sequence_length=sequence_length,
                            is_reverse=False)

            self.assertRaises(TypeError, test_input_list)

            def test_initial_states_type():
                cell = GRUCell(hidden_size, name="GRUCell_for_rnn")
                error_initial_states = np.random.random(
                    (batch_size, hidden_size)).astype("float32")
                dynamic_rnn(cell=cell,
                            inputs=inputs_dynamic_rnn,
                            initial_states=error_initial_states,
                            sequence_length=sequence_length,
                            is_reverse=False)

            self.assertRaises(TypeError, test_initial_states_type)

            def test_initial_states_list():
                error_initial_states = [
                    np.random.random(
                        (batch_size, hidden_size)).astype("float32"),
                    np.random.random(
                        (batch_size, hidden_size)).astype("float32")
                ]
                dynamic_rnn(cell=cell,
                            inputs=inputs_dynamic_rnn,
                            initial_states=error_initial_states,
                            sequence_length=sequence_length,
                            is_reverse=False)

            self.assertRaises(TypeError, test_initial_states_type)

            def test_sequence_length_type():
                np_sequence_length = np.random.random(
                    (batch_size)).astype("float32")
                dynamic_rnn(cell=cell,
                            inputs=inputs_dynamic_rnn,
                            sequence_length=np_sequence_length,
                            is_reverse=False)

            self.assertRaises(TypeError, test_sequence_length_type)


class TestRnn(unittest.TestCase):

    def setUp(self):
        self.batch_size = 4
        self.input_size = 16
        self.hidden_size = 16
        self.seq_len = 4

    def test_run(self):
        inputs_basic_lstm = fluid.data(name='inputs_basic_lstm',
                                       shape=[None, None, self.input_size],
                                       dtype='float32')
        sequence_length = fluid.data(name="sequence_length",
                                     shape=[None],
                                     dtype='int64')

        inputs_dynamic_rnn = layers.transpose(inputs_basic_lstm, perm=[1, 0, 2])
        cell = LSTMCell(self.hidden_size, name="LSTMCell_for_rnn")
        output, final_state = dynamic_rnn(cell=cell,
                                          inputs=inputs_dynamic_rnn,
                                          sequence_length=sequence_length,
                                          is_reverse=False)
        output_new = layers.transpose(output, perm=[1, 0, 2])

        rnn_out, last_hidden, last_cell = basic_lstm(inputs_basic_lstm, None, None, self.hidden_size, num_layers=1, \
                batch_first = False, bidirectional=False, sequence_length=sequence_length, forget_bias = 1.0)

        if core.is_compiled_with_cuda():
            place = core.CUDAPlace(0)
        else:
            place = core.CPUPlace()
        exe = Executor(place)
        exe.run(framework.default_startup_program())

        inputs_basic_lstm_np = np.random.uniform(
            -0.1, 0.1,
            (self.seq_len, self.batch_size, self.input_size)).astype('float32')
        sequence_length_np = np.ones(self.batch_size,
                                     dtype='int64') * self.seq_len

        inputs_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.input_size)).astype('float32')
        pre_hidden_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.hidden_size)).astype('float32')
        pre_cell_np = np.random.uniform(
            -0.1, 0.1, (self.batch_size, self.hidden_size)).astype('float32')

        param_names = [[
            "LSTMCell_for_rnn/BasicLSTMUnit_0.w_0",
            "basic_lstm_layers_0/BasicLSTMUnit_0.w_0"
        ],
                       [
                           "LSTMCell_for_rnn/BasicLSTMUnit_0.b_0",
                           "basic_lstm_layers_0/BasicLSTMUnit_0.b_0"
                       ]]

        for names in param_names:
            param = np.array(fluid.global_scope().find_var(
                names[0]).get_tensor())
            param = np.random.uniform(-0.1, 0.1,
                                      size=param.shape).astype('float32')
            fluid.global_scope().find_var(names[0]).get_tensor().set(
                param, place)
            fluid.global_scope().find_var(names[1]).get_tensor().set(
                param, place)

        out = exe.run(feed={
            'inputs_basic_lstm': inputs_basic_lstm_np,
            'sequence_length': sequence_length_np,
            'inputs': inputs_np,
            'pre_hidden': pre_hidden_np,
            'pre_cell': pre_cell_np
        },
                      fetch_list=[output_new, rnn_out])

        self.assertTrue(np.allclose(out[0], out[1], rtol=1e-4))


class TestRnnUtil(unittest.TestCase):
    """
    Test cases for rnn apis' utility methods for coverage.
    """

    def test_case(self):
        inputs = {"key1": 1, "key2": 2}
        func = lambda x: x + 1
        outputs = utils.map_structure(func, inputs)
        utils.assert_same_structure(inputs, outputs)
        try:
            inputs["key3"] = 3
            utils.assert_same_structure(inputs, outputs)
        except ValueError as identifier:
            pass


class EncoderCell(RNNCell):
    """Encoder Cell"""

    def __init__(
        self,
        num_layers,
        hidden_size,
        dropout_prob=0.,
        init_scale=0.1,
    ):
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.dropout_prob = dropout_prob
        self.lstm_cells = []

        for i in range(num_layers):
            self.lstm_cells.append(LSTMCell(hidden_size))

    def call(self, step_input, states):
        new_states = []
        for i in range(self.num_layers):
            out, new_state = self.lstm_cells[i](step_input, states[i])
            step_input = layers.dropout(
                out,
                self.dropout_prob,
            ) if self.dropout_prob else out
            new_states.append(new_state)
        return step_input, new_states

    @property
    def state_shape(self):
        return [cell.state_shape for cell in self.lstm_cells]


class DecoderCell(RNNCell):
    """Decoder Cell"""

    def __init__(self, num_layers, hidden_size, dropout_prob=0.):
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.dropout_prob = dropout_prob
        self.lstm_cells = []
        for i in range(num_layers):
            self.lstm_cells.append(LSTMCell(hidden_size))

    def call(self, step_input, states):
        new_lstm_states = []
        for i in range(self.num_layers):
            out, new_lstm_state = self.lstm_cells[i](step_input, states[i])
            step_input = layers.dropout(
                out,
                self.dropout_prob,
            ) if self.dropout_prob else out
            new_lstm_states.append(new_lstm_state)
        return step_input, new_lstm_states


def def_seq2seq_model(num_layers, hidden_size, dropout_prob, src_vocab_size,
                      trg_vocab_size):
    "vanilla seq2seq model"
    # data
    source = fluid.data(name="src", shape=[None, None], dtype="int64")
    source_length = fluid.data(name="src_sequence_length",
                               shape=[None],
                               dtype="int64")
    target = fluid.data(name="trg", shape=[None, None], dtype="int64")
    target_length = fluid.data(name="trg_sequence_length",
                               shape=[None],
                               dtype="int64")
    label = fluid.data(name="label", shape=[None, None, 1], dtype="int64")

    # embedding
    src_emb = fluid.embedding(source, (src_vocab_size, hidden_size))
    tar_emb = fluid.embedding(target, (src_vocab_size, hidden_size))

    # encoder
    enc_cell = EncoderCell(num_layers, hidden_size, dropout_prob)
    enc_output, enc_final_state = dynamic_rnn(cell=enc_cell,
                                              inputs=src_emb,
                                              sequence_length=source_length)

    # decoder
    dec_cell = DecoderCell(num_layers, hidden_size, dropout_prob)
    dec_output, dec_final_state = dynamic_rnn(cell=dec_cell,
                                              inputs=tar_emb,
                                              initial_states=enc_final_state)
    logits = layers.fc(dec_output,
                       size=trg_vocab_size,
                       num_flatten_dims=len(dec_output.shape) - 1,
                       bias_attr=False)

    # loss
    loss = layers.softmax_with_cross_entropy(logits=logits,
                                             label=label,
                                             soft_label=False)
    loss = layers.unsqueeze(loss, axes=[2])
    max_tar_seq_len = layers.shape(target)[1]
    tar_mask = layers.sequence_mask(target_length,
                                    maxlen=max_tar_seq_len,
                                    dtype="float32")
    loss = loss * tar_mask
    loss = layers.reduce_mean(loss, dim=[0])
    loss = layers.reduce_sum(loss)

    # optimizer
    optimizer = fluid.optimizer.Adam(0.001)
    optimizer.minimize(loss)
    return loss


class TestSeq2SeqModel(unittest.TestCase):
    """
    Test cases to confirm seq2seq api training correctly.
    """

    def setUp(self):
        np.random.seed(123)
        self.model_hparams = {
            "num_layers": 2,
            "hidden_size": 128,
            "dropout_prob": 0.1,
            "src_vocab_size": 100,
            "trg_vocab_size": 100
        }

        self.iter_num = iter_num = 2
        self.batch_size = batch_size = 4
        src_seq_len = 10
        trg_seq_len = 12
        self.data = {
            "src":
            np.random.randint(
                2, self.model_hparams["src_vocab_size"],
                (iter_num * batch_size, src_seq_len)).astype("int64"),
            "src_sequence_length":
            np.random.randint(1, src_seq_len,
                              (iter_num * batch_size, )).astype("int64"),
            "trg":
            np.random.randint(
                2, self.model_hparams["src_vocab_size"],
                (iter_num * batch_size, trg_seq_len)).astype("int64"),
            "trg_sequence_length":
            np.random.randint(1, trg_seq_len,
                              (iter_num * batch_size, )).astype("int64"),
            "label":
            np.random.randint(
                2, self.model_hparams["src_vocab_size"],
                (iter_num * batch_size, trg_seq_len, 1)).astype("int64"),
        }

        place = core.CUDAPlace(
            0) if core.is_compiled_with_cuda() else core.CPUPlace()
        self.exe = Executor(place)

    def test_seq2seq_model(self):
        main_program = fluid.Program()
        startup_program = fluid.Program()
        with fluid.program_guard(main_program, startup_program):
            cost = def_seq2seq_model(**self.model_hparams)
            self.exe.run(startup_program)
            for iter_idx in range(self.iter_num):
                cost_val = self.exe.run(feed={
                    "src":
                    self.data["src"][iter_idx * self.batch_size:(iter_idx + 1) *
                                     self.batch_size, :],
                    "src_sequence_length":
                    self.data["src_sequence_length"][iter_idx *
                                                     self.batch_size:(iter_idx +
                                                                      1) *
                                                     self.batch_size],
                    "trg":
                    self.data["trg"][iter_idx * self.batch_size:(iter_idx + 1) *
                                     self.batch_size, :],
                    "trg_sequence_length":
                    self.data["trg_sequence_length"][iter_idx *
                                                     self.batch_size:(iter_idx +
                                                                      1) *
                                                     self.batch_size],
                    "label":
                    self.data["label"][iter_idx *
                                       self.batch_size:(iter_idx + 1) *
                                       self.batch_size]
                },
                                        fetch_list=[cost])[0]
                print("iter_idx: %d, cost: %f" % (iter_idx, cost_val))


if __name__ == '__main__':
    unittest.main()
