# Copyright 2016 The TensorFlow 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.
# ==============================================================================
"""Tests for convolutional recurrent layers."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.python.keras._impl import keras
from tensorflow.python.keras._impl.keras import testing_utils
from tensorflow.python.platform import test


class ConvLSTMTest(test.TestCase):

  def test_conv_lstm(self):
    num_row = 3
    num_col = 3
    filters = 2
    num_samples = 1
    input_channel = 2
    input_num_row = 5
    input_num_col = 5
    sequence_len = 2
    for data_format in ['channels_first', 'channels_last']:
      if data_format == 'channels_first':
        inputs = np.random.rand(num_samples, sequence_len,
                                input_channel,
                                input_num_row, input_num_col)
      else:
        inputs = np.random.rand(num_samples, sequence_len,
                                input_num_row, input_num_col,
                                input_channel)

      for return_sequences in [True, False]:
        with self.test_session():
          # test for return state:
          x = keras.Input(batch_shape=inputs.shape)
          kwargs = {'data_format': data_format,
                    'return_sequences': return_sequences,
                    'return_state': True,
                    'stateful': True,
                    'filters': filters,
                    'kernel_size': (num_row, num_col),
                    'padding': 'valid'}
          layer = keras.layers.ConvLSTM2D(**kwargs)
          layer.build(inputs.shape)
          outputs = layer(x)
          _, states = outputs[0], outputs[1:]
          self.assertEqual(len(states), 2)
          model = keras.models.Model(x, states[0])
          state = model.predict(inputs)
          self.assertAllClose(
              keras.backend.eval(layer.states[0]), state, atol=1e-4)

          # test for output shape:
          testing_utils.layer_test(
              keras.layers.ConvLSTM2D,
              kwargs={'data_format': data_format,
                      'return_sequences': return_sequences,
                      'filters': filters,
                      'kernel_size': (num_row, num_col),
                      'padding': 'valid'},
              input_shape=inputs.shape)

  def test_conv_lstm_statefulness(self):
    # Tests for statefulness
    num_row = 3
    num_col = 3
    filters = 2
    num_samples = 1
    input_channel = 2
    input_num_row = 5
    input_num_col = 5
    sequence_len = 2
    inputs = np.random.rand(num_samples, sequence_len,
                            input_num_row, input_num_col,
                            input_channel)

    with self.test_session():
      model = keras.models.Sequential()
      kwargs = {'data_format': 'channels_last',
                'return_sequences': False,
                'filters': filters,
                'kernel_size': (num_row, num_col),
                'stateful': True,
                'batch_input_shape': inputs.shape,
                'padding': 'same'}
      layer = keras.layers.ConvLSTM2D(**kwargs)

      model.add(layer)
      model.compile(optimizer='sgd', loss='mse')
      out1 = model.predict(np.ones_like(inputs))

      # train once so that the states change
      model.train_on_batch(np.ones_like(inputs),
                           np.random.random(out1.shape))
      out2 = model.predict(np.ones_like(inputs))

      # if the state is not reset, output should be different
      self.assertNotEqual(out1.max(), out2.max())

      # check that output changes after states are reset
      # (even though the model itself didn't change)
      layer.reset_states()
      out3 = model.predict(np.ones_like(inputs))
      self.assertNotEqual(out3.max(), out2.max())

      # check that container-level reset_states() works
      model.reset_states()
      out4 = model.predict(np.ones_like(inputs))
      self.assertAllClose(out3, out4, atol=1e-5)

      # check that the call to `predict` updated the states
      out5 = model.predict(np.ones_like(inputs))
      self.assertNotEqual(out4.max(), out5.max())

  def test_conv_lstm_regularizers(self):
    # check regularizers
    num_row = 3
    num_col = 3
    filters = 2
    num_samples = 1
    input_channel = 2
    input_num_row = 5
    input_num_col = 5
    sequence_len = 2
    inputs = np.random.rand(num_samples, sequence_len,
                            input_num_row, input_num_col,
                            input_channel)

    with self.test_session():
      kwargs = {'data_format': 'channels_last',
                'return_sequences': False,
                'kernel_size': (num_row, num_col),
                'stateful': True,
                'filters': filters,
                'batch_input_shape': inputs.shape,
                'kernel_regularizer': keras.regularizers.L1L2(l1=0.01),
                'recurrent_regularizer': keras.regularizers.L1L2(l1=0.01),
                'activity_regularizer': 'l2',
                'bias_regularizer': 'l2',
                'kernel_constraint': 'max_norm',
                'recurrent_constraint': 'max_norm',
                'bias_constraint': 'max_norm',
                'padding': 'same'}

      layer = keras.layers.ConvLSTM2D(**kwargs)
      layer.build(inputs.shape)
      self.assertEqual(len(layer.losses), 3)
      layer(keras.backend.variable(np.ones(inputs.shape)))
      self.assertEqual(len(layer.losses), 4)

  def test_conv_lstm_dropout(self):
    # check dropout
    with self.test_session():
      testing_utils.layer_test(
          keras.layers.ConvLSTM2D,
          kwargs={'data_format': 'channels_last',
                  'return_sequences': False,
                  'filters': 2,
                  'kernel_size': (3, 3),
                  'padding': 'same',
                  'dropout': 0.1,
                  'recurrent_dropout': 0.1},
          input_shape=(1, 2, 5, 5, 2))


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