#
# Copyright 2016 The BigDL Authors.
#
# 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 pytest
import tempfile
import os

from unittest import TestCase
import numpy as np
import tensorflow as tf


def create_data(tf_data=False, batch_size=32):
    train_num_samples = 1000
    test_num_samples = 400
    input_feature_num = 10
    output_feature_num = 2
    past_seq_len = 10
    future_seq_len = 2
    
    def get_x_y(num_sample):
        x = np.random.randn(num_sample, past_seq_len, input_feature_num)
        y = np.random.randn(num_sample, future_seq_len, output_feature_num)
        return x, y

    train_data = get_x_y(train_num_samples)
    test_data = get_x_y(test_num_samples)

    if tf_data:
        from_tensor_slices = tf.data.Dataset.from_tensor_slices
        train_data = from_tensor_slices(train_data).cache()\
                                                   .shuffle(train_num_samples)\
                                                   .batch(batch_size)\
                                                   .prefetch(tf.data.AUTOTUNE)
        test_data = from_tensor_slices(test_data).cache()\
                                                 .batch(batch_size)\
                                                 .prefetch(tf.data.AUTOTUNE)
    return train_data, test_data


@pytest.mark.skipif(tf.__version__ < '2.0.0', reason="Run only when tf > 2.0.0.")
class TestSeq2SeqForecaster(TestCase):
    
    def setUp(self):
        from bigdl.chronos.forecaster.tf.seq2seq_forecaster import Seq2SeqForecaster
        self.forecaster = Seq2SeqForecaster(past_seq_len=10,
                                            future_seq_len=2,
                                            input_feature_num=10,
                                            output_feature_num=2)

    def tearDown(self):
        pass

    def test_seq2seq_fit_predict_evaluate(self):
        train_data, test_data = create_data()
        self.forecaster.fit(train_data,
                            epochs=2,
                            batch_size=32)
        yhat = self.forecaster.predict(test_data[0])
        assert yhat.shape == (400, 2, 2)
        mse = self.forecaster.evaluate(test_data, multioutput="raw_values")
        assert mse[0].shape == test_data[-1].shape[1:]
    
    def test_seq2seq_fit_tf_data(self):
        train_data, test_data = create_data(tf_data=True)
        self.forecaster.fit(train_data,
                            epochs=2)
        yhat = self.forecaster.predict(test_data)
        assert yhat.shape == (400, 2, 2)

    def test_seq2seq_save_load(self):
        train_data, test_data = create_data()
        self.forecaster.fit(train_data,
                            epochs=2,
                            batch_size=32)
        yhat = self.forecaster.predict(test_data[0])
        with tempfile.TemporaryDirectory() as tmp_dir_file:
            tmp_dir_file = os.path.join(tmp_dir_file, 'seq2seq.ckpt')
            self.forecaster.save(tmp_dir_file)
            self.forecaster.load(tmp_dir_file)
            from bigdl.chronos.model.tf2.Seq2Seq_keras import LSTMSeq2Seq
            assert isinstance(self.forecaster.internal, LSTMSeq2Seq)
        load_model_yhat = self.forecaster.predict(test_data[0])
        assert yhat.shape == (400, 2, 2)
        np.testing.assert_almost_equal(yhat, load_model_yhat, decimal=5)


if __name__ == '__main__':
    pytest.main([__file__])

[((53, 1, 53, 81), 'pytest.mark.skipif', 'pytest.mark.skipif', (), '', False, 'import pytest\n'), ((101, 4, 101, 27), 'pytest.main', 'pytest.main', ({(101, 16, 101, 26): '[__file__]'}, {}), '([__file__])', False, 'import pytest\n'), ((34, 12, 34, 72), 'numpy.random.randn', 'np.random.randn', ({(34, 28, 34, 38): 'num_sample', (34, 40, 34, 52): 'past_seq_len', (34, 54, 34, 71): 'input_feature_num'}, {}), '(num_sample, past_seq_len, input_feature_num)', True, 'import numpy as np\n'), ((35, 12, 35, 75), 'numpy.random.randn', 'np.random.randn', ({(35, 28, 35, 38): 'num_sample', (35, 40, 35, 54): 'future_seq_len', (35, 56, 35, 74): 'output_feature_num'}, {}), '(num_sample, future_seq_len, output_feature_num)', True, 'import numpy as np\n'), ((58, 26, 61, 65), 'bigdl.chronos.forecaster.tf.seq2seq_forecaster.Seq2SeqForecaster', 'Seq2SeqForecaster', (), '', False, 'from bigdl.chronos.forecaster.tf.seq2seq_forecaster import Seq2SeqForecaster\n'), ((97, 8, 97, 72), 'numpy.testing.assert_almost_equal', 'np.testing.assert_almost_equal', (), '', True, 'import numpy as np\n'), ((89, 13, 89, 42), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ({}, {}), '()', False, 'import tempfile\n'), ((90, 27, 90, 69), 'os.path.join', 'os.path.join', ({(90, 40, 90, 52): 'tmp_dir_file', (90, 54, 90, 68): '"""seq2seq.ckpt"""'}, {}), "(tmp_dir_file, 'seq2seq.ckpt')", False, 'import os\n')]