# Copyright 2019 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.
# ==============================================================================
"""DataAdapter tests."""

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

import math

from absl.testing import parameterized
import numpy as np

from tensorflow.python import keras
from tensorflow.python.data.experimental.ops import cardinality
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.keras.engine import data_adapter
from tensorflow.python.keras.utils import data_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.platform import test
from tensorflow.python.util import nest

class DummyArrayLike(object):
  """Dummy array-like object."""

  def __init__(self, data):
    self.data = data

  def __len__(self):
    return len(self.data)

  def __getitem__(self, key):
    return self.data[key]

  @property
  def shape(self):
    return self.data.shape

  @property
  def dtype(self):
    return self.data.dtype

def fail_on_convert(x, **kwargs):
  _ = x
  _ = kwargs
  raise TypeError('Cannot convert DummyArrayLike to a tensor')
ops.register_tensor_conversion_function(DummyArrayLike, fail_on_convert)

class DataAdapterTestBase(keras_parameterized.TestCase):

  def setUp(self):
    super(DataAdapterTestBase, self).setUp()
    self.batch_size = 5
    self.numpy_input = np.zeros((50, 10))
    self.numpy_target = np.ones(50)
    self.tensor_input = constant_op.constant(2.0, shape=(50, 10))
    self.tensor_target = array_ops.ones((50,))
    self.arraylike_input = DummyArrayLike(self.numpy_input)
    self.arraylike_target = DummyArrayLike(self.numpy_target)
    self.dataset_input = dataset_ops.DatasetV2.from_tensor_slices(
        (self.numpy_input, self.numpy_target)).shuffle(50).batch(
            self.batch_size)

    def generator():
      while True:
        yield (np.zeros((self.batch_size, 10)), np.ones(self.batch_size))
    self.generator_input = generator()
    self.iterator_input = data_utils.threadsafe_generator(generator)()
    self.sequence_input = TestSequence(batch_size=self.batch_size,
                                       feature_shape=10)
    self.model = keras.models.Sequential(
        [keras.layers.Dense(8, input_shape=(10,), activation='softmax')])

class TestSequence(data_utils.Sequence):

  def __init__(self, batch_size, feature_shape):
    self.batch_size = batch_size
    self.feature_shape = feature_shape

  def __getitem__(self, item):
    return (np.zeros((self.batch_size, self.feature_shape)),
            np.ones((self.batch_size,)))

  def __len__(self):
    return 10

class TensorLikeDataAdapterTest(DataAdapterTestBase):

  def setUp(self):
    super(TensorLikeDataAdapterTest, self).setUp()
    self.adapter_cls = data_adapter.TensorLikeDataAdapter

  def test_can_handle_numpy(self):
    self.assertTrue(self.adapter_cls.can_handle(self.numpy_input))
    self.assertTrue(
        self.adapter_cls.can_handle(self.numpy_input, self.numpy_target))

    self.assertFalse(self.adapter_cls.can_handle(self.dataset_input))
    self.assertFalse(self.adapter_cls.can_handle(self.generator_input))
    self.assertFalse(self.adapter_cls.can_handle(self.sequence_input))

  def test_size_numpy(self):
    adapter = self.adapter_cls(
        self.numpy_input, self.numpy_target, batch_size=5)
    self.assertEqual(adapter.get_size(), 10)
    self.assertFalse(adapter.has_partial_batch())

  def test_batch_size_numpy(self):
    adapter = self.adapter_cls(
        self.numpy_input, self.numpy_target, batch_size=5)
    self.assertEqual(adapter.batch_size(), 5)

  def test_partial_batch_numpy(self):
    adapter = self.adapter_cls(
        self.numpy_input, self.numpy_target, batch_size=4)
    self.assertEqual(adapter.get_size(), 13)   # 50/4
    self.assertTrue(adapter.has_partial_batch())
    self.assertEqual(adapter.partial_batch_size(), 2)

  def test_epochs(self):
    num_epochs = 3
    adapter = self.adapter_cls(
        self.numpy_input, self.numpy_target, batch_size=5, epochs=num_epochs)
    ds_iter = iter(adapter.get_dataset())
    num_batches_per_epoch = self.numpy_input.shape[0] // 5
    for _ in range(num_batches_per_epoch * num_epochs):
      next(ds_iter)
    with self.assertRaises(StopIteration):
      next(ds_iter)

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training_numpy(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.numpy_input, self.numpy_target, batch_size=5)

  def test_can_handle_pandas(self):
    try:
      import pandas as pd  # pylint: disable=g-import-not-at-top
    except ImportError:
      self.skipTest('Skipping test because pandas is not installed.')
    self.assertTrue(self.adapter_cls.can_handle(pd.DataFrame(self.numpy_input)))
    self.assertTrue(
        self.adapter_cls.can_handle(pd.DataFrame(self.numpy_input)[0]))
    self.assertTrue(
        self.adapter_cls.can_handle(
            pd.DataFrame(self.numpy_input),
            pd.DataFrame(self.numpy_input)[0]))

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training_pandas(self):
    try:
      import pandas as pd  # pylint: disable=g-import-not-at-top
    except ImportError:
      self.skipTest('Skipping test because pandas is not installed.')
    input_a = keras.Input(shape=(3,), name='input_a')
    input_b = keras.Input(shape=(3,), name='input_b')
    input_c = keras.Input(shape=(1,), name='input_b')

    x = keras.layers.Dense(4, name='dense_1')(input_a)
    y = keras.layers.Dense(3, name='dense_2')(input_b)
    z = keras.layers.Dense(1, name='dense_3')(input_c)

    model_1 = keras.Model(inputs=input_a, outputs=x)
    model_2 = keras.Model(inputs=[input_a, input_b], outputs=[x, y])
    model_3 = keras.Model(inputs=input_c, outputs=z)

    model_1.compile(optimizer='rmsprop', loss='mse')
    model_2.compile(optimizer='rmsprop', loss='mse')

    input_a_np = np.random.random((10, 3))
    input_b_np = np.random.random((10, 3))
    input_a_df = pd.DataFrame(input_a_np)
    input_b_df = pd.DataFrame(input_b_np)

    output_a_df = pd.DataFrame(np.random.random((10, 4)))
    output_b_df = pd.DataFrame(np.random.random((10, 3)))

    model_1.fit(input_a_df,
                output_a_df)
    model_2.fit([input_a_df, input_b_df],
                [output_a_df, output_b_df])
    model_1.fit([input_a_df],
                [output_a_df])
    model_1.fit({'input_a': input_a_df},
                output_a_df)
    model_2.fit({'input_a': input_a_df, 'input_b': input_b_df},
                [output_a_df, output_b_df])

    model_1.evaluate(input_a_df,
                     output_a_df)
    model_2.evaluate([input_a_df, input_b_df],
                     [output_a_df, output_b_df])
    model_1.evaluate([input_a_df],
                     [output_a_df])
    model_1.evaluate({'input_a': input_a_df},
                     output_a_df)
    model_2.evaluate({'input_a': input_a_df, 'input_b': input_b_df},
                     [output_a_df, output_b_df])

    # Verify predicting on pandas vs numpy returns the same result
    predict_1_pandas = model_1.predict(input_a_df)
    predict_2_pandas = model_2.predict([input_a_df, input_b_df])
    predict_3_pandas = model_3.predict(input_a_df[0])

    predict_1_numpy = model_1.predict(input_a_np)
    predict_2_numpy = model_2.predict([input_a_np, input_b_np])
    predict_3_numpy = model_3.predict(np.asarray(input_a_df[0]))

    self.assertAllClose(predict_1_numpy, predict_1_pandas)
    self.assertAllClose(predict_2_numpy, predict_2_pandas)
    self.assertAllClose(predict_3_numpy, predict_3_pandas)

    # Extra ways to pass in dataframes
    model_1.predict([input_a_df])
    model_1.predict({'input_a': input_a_df})
    model_2.predict({'input_a': input_a_df, 'input_b': input_b_df})

  def test_can_handle(self):
    self.assertTrue(self.adapter_cls.can_handle(self.tensor_input))
    self.assertTrue(
        self.adapter_cls.can_handle(self.tensor_input, self.tensor_target))

    self.assertFalse(self.adapter_cls.can_handle(self.arraylike_input))
    self.assertFalse(
        self.adapter_cls.can_handle(self.arraylike_input,
                                    self.arraylike_target))
    self.assertFalse(self.adapter_cls.can_handle(self.dataset_input))
    self.assertFalse(self.adapter_cls.can_handle(self.generator_input))
    self.assertFalse(self.adapter_cls.can_handle(self.sequence_input))

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.tensor_input, self.tensor_target, batch_size=5)

  def test_size(self):
    adapter = self.adapter_cls(
        self.tensor_input, self.tensor_target, batch_size=5)
    self.assertEqual(adapter.get_size(), 10)
    self.assertFalse(adapter.has_partial_batch())

  def test_shuffle_correctness(self):
    with context.eager_mode():
      num_samples = 100
      batch_size = 32
      x = np.arange(num_samples)
      np.random.seed(99)
      adapter = self.adapter_cls(
          x, y=None, batch_size=batch_size, shuffle=True, epochs=2)

      def _get_epoch(ds_iter):
        ds_data = []
        for _ in range(int(math.ceil(num_samples / batch_size))):
          ds_data.append(next(ds_iter)[0].numpy())
        return np.concatenate(ds_data)

      ds_iter = iter(adapter.get_dataset())

      # First epoch.
      epoch_data = _get_epoch(ds_iter)
      # Check that shuffling occurred.
      self.assertNotAllClose(x, epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(epoch_data))

      # Second epoch.
      second_epoch_data = _get_epoch(ds_iter)
      # Check that shuffling occurred.
      self.assertNotAllClose(x, second_epoch_data)
      # Check that shuffling is different across epochs.
      self.assertNotAllClose(epoch_data, second_epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(second_epoch_data))

  def test_batch_shuffle_correctness(self):
    with context.eager_mode():
      num_samples = 100
      batch_size = 6
      x = np.arange(num_samples)
      np.random.seed(99)
      adapter = self.adapter_cls(
          x, y=None, batch_size=batch_size, shuffle='batch', epochs=2)

      def _get_epoch_batches(ds_iter):
        ds_data = []
        for _ in range(int(math.ceil(num_samples / batch_size))):
          ds_data.append(next(ds_iter)[0].numpy())
        return ds_data

      ds_iter = iter(adapter.get_dataset())

      # First epoch.
      epoch_batch_data = _get_epoch_batches(ds_iter)
      epoch_data = np.concatenate(epoch_batch_data)

      def _verify_batch(batch):
        # Verify that a batch contains only contiguous data, and that it has
        # been shuffled.
        shuffled_batch = np.sort(batch)
        self.assertNotAllClose(batch, shuffled_batch)
        for i in range(1, len(batch)):
          self.assertEqual(shuffled_batch[i-1] + 1, shuffled_batch[i])

      # Assert that the data within each batch remains contiguous
      for batch in epoch_batch_data:
        _verify_batch(batch)

      # Check that individual batches are unshuffled
      # Check that shuffling occurred.
      self.assertNotAllClose(x, epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(epoch_data))

      # Second epoch.
      second_epoch_batch_data = _get_epoch_batches(ds_iter)
      second_epoch_data = np.concatenate(second_epoch_batch_data)

      # Assert that the data within each batch remains contiguous
      for batch in second_epoch_batch_data:
        _verify_batch(batch)

      # Check that shuffling occurred.
      self.assertNotAllClose(x, second_epoch_data)
      # Check that shuffling is different across epochs.
      self.assertNotAllClose(epoch_data, second_epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(second_epoch_data))

  @parameterized.named_parameters(
      ('batch_size_5', 5, None, 5),
      ('batch_size_50', 50, 4, 50),  # Sanity check: batch_size takes precedence
      ('steps_1', None, 1, 50),
      ('steps_4', None, 4, 13),
      )
  def test_batch_size(self, batch_size_in, steps, batch_size_out):
    adapter = self.adapter_cls(
        self.tensor_input, self.tensor_target, batch_size=batch_size_in,
        steps=steps)
    self.assertEqual(adapter.batch_size(), batch_size_out)

  @parameterized.named_parameters(
      ('batch_size_5', 5, None, 10, 0),
      ('batch_size_4', 4, None, 13, 2),
      ('steps_1', None, 1, 1, 0),
      ('steps_5', None, 5, 5, 0),
      ('steps_4', None, 4, 4, 11),
      )
  def test_partial_batch(
      self, batch_size_in, steps, size, partial_batch_size):
    adapter = self.adapter_cls(
        self.tensor_input, self.tensor_target, batch_size=batch_size_in,
        steps=steps)
    self.assertEqual(adapter.get_size(), size)   # 50/steps
    self.assertEqual(adapter.has_partial_batch(), bool(partial_batch_size))
    self.assertEqual(adapter.partial_batch_size(), partial_batch_size or None)

class GenericArrayLikeDataAdapterTest(DataAdapterTestBase):

  def setUp(self):
    super(GenericArrayLikeDataAdapterTest, self).setUp()
    self.adapter_cls = data_adapter.GenericArrayLikeDataAdapter

  def test_can_handle_some_numpy(self):
    self.assertTrue(self.adapter_cls.can_handle(
        self.arraylike_input))
    self.assertTrue(
        self.adapter_cls.can_handle(self.arraylike_input,
                                    self.arraylike_target))

    # Because adapters are mutually exclusive, don't handle cases
    # where all the data is numpy or an eagertensor
    self.assertFalse(self.adapter_cls.can_handle(self.numpy_input))
    self.assertFalse(
        self.adapter_cls.can_handle(self.numpy_input,
                                    self.numpy_target))
    self.assertFalse(self.adapter_cls.can_handle(self.tensor_input))
    self.assertFalse(
        self.adapter_cls.can_handle(self.tensor_input, self.tensor_target))

    # But do handle mixes that include generic arraylike data
    self.assertTrue(
        self.adapter_cls.can_handle(self.numpy_input,
                                    self.arraylike_target))
    self.assertTrue(
        self.adapter_cls.can_handle(self.arraylike_input,
                                    self.numpy_target))
    self.assertTrue(
        self.adapter_cls.can_handle(self.arraylike_input,
                                    self.tensor_target))
    self.assertTrue(
        self.adapter_cls.can_handle(self.tensor_input,
                                    self.arraylike_target))

    self.assertFalse(self.adapter_cls.can_handle(self.dataset_input))
    self.assertFalse(self.adapter_cls.can_handle(self.generator_input))
    self.assertFalse(self.adapter_cls.can_handle(self.sequence_input))

  def test_size(self):
    adapter = self.adapter_cls(
        self.arraylike_input,
        self.arraylike_target, batch_size=5)
    self.assertEqual(adapter.get_size(), 10)
    self.assertFalse(adapter.has_partial_batch())

  def test_epochs(self):
    num_epochs = 3
    adapter = self.adapter_cls(
        self.arraylike_input,
        self.numpy_target, batch_size=5, epochs=num_epochs)
    ds_iter = iter(adapter.get_dataset())
    num_batches_per_epoch = self.numpy_input.shape[0] // 5
    for _ in range(num_batches_per_epoch * num_epochs):
      next(ds_iter)
    with self.assertRaises(StopIteration):
      next(ds_iter)

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training(self):
    # First verify that DummyArrayLike can't be converted to a Tensor
    with self.assertRaises(TypeError):
      ops.convert_to_tensor_v2(self.arraylike_input)

    # Then train on the array like.
    # It should not be converted to a tensor directly (which would force it into
    # memory), only the sliced data should be converted.
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.arraylike_input,
                   self.arraylike_target, batch_size=5)
    self.model.fit(self.arraylike_input,
                   self.arraylike_target,
                   shuffle=True, batch_size=5)
    self.model.fit(self.arraylike_input,
                   self.arraylike_target,
                   shuffle='batch', batch_size=5)
    self.model.evaluate(self.arraylike_input,
                        self.arraylike_target, batch_size=5)
    self.model.predict(self.arraylike_input, batch_size=5)

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training_numpy_target(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.arraylike_input,
                   self.numpy_target, batch_size=5)
    self.model.fit(self.arraylike_input,
                   self.numpy_target, shuffle=True,
                   batch_size=5)
    self.model.fit(self.arraylike_input,
                   self.numpy_target, shuffle='batch',
                   batch_size=5)
    self.model.evaluate(self.arraylike_input,
                        self.numpy_target, batch_size=5)

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training_tensor_target(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.arraylike_input,
                   self.tensor_target, batch_size=5)
    self.model.fit(self.arraylike_input,
                   self.tensor_target, shuffle=True,
                   batch_size=5)
    self.model.fit(self.arraylike_input,
                   self.tensor_target, shuffle='batch',
                   batch_size=5)
    self.model.evaluate(self.arraylike_input,
                        self.tensor_target, batch_size=5)

  def test_shuffle_correctness(self):
    with context.eager_mode():
      num_samples = 100
      batch_size = 32
      x = DummyArrayLike(np.arange(num_samples))
      np.random.seed(99)
      adapter = self.adapter_cls(
          x, y=None, batch_size=batch_size, shuffle=True, epochs=2)

      def _get_epoch(ds_iter):
        ds_data = []
        for _ in range(int(math.ceil(num_samples / batch_size))):
          ds_data.append(next(ds_iter)[0].numpy())
        return np.concatenate(ds_data)

      ds_iter = iter(adapter.get_dataset())

      # First epoch.
      epoch_data = _get_epoch(ds_iter)
      # Check that shuffling occurred.
      self.assertNotAllClose(x, epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(epoch_data))

      # Second epoch.
      second_epoch_data = _get_epoch(ds_iter)
      # Check that shuffling occurred.
      self.assertNotAllClose(x, second_epoch_data)
      # Check that shuffling is different across epochs.
      self.assertNotAllClose(epoch_data, second_epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(second_epoch_data))

  def test_batch_shuffle_correctness(self):
    with context.eager_mode():
      num_samples = 100
      batch_size = 6
      x = DummyArrayLike(np.arange(num_samples))
      np.random.seed(99)
      adapter = self.adapter_cls(
          x, y=None, batch_size=batch_size, shuffle='batch', epochs=2)

      def _get_epoch_batches(ds_iter):
        ds_data = []
        for _ in range(int(math.ceil(num_samples / batch_size))):
          ds_data.append(next(ds_iter)[0].numpy())
        return ds_data

      ds_iter = iter(adapter.get_dataset())

      # First epoch.
      epoch_batch_data = _get_epoch_batches(ds_iter)
      epoch_data = np.concatenate(epoch_batch_data)

      def _verify_batch(batch):
        # Verify that a batch contains only contiguous data, but that it has
        # been shuffled.
        shuffled_batch = np.sort(batch)
        self.assertNotAllClose(batch, shuffled_batch)
        for i in range(1, len(batch)):
          self.assertEqual(shuffled_batch[i-1] + 1, shuffled_batch[i])

      # Assert that the data within each batch is shuffled contiguous data
      for batch in epoch_batch_data:
        _verify_batch(batch)

      # Check that individual batches are unshuffled
      # Check that shuffling occurred.
      self.assertNotAllClose(x, epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(epoch_data))

      # Second epoch.
      second_epoch_batch_data = _get_epoch_batches(ds_iter)
      second_epoch_data = np.concatenate(second_epoch_batch_data)

      # Assert that the data within each batch remains contiguous
      for batch in second_epoch_batch_data:
        _verify_batch(batch)

      # Check that shuffling occurred.
      self.assertNotAllClose(x, second_epoch_data)
      # Check that shuffling is different across epochs.
      self.assertNotAllClose(epoch_data, second_epoch_data)
      # Check that each elements appears, and only once.
      self.assertAllClose(x, np.sort(second_epoch_data))

  @parameterized.named_parameters(
      ('batch_size_5', 5, None, 5),
      ('batch_size_50', 50, 4, 50),  # Sanity check: batch_size takes precedence
      ('steps_1', None, 1, 50),
      ('steps_4', None, 4, 13),
  )
  def test_batch_size(self, batch_size_in, steps, batch_size_out):
    adapter = self.adapter_cls(
        self.arraylike_input,
        self.arraylike_target, batch_size=batch_size_in,
        steps=steps)
    self.assertEqual(adapter.batch_size(), batch_size_out)

  @parameterized.named_parameters(
      ('batch_size_5', 5, None, 10, 0),
      ('batch_size_4', 4, None, 13, 2),
      ('steps_1', None, 1, 1, 0),
      ('steps_5', None, 5, 5, 0),
      ('steps_4', None, 4, 4, 11),
  )
  def test_partial_batch(
      self, batch_size_in, steps, size, partial_batch_size):
    adapter = self.adapter_cls(
        self.arraylike_input, self.arraylike_target,
        batch_size=batch_size_in,
        steps=steps)
    self.assertEqual(adapter.get_size(), size)   # 50/steps
    self.assertEqual(adapter.has_partial_batch(), bool(partial_batch_size))
    self.assertEqual(adapter.partial_batch_size(), partial_batch_size or None)

class DatasetAdapterTest(DataAdapterTestBase):

  def setUp(self):
    super(DatasetAdapterTest, self).setUp()
    self.adapter_cls = data_adapter.DatasetAdapter

  def test_can_handle(self):
    self.assertFalse(self.adapter_cls.can_handle(self.numpy_input))
    self.assertFalse(self.adapter_cls.can_handle(self.tensor_input))
    self.assertTrue(self.adapter_cls.can_handle(self.dataset_input))
    self.assertFalse(self.adapter_cls.can_handle(self.generator_input))
    self.assertFalse(self.adapter_cls.can_handle(self.sequence_input))

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training(self):
    dataset = self.adapter_cls(self.dataset_input).get_dataset()
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(dataset)

  def test_size(self):
    adapter = self.adapter_cls(self.dataset_input)
    self.assertIsNone(adapter.get_size())

  def test_batch_size(self):
    adapter = self.adapter_cls(self.dataset_input)
    self.assertIsNone(adapter.batch_size())

  def test_partial_batch(self):
    adapter = self.adapter_cls(self.dataset_input)
    self.assertFalse(adapter.has_partial_batch())
    self.assertIsNone(adapter.partial_batch_size())

  def test_invalid_targets_argument(self):
    with self.assertRaisesRegexp(ValueError, r'`y` argument is not supported'):
      self.adapter_cls(self.dataset_input, y=self.dataset_input)

  def test_invalid_sample_weights_argument(self):
    with self.assertRaisesRegexp(ValueError,
                                 r'`sample_weight` argument is not supported'):
      self.adapter_cls(self.dataset_input, sample_weights=self.dataset_input)

class GeneratorDataAdapterTest(DataAdapterTestBase):

  def setUp(self):
    super(GeneratorDataAdapterTest, self).setUp()
    self.adapter_cls = data_adapter.GeneratorDataAdapter

  def test_can_handle(self):
    self.assertFalse(self.adapter_cls.can_handle(self.numpy_input))
    self.assertFalse(self.adapter_cls.can_handle(self.tensor_input))
    self.assertFalse(self.adapter_cls.can_handle(self.dataset_input))
    self.assertTrue(self.adapter_cls.can_handle(self.generator_input))
    self.assertFalse(self.adapter_cls.can_handle(self.sequence_input))

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.generator_input, steps_per_epoch=10)

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  @test_util.run_v2_only
  @data_utils.dont_use_multiprocessing_pool
  def test_with_multiprocessing_training(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.iterator_input, workers=1, use_multiprocessing=True,
                   max_queue_size=10, steps_per_epoch=10)
    # Fit twice to ensure there isn't any duplication that prevent the worker
    # from starting.
    self.model.fit(self.iterator_input, workers=1, use_multiprocessing=True,
                   max_queue_size=10, steps_per_epoch=10)

  def test_size(self):
    adapter = self.adapter_cls(self.generator_input)
    self.assertIsNone(adapter.get_size())

  def test_batch_size(self):
    adapter = self.adapter_cls(self.generator_input)
    self.assertEqual(adapter.batch_size(), None)
    self.assertEqual(adapter.representative_batch_size(), 5)

  def test_partial_batch(self):
    adapter = self.adapter_cls(self.generator_input)
    self.assertFalse(adapter.has_partial_batch())
    self.assertIsNone(adapter.partial_batch_size())

  def test_invalid_targets_argument(self):
    with self.assertRaisesRegexp(ValueError, r'`y` argument is not supported'):
      self.adapter_cls(self.generator_input, y=self.generator_input)

  def test_invalid_sample_weights_argument(self):
    with self.assertRaisesRegexp(ValueError,
                                 r'`sample_weight` argument is not supported'):
      self.adapter_cls(
          self.generator_input, sample_weights=self.generator_input)

  def test_not_shuffled(self):
    def generator():
      for i in range(10):
        yield np.ones((1, 1)) * i

    adapter = self.adapter_cls(generator(), shuffle=True)
    with context.eager_mode():
      for i, data in enumerate(adapter.get_dataset()):
        self.assertEqual(i, data[0].numpy().flatten())

class KerasSequenceAdapterTest(DataAdapterTestBase):

  def setUp(self):
    super(KerasSequenceAdapterTest, self).setUp()
    self.adapter_cls = data_adapter.KerasSequenceAdapter

  def test_can_handle(self):
    self.assertFalse(self.adapter_cls.can_handle(self.numpy_input))
    self.assertFalse(self.adapter_cls.can_handle(self.tensor_input))
    self.assertFalse(self.adapter_cls.can_handle(self.dataset_input))
    self.assertFalse(self.adapter_cls.can_handle(self.generator_input))
    self.assertTrue(self.adapter_cls.can_handle(self.sequence_input))

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  def test_training(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.sequence_input)

  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  @test_util.run_v2_only
  @data_utils.dont_use_multiprocessing_pool
  def test_with_multiprocessing_training(self):
    self.model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd',
                       run_eagerly=testing_utils.should_run_eagerly())
    self.model.fit(self.sequence_input, workers=1, use_multiprocessing=True,
                   max_queue_size=10, steps_per_epoch=10)
    # Fit twice to ensure there isn't any duplication that prevent the worker
    # from starting.
    self.model.fit(self.sequence_input, workers=1, use_multiprocessing=True,
                   max_queue_size=10, steps_per_epoch=10)

  def test_size(self):
    adapter = self.adapter_cls(self.sequence_input)
    self.assertEqual(adapter.get_size(), 10)

  def test_batch_size(self):
    adapter = self.adapter_cls(self.sequence_input)
    self.assertEqual(adapter.batch_size(), None)
    self.assertEqual(adapter.representative_batch_size(), 5)

  def test_partial_batch(self):
    adapter = self.adapter_cls(self.sequence_input)
    self.assertFalse(adapter.has_partial_batch())
    self.assertIsNone(adapter.partial_batch_size())

  def test_invalid_targets_argument(self):
    with self.assertRaisesRegexp(ValueError, r'`y` argument is not supported'):
      self.adapter_cls(self.sequence_input, y=self.sequence_input)

  def test_invalid_sample_weights_argument(self):
    with self.assertRaisesRegexp(ValueError,
                                 r'`sample_weight` argument is not supported'):
      self.adapter_cls(self.sequence_input, sample_weights=self.sequence_input)

class DataHandlerTest(keras_parameterized.TestCase):

  def test_finite_dataset_with_steps_per_epoch(self):
    data = dataset_ops.Dataset.from_tensor_slices([0, 1, 2, 3]).batch(1)
    # User can choose to only partially consume `Dataset`.
    data_handler = data_adapter.DataHandler(
        data, initial_epoch=0, epochs=2, steps_per_epoch=2)
    self.assertEqual(data_handler.inferred_steps, 2)
    self.assertFalse(data_handler._adapter.should_recreate_iterator())
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator).numpy())
      returned_data.append(epoch_data)
    self.assertEqual(returned_data, [[0, 1], [2, 3]])

  def test_finite_dataset_without_steps_per_epoch(self):
    data = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]).batch(1)
    data_handler = data_adapter.DataHandler(data, initial_epoch=0, epochs=2)
    self.assertEqual(data_handler.inferred_steps, 3)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator).numpy())
      returned_data.append(epoch_data)
    self.assertEqual(returned_data, [[0, 1, 2], [0, 1, 2]])

  def test_finite_dataset_with_steps_per_epoch_exact_size(self):
    data = dataset_ops.Dataset.from_tensor_slices([0, 1, 2, 3]).batch(1)
    # If user specifies exact size of `Dataset` as `steps_per_epoch`,
    # create a new iterator each epoch.
    data_handler = data_adapter.DataHandler(
        data, initial_epoch=0, epochs=2, steps_per_epoch=4)
    self.assertTrue(data_handler._adapter.should_recreate_iterator())
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator).numpy())
      returned_data.append(epoch_data)
    self.assertEqual(returned_data, [[0, 1, 2, 3], [0, 1, 2, 3]])

  def test_infinite_dataset_with_steps_per_epoch(self):
    data = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]).batch(1).repeat()
    data_handler = data_adapter.DataHandler(
        data, initial_epoch=0, epochs=2, steps_per_epoch=3)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator).numpy())
      returned_data.append(epoch_data)
    self.assertEqual(returned_data, [[0, 1, 2], [0, 1, 2]])

  def test_unknown_cardinality_dataset_with_steps_per_epoch(self):
    ds = dataset_ops.DatasetV2.from_tensor_slices([0, 1, 2, 3, 4, 5, 6])
    filtered_ds = ds.filter(lambda x: x < 4)
    self.assertEqual(
        cardinality.cardinality(filtered_ds).numpy(), cardinality.UNKNOWN)

    # User can choose to only partially consume `Dataset`.
    data_handler = data_adapter.DataHandler(
        filtered_ds, initial_epoch=0, epochs=2, steps_per_epoch=2)
    self.assertFalse(data_handler._adapter.should_recreate_iterator())
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(returned_data)
    self.assertEqual(returned_data, [[0, 1], [2, 3]])
    self.assertEqual(data_handler.inferred_steps, 2)

  def test_unknown_cardinality_dataset_without_steps_per_epoch(self):
    ds = dataset_ops.DatasetV2.from_tensor_slices([0, 1, 2, 3, 4, 5, 6])
    filtered_ds = ds.filter(lambda x: x < 4)
    self.assertEqual(
        cardinality.cardinality(filtered_ds).numpy(), cardinality.UNKNOWN)

    data_handler = data_adapter.DataHandler(
        filtered_ds, initial_epoch=0, epochs=2)
    self.assertEqual(data_handler.inferred_steps, None)
    self.assertTrue(data_handler._adapter.should_recreate_iterator())
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      with data_handler.catch_stop_iteration():
        for _ in data_handler.steps():
          epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(returned_data)
    self.assertEqual(returned_data, [[0, 1, 2, 3], [0, 1, 2, 3]])
    self.assertEqual(data_handler.inferred_steps, 4)

  def test_insufficient_data(self):
    ds = dataset_ops.DatasetV2.from_tensor_slices([0, 1])
    ds = ds.filter(lambda *args, **kwargs: True)
    data_handler = data_adapter.DataHandler(
        ds, initial_epoch=0, epochs=2, steps_per_epoch=3)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        with data_handler.catch_stop_iteration():
          epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(returned_data)
    self.assertTrue(data_handler._insufficient_data)
    self.assertEqual(returned_data, [[0, 1]])

  def test_numpy(self):
    x = np.array([0, 1, 2])
    y = np.array([0, 2, 4])
    sw = np.array([0, 4, 8])
    data_handler = data_adapter.DataHandler(
        x=x, y=y, sample_weight=sw, batch_size=1, epochs=2)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(returned_data)
    self.assertEqual(returned_data,
                     [[(0, 0, 0), (1, 2, 4),
                       (2, 4, 8)], [(0, 0, 0), (1, 2, 4), (2, 4, 8)]])

  def test_generator(self):

    def generator():
      for _ in range(2):
        for step in range(3):
          yield (ops.convert_to_tensor_v2([step]),)

    data_handler = data_adapter.DataHandler(
        generator(), epochs=2, steps_per_epoch=3)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(returned_data)
    self.assertEqual(returned_data, [[([0],), ([1],),
                                      ([2],)], [([0],), ([1],), ([2],)]])

  def test_composite_tensor(self):
    st = sparse_tensor.SparseTensor(
        indices=[[0, 0], [1, 0], [2, 0]], values=[0, 1, 2], dense_shape=[3, 1])
    data_handler = data_adapter.DataHandler(st, epochs=2, steps_per_epoch=3)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(
        nest.map_structure(sparse_ops.sparse_tensor_to_dense, returned_data))
    self.assertEqual(returned_data, [[([0],), ([1],),
                                      ([2],)], [([0],), ([1],), ([2],)]])

  def test_list_of_scalars(self):
    data_handler = data_adapter.DataHandler([[0], [1], [2]],
                                            epochs=2,
                                            steps_per_epoch=3)
    returned_data = []
    for _, iterator in data_handler.enumerate_epochs():
      epoch_data = []
      for _ in data_handler.steps():
        epoch_data.append(next(iterator))
      returned_data.append(epoch_data)
    returned_data = self.evaluate(returned_data)
    self.assertEqual(returned_data, [[([0],), ([1],),
                                      ([2],)], [([0],), ([1],), ([2],)]])

  def test_class_weight_user_errors(self):
    with self.assertRaisesRegexp(ValueError, 'to be a dict with keys'):
      data_adapter.DataHandler(
          x=[[0], [1], [2]],
          y=[[2], [1], [0]],
          batch_size=1,
          sample_weight=[[1.], [2.], [4.]],
          class_weight={
              0: 0.5,
              1: 1.,
              3: 1.5  # Skips class `2`.
          })

    with self.assertRaisesRegexp(ValueError, 'with a single output'):
      data_adapter.DataHandler(
          x=np.ones((10, 1)),
          y=[np.ones((10, 1)), np.zeros((10, 1))],
          batch_size=2,
          class_weight={
              0: 0.5,
              1: 1.,
              2: 1.5
          })

class TestValidationSplit(keras_parameterized.TestCase):

  @parameterized.named_parameters(('numpy_arrays', True), ('tensors', False))
  def test_validation_split_shuffled(self, use_numpy):
    if use_numpy:
      x = np.array([0, 1, 2, 3, 4])
      y = np.array([0, 2, 4, 6, 8])
      sw = np.array([0, 4, 8, 12, 16])
    else:
      x = ops.convert_to_tensor_v2([0, 1, 2, 3, 4])
      y = ops.convert_to_tensor_v2([0, 2, 4, 6, 8])
      sw = ops.convert_to_tensor_v2([0, 4, 8, 12, 16])

    (train_x, train_y, train_sw), (val_x, val_y, val_sw) = (
        data_adapter.train_validation_split((x, y, sw), validation_split=0.2))

    self.assertEqual(int(train_x.shape[0]), 4)
    self.assertEqual(int(train_y.shape[0]), 4)
    self.assertEqual(int(train_sw.shape[0]), 4)
    for i in range(4):
      # Check that all arrays were shuffled in identical order.
      self.assertEqual(2 * train_x[i].numpy(), train_y[i].numpy())
      self.assertEqual(2 * train_y[i].numpy(), train_sw[i].numpy())

    self.assertEqual(int(val_x.shape[0]), 1)
    self.assertEqual(int(val_y.shape[0]), 1)
    self.assertEqual(int(val_sw.shape[0]), 1)
    for i in range(1):
      # Check that all arrays were shuffled in identical order.
      self.assertEqual(2 * train_x[i].numpy(), train_y[i].numpy())
      self.assertEqual(2 * train_y[i].numpy(), train_sw[i].numpy())

    # Check that arrays contain expected values.
    self.assertEqual(
        sorted(array_ops.concat([train_x, val_x], axis=0).numpy().tolist()),
        sorted(ops.convert_to_tensor_v2(x).numpy().tolist()))
    self.assertEqual(
        sorted(array_ops.concat([train_y, val_y], axis=0).numpy().tolist()),
        sorted(ops.convert_to_tensor_v2(y).numpy().tolist()))
    self.assertEqual(
        sorted(array_ops.concat([train_sw, val_sw], axis=0).numpy().tolist()),
        sorted(ops.convert_to_tensor_v2(sw).numpy().tolist()))

  @parameterized.named_parameters(('numpy_arrays', True), ('tensors', False))
  def test_validation_split_unshuffled(self, use_numpy):
    if use_numpy:
      x = np.array([0, 1, 2, 3, 4])
      y = np.array([0, 2, 4, 6, 8])
      sw = np.array([0, 4, 8, 12, 16])
    else:
      x = ops.convert_to_tensor_v2([0, 1, 2, 3, 4])
      y = ops.convert_to_tensor_v2([0, 2, 4, 6, 8])
      sw = ops.convert_to_tensor_v2([0, 4, 8, 12, 16])

    (train_x, train_y, train_sw), (val_x, val_y, val_sw) = (
        data_adapter.train_validation_split((x, y, sw),
                                            validation_split=0.2,
                                            shuffle=False))

    self.assertEqual(train_x.numpy().tolist(), [0, 1, 2, 3])
    self.assertEqual(train_y.numpy().tolist(), [0, 2, 4, 6])
    self.assertEqual(train_sw.numpy().tolist(), [0, 4, 8, 12])

    self.assertEqual(val_x.numpy().tolist(), [4])
    self.assertEqual(val_y.numpy().tolist(), [8])
    self.assertEqual(val_sw.numpy().tolist(), [16])

  def test_validation_split_user_error(self):
    with self.assertRaisesRegexp(ValueError, 'is only supported for Tensors'):
      data_adapter.train_validation_split(
          lambda: np.ones((10, 1)), validation_split=0.2)

  def test_validation_split_examples_too_few(self):
    with self.assertRaisesRegexp(
        ValueError, 'not sufficient to split it'):
      data_adapter.train_validation_split(
          np.ones((1, 10)), validation_split=0.2)

  def test_validation_split_none(self):
    train_sw, val_sw = data_adapter.train_validation_split(
        None, validation_split=0.2)
    self.assertIsNone(train_sw)
    self.assertIsNone(val_sw)

    (_, train_sw), (_, val_sw) = data_adapter.train_validation_split(
        (np.ones((10, 1)), None), validation_split=0.2)
    self.assertIsNone(train_sw)
    self.assertIsNone(val_sw)

class TestUtils(keras_parameterized.TestCase):

  def test_expand_1d_sparse_tensors_untouched(self):
    st = sparse_tensor.SparseTensor(
        indices=[[0], [10]], values=[1, 2], dense_shape=[10])
    st = data_adapter.expand_1d(st)
    self.assertEqual(st.shape.rank, 1)

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


#  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.
"""Example of DNNClassifier for Iris plant dataset, with early stopping."""

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

import shutil

from sklearn import datasets
from sklearn import metrics
from sklearn.cross_validation import train_test_split
import tensorflow as tf

learn = tf.contrib.learn

def clean_folder(folder):
  """Cleans the given folder if it exists."""
  try:
    shutil.rmtree(folder)
  except OSError:
    pass

def main(unused_argv):
  iris = datasets.load_iris()
  x_train, x_test, y_train, y_test = train_test_split(
      iris.data, iris.target, test_size=0.2, random_state=42)

  x_train, x_val, y_train, y_val = train_test_split(
      x_train, y_train, test_size=0.2, random_state=42)
  val_monitor = learn.monitors.ValidationMonitor(
      x_val, y_val, early_stopping_rounds=200)

  model_dir = '/tmp/iris_model'
  clean_folder(model_dir)

  # classifier with early stopping on training data
  classifier1 = learn.DNNClassifier(
      feature_columns=learn.infer_real_valued_columns_from_input(x_train),
      hidden_units=[10, 20, 10],
      n_classes=3,
      model_dir=model_dir)
  classifier1.fit(x=x_train, y=y_train, steps=2000)
  predictions1 = list(classifier1.predict(x_test, as_iterable=True))
  score1 = metrics.accuracy_score(y_test, predictions1)

  model_dir = '/tmp/iris_model_val'
  clean_folder(model_dir)

  # classifier with early stopping on validation data, save frequently for
  # monitor to pick up new checkpoints.
  classifier2 = learn.DNNClassifier(
      feature_columns=learn.infer_real_valued_columns_from_input(x_train),
      hidden_units=[10, 20, 10],
      n_classes=3,
      model_dir=model_dir,
      config=tf.contrib.learn.RunConfig(save_checkpoints_secs=1))
  classifier2.fit(x=x_train, y=y_train, steps=2000, monitors=[val_monitor])
  predictions2 = list(classifier2.predict(x_test, as_iterable=True))
  score2 = metrics.accuracy_score(y_test, predictions2)

  # In many applications, the score is improved by using early stopping
  print('score1: ', score1)
  print('score2: ', score2)
  print('score2 > score1: ', score2 > score1)

if __name__ == '__main__':
  tf.app.run()


"""
===================================
Compare cross decomposition methods
===================================

Simple usage of various cross decomposition algorithms:
- PLSCanonical
- PLSRegression, with multivariate response, a.k.a. PLS2
- PLSRegression, with univariate response, a.k.a. PLS1
- CCA

Given 2 multivariate covarying two-dimensional datasets, X, and Y,
PLS extracts the 'directions of covariance', i.e. the components of each
datasets that explain the most shared variance between both datasets.
This is apparent on the **scatterplot matrix** display: components 1 in
dataset X and dataset Y are maximally correlated (points lie around the
first diagonal). This is also true for components 2 in both dataset,
however, the correlation across datasets for different components is
weak: the point cloud is very spherical.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cross_decomposition import PLSCanonical, PLSRegression, CCA

###############################################################################
# Dataset based latent variables model

n = 500
# 2 latents vars:
l1 = np.random.normal(size=n)
l2 = np.random.normal(size=n)

latents = np.array([l1, l1, l2, l2]).T
X = latents + np.random.normal(size=4 * n).reshape((n, 4))
Y = latents + np.random.normal(size=4 * n).reshape((n, 4))

X_train = X[:n / 2]
Y_train = Y[:n / 2]
X_test = X[n / 2:]
Y_test = Y[n / 2:]

print("Corr(X)")
print(np.round(np.corrcoef(X.T), 2))
print("Corr(Y)")
print(np.round(np.corrcoef(Y.T), 2))

###############################################################################
# Canonical (symmetric) PLS

# Transform data
# ~~~~~~~~~~~~~~
plsca = PLSCanonical(n_components=2)
plsca.fit(X_train, Y_train)
X_train_r, Y_train_r = plsca.transform(X_train, Y_train)
X_test_r, Y_test_r = plsca.transform(X_test, Y_test)

# Scatter plot of scores
# ~~~~~~~~~~~~~~~~~~~~~~
# 1) On diagonal plot X vs Y scores on each components
plt.figure(figsize=(12, 8))
plt.subplot(221)
plt.plot(X_train_r[:, 0], Y_train_r[:, 0], "ob", label="train")
plt.plot(X_test_r[:, 0], Y_test_r[:, 0], "or", label="test")
plt.xlabel("x scores")
plt.ylabel("y scores")
plt.title('Comp. 1: X vs Y (test corr = %.2f)' %
          np.corrcoef(X_test_r[:, 0], Y_test_r[:, 0])[0, 1])
plt.xticks(())
plt.yticks(())
plt.legend(loc="best")

plt.subplot(224)
plt.plot(X_train_r[:, 1], Y_train_r[:, 1], "ob", label="train")
plt.plot(X_test_r[:, 1], Y_test_r[:, 1], "or", label="test")
plt.xlabel("x scores")
plt.ylabel("y scores")
plt.title('Comp. 2: X vs Y (test corr = %.2f)' %
          np.corrcoef(X_test_r[:, 1], Y_test_r[:, 1])[0, 1])
plt.xticks(())
plt.yticks(())
plt.legend(loc="best")

# 2) Off diagonal plot components 1 vs 2 for X and Y
plt.subplot(222)
plt.plot(X_train_r[:, 0], X_train_r[:, 1], "*b", label="train")
plt.plot(X_test_r[:, 0], X_test_r[:, 1], "*r", label="test")
plt.xlabel("X comp. 1")
plt.ylabel("X comp. 2")
plt.title('X comp. 1 vs X comp. 2 (test corr = %.2f)'
          % np.corrcoef(X_test_r[:, 0], X_test_r[:, 1])[0, 1])
plt.legend(loc="best")
plt.xticks(())
plt.yticks(())

plt.subplot(223)
plt.plot(Y_train_r[:, 0], Y_train_r[:, 1], "*b", label="train")
plt.plot(Y_test_r[:, 0], Y_test_r[:, 1], "*r", label="test")
plt.xlabel("Y comp. 1")
plt.ylabel("Y comp. 2")
plt.title('Y comp. 1 vs Y comp. 2 , (test corr = %.2f)'
          % np.corrcoef(Y_test_r[:, 0], Y_test_r[:, 1])[0, 1])
plt.legend(loc="best")
plt.xticks(())
plt.yticks(())
plt.show()

###############################################################################
# PLS regression, with multivariate response, a.k.a. PLS2

n = 1000
q = 3
p = 10
X = np.random.normal(size=n * p).reshape((n, p))
B = np.array([[1, 2] + [0] * (p - 2)] * q).T
# each Yj = 1*X1 + 2*X2 + noize
Y = np.dot(X, B) + np.random.normal(size=n * q).reshape((n, q)) + 5

pls2 = PLSRegression(n_components=3)
pls2.fit(X, Y)
print("True B (such that: Y = XB + Err)")
print(B)
# compare pls2.coefs with B
print("Estimated B")
print(np.round(pls2.coefs, 1))
pls2.predict(X)

###############################################################################
# PLS regression, with univariate response, a.k.a. PLS1

n = 1000
p = 10
X = np.random.normal(size=n * p).reshape((n, p))
y = X[:, 0] + 2 * X[:, 1] + np.random.normal(size=n * 1) + 5
pls1 = PLSRegression(n_components=3)
pls1.fit(X, y)
# note that the number of compements exceeds 1 (the dimension of y)
print("Estimated betas")
print(np.round(pls1.coefs, 1))

###############################################################################
# CCA (PLS mode B with symmetric deflation)

cca = CCA(n_components=2)
cca.fit(X_train, Y_train)
X_train_r, Y_train_r = plsca.transform(X_train, Y_train)
X_test_r, Y_test_r = plsca.transform(X_test, Y_test)


"""
Load an analysis file and redo the predictions on the validation set / test set,
this time with augmented data and averaging. Store them as numpy files.
"""

import numpy as np
# import pandas as pd
import theano
import theano.tensor as T
import layers
import cc_layers
import custom
import load_data
import realtime_augmentation as ra
import time
import csv
import os
import cPickle as pickle

BATCH_SIZE = 32 # 16
NUM_INPUT_FEATURES = 3

CHUNK_SIZE = 8000 # 10000 # this should be a multiple of the batch size

# ANALYSIS_PATH = "analysis/try_convnet_cc_multirot_3x69r45_untied_bias.pkl"
ANALYSIS_PATH = "analysis/final/try_convnet_cc_multirotflip_3x69r45_maxout2048_pysex.pkl"

DO_VALID = True # disable this to not bother with the validation set evaluation
DO_TEST = True # disable this to not generate predictions on the testset

target_filename = os.path.basename(ANALYSIS_PATH).replace(".pkl", ".npy.gz")
target_path_valid = os.path.join("predictions/final/augmented/valid", target_filename)
target_path_test = os.path.join("predictions/final/augmented/test", target_filename)

print "Loading model data etc."
analysis = np.load(ANALYSIS_PATH)

input_sizes = [(69, 69), (69, 69)]

ds_transforms = [
    ra.build_ds_transform(3.0, target_size=input_sizes[0]),
    ra.build_ds_transform(3.0, target_size=input_sizes[1]) + ra.build_augmentation_transform(rotation=45)]

num_input_representations = len(ds_transforms)

# split training data into training + a small validation set
num_train = load_data.num_train
num_valid = num_train // 10 # integer division
num_train -= num_valid
num_test = load_data.num_test

valid_ids = load_data.train_ids[num_train:]
train_ids = load_data.train_ids[:num_train]
test_ids = load_data.test_ids

train_indices = np.arange(num_train)
valid_indices = np.arange(num_train, num_train+num_valid)
test_indices = np.arange(num_test)

y_valid = np.load("data/solutions_train.npy")[num_train:]

print "Build model"
l0 = layers.Input2DLayer(BATCH_SIZE, NUM_INPUT_FEATURES, input_sizes[0][0], input_sizes[0][1])
l0_45 = layers.Input2DLayer(BATCH_SIZE, NUM_INPUT_FEATURES, input_sizes[1][0], input_sizes[1][1])

l0r = layers.MultiRotSliceLayer([l0, l0_45], part_size=45, include_flip=True)

l0s = cc_layers.ShuffleBC01ToC01BLayer(l0r)

l1a = cc_layers.CudaConvnetConv2DLayer(l0s, n_filters=32, filter_size=6, weights_std=0.01, init_bias_value=0.1, dropout=0.0, partial_sum=1, untie_biases=True)
l1 = cc_layers.CudaConvnetPooling2DLayer(l1a, pool_size=2)

l2a = cc_layers.CudaConvnetConv2DLayer(l1, n_filters=64, filter_size=5, weights_std=0.01, init_bias_value=0.1, dropout=0.0, partial_sum=1, untie_biases=True)
l2 = cc_layers.CudaConvnetPooling2DLayer(l2a, pool_size=2)

l3a = cc_layers.CudaConvnetConv2DLayer(l2, n_filters=128, filter_size=3, weights_std=0.01, init_bias_value=0.1, dropout=0.0, partial_sum=1, untie_biases=True)
l3b = cc_layers.CudaConvnetConv2DLayer(l3a, n_filters=128, filter_size=3, pad=0, weights_std=0.1, init_bias_value=0.1, dropout=0.0, partial_sum=1, untie_biases=True)
l3 = cc_layers.CudaConvnetPooling2DLayer(l3b, pool_size=2)

l3s = cc_layers.ShuffleC01BToBC01Layer(l3)

j3 = layers.MultiRotMergeLayer(l3s, num_views=4) # 2) # merge convolutional parts

# l4 = layers.DenseLayer(j3, n_outputs=4096, weights_std=0.001, init_bias_value=0.01, dropout=0.5)

l4a = layers.DenseLayer(j3, n_outputs=4096, weights_std=0.001, init_bias_value=0.01, dropout=0.5, nonlinearity=layers.identity)
l4 = layers.FeatureMaxPoolingLayer(l4a, pool_size=2, feature_dim=1, implementation='reshape')

# l5 = layers.DenseLayer(l4, n_outputs=37, weights_std=0.01, init_bias_value=0.0, dropout=0.5, nonlinearity=custom.clip_01) #  nonlinearity=layers.identity)
l5 = layers.DenseLayer(l4, n_outputs=37, weights_std=0.01, init_bias_value=0.1, dropout=0.5, nonlinearity=layers.identity)

# l6 = layers.OutputLayer(l5, error_measure='mse')
l6 = custom.OptimisedDivGalaxyOutputLayer(l5) # this incorporates the constraints on the output (probabilities sum to one, weighting, etc.)

xs_shared = [theano.shared(np.zeros((1,1,1,1), dtype=theano.config.floatX)) for _ in xrange(num_input_representations)]

idx = T.lscalar('idx')

givens = {
    l0.input_var: xs_shared[0][idx*BATCH_SIZE:(idx+1)*BATCH_SIZE],
    l0_45.input_var: xs_shared[1][idx*BATCH_SIZE:(idx+1)*BATCH_SIZE],
}

compute_output = theano.function([idx], l6.predictions(dropout_active=False), givens=givens)

print "Load model parameters"
layers.set_param_values(l6, analysis['param_values'])

print "Create generators"
# set here which transforms to use to make predictions
augmentation_transforms = []
for zoom in [1 / 1.2, 1.0, 1.2]:
    for angle in np.linspace(0, 360, 10, endpoint=False):
        augmentation_transforms.append(ra.build_augmentation_transform(rotation=angle, zoom=zoom))
        augmentation_transforms.append(ra.build_augmentation_transform(rotation=(angle + 180), zoom=zoom, shear=180)) # flipped

print "  %d augmentation transforms." % len(augmentation_transforms)

augmented_data_gen_valid = ra.realtime_fixed_augmented_data_gen(valid_indices, 'train', augmentation_transforms=augmentation_transforms, chunk_size=CHUNK_SIZE, target_sizes=input_sizes, ds_transforms=ds_transforms, processor_class=ra.LoadAndProcessFixedPysexCenteringRescaling)
valid_gen = load_data.buffered_gen_mp(augmented_data_gen_valid, buffer_size=1)

augmented_data_gen_test = ra.realtime_fixed_augmented_data_gen(test_indices, 'test', augmentation_transforms=augmentation_transforms, chunk_size=CHUNK_SIZE, target_sizes=input_sizes, ds_transforms=ds_transforms, processor_class=ra.LoadAndProcessFixedPysexCenteringRescaling)
test_gen = load_data.buffered_gen_mp(augmented_data_gen_test, buffer_size=1)

approx_num_chunks_valid = int(np.ceil(num_valid * len(augmentation_transforms) / float(CHUNK_SIZE)))
approx_num_chunks_test = int(np.ceil(num_test * len(augmentation_transforms) / float(CHUNK_SIZE)))

print "Approximately %d chunks for the validation set" % approx_num_chunks_valid
print "Approximately %d chunks for the test set" % approx_num_chunks_test

if DO_VALID:
    print
    print "VALIDATION SET"
    print "Compute predictions"
    predictions_list = []
    start_time = time.time()

    for e, (chunk_data, chunk_length) in enumerate(valid_gen):
        print "Chunk %d" % (e + 1)
        xs_chunk = chunk_data

        # need to transpose the chunks to move the 'channels' dimension up
        xs_chunk = [x_chunk.transpose(0, 3, 1, 2) for x_chunk in xs_chunk]

        print "  load data onto GPU"
        for x_shared, x_chunk in zip(xs_shared, xs_chunk):
            x_shared.set_value(x_chunk)
        num_batches_chunk = int(np.ceil(chunk_length / float(BATCH_SIZE)))

        # make predictions, don't forget to cute off the zeros at the end
        predictions_chunk_list = []
        for b in xrange(num_batches_chunk):
            if b % 1000 == 0:
                print "  batch %d/%d" % (b + 1, num_batches_chunk)

            predictions = compute_output(b)
            predictions_chunk_list.append(predictions)

        predictions_chunk = np.vstack(predictions_chunk_list)
        predictions_chunk = predictions_chunk[:chunk_length] # cut off zeros / padding

        print "  compute average over transforms"
        predictions_chunk_avg = predictions_chunk.reshape(-1, len(augmentation_transforms), 37).mean(1)

        predictions_list.append(predictions_chunk_avg)

        time_since_start = time.time() - start_time
        print "  %s since start" % load_data.hms(time_since_start)

    all_predictions = np.vstack(predictions_list)

    print "Write predictions to %s" % target_path_valid
    load_data.save_gz(target_path_valid, all_predictions)

    print "Evaluate"
    rmse_valid = analysis['losses_valid'][-1]
    rmse_augmented = np.sqrt(np.mean((y_valid - all_predictions)**2))
    print "  MSE (last iteration):\t%.6f" % rmse_valid
    print "  MSE (augmented):\t%.6f" % rmse_augmented

if DO_TEST:
    print
    print "TEST SET"
    print "Compute predictions"
    predictions_list = []
    start_time = time.time()

    for e, (chunk_data, chunk_length) in enumerate(test_gen):
        print "Chunk %d" % (e + 1)
        xs_chunk = chunk_data

        # need to transpose the chunks to move the 'channels' dimension up
        xs_chunk = [x_chunk.transpose(0, 3, 1, 2) for x_chunk in xs_chunk]

        print "  load data onto GPU"
        for x_shared, x_chunk in zip(xs_shared, xs_chunk):
            x_shared.set_value(x_chunk)
        num_batches_chunk = int(np.ceil(chunk_length / float(BATCH_SIZE)))

        # make predictions, don't forget to cute off the zeros at the end
        predictions_chunk_list = []
        for b in xrange(num_batches_chunk):
            if b % 1000 == 0:
                print "  batch %d/%d" % (b + 1, num_batches_chunk)

            predictions = compute_output(b)
            predictions_chunk_list.append(predictions)

        predictions_chunk = np.vstack(predictions_chunk_list)
        predictions_chunk = predictions_chunk[:chunk_length] # cut off zeros / padding

        print "  compute average over transforms"
        predictions_chunk_avg = predictions_chunk.reshape(-1, len(augmentation_transforms), 37).mean(1)

        predictions_list.append(predictions_chunk_avg)

        time_since_start = time.time() - start_time
        print "  %s since start" % load_data.hms(time_since_start)

    all_predictions = np.vstack(predictions_list)

    print "Write predictions to %s" % target_path_test
    load_data.save_gz(target_path_test, all_predictions)

    print "Done!"


# overlap coefficient join
from joblib import delayed, Parallel
from six import iteritems
import pandas as pd
import pyprind

from py_stringsimjoin.filter.overlap_filter import OverlapFilter
from py_stringsimjoin.index.inverted_index import InvertedIndex
from py_stringsimjoin.utils.generic_helper import convert_dataframe_to_array, \
    find_output_attribute_indices, get_attrs_to_project, \
    get_num_processes_to_launch, get_output_header_from_tables, \
    get_output_row_from_tables, remove_redundant_attrs, split_table, COMP_OP_MAP
from py_stringsimjoin.utils.missing_value_handler import \
    get_pairs_with_missing_value
from py_stringsimjoin.utils.validation import validate_attr, \
    validate_attr_type, validate_comp_op_for_sim_measure, validate_key_attr, \
    validate_input_table, validate_threshold, validate_tokenizer, \
    validate_output_attrs

def overlap_coefficient_join_py(ltable, rtable,
                                l_key_attr, r_key_attr,
                                l_join_attr, r_join_attr,
                                tokenizer, threshold, comp_op='>=',
                                allow_empty=True, allow_missing=False,
                                l_out_attrs=None, r_out_attrs=None,
                                l_out_prefix='l_', r_out_prefix='r_',
                                out_sim_score=True, n_jobs=1, show_progress=True):
    """Join two tables using overlap coefficient.

    For two sets X and Y, the overlap coefficient between them is given by:

        :math:`overlap\\_coefficient(X, Y) = \\frac{|X \\cap Y|}{\\min(|X|, |Y|)}`

    In the case where one of X and Y is an empty set and the other is a
    non-empty set, we define their overlap coefficient to be 0. In the case
    where both X and Y are empty sets, we define their overlap coefficient to
    be 1.

    Finds tuple pairs from left table and right table such that the overlap
    coefficient between the join attributes satisfies the condition on input
    threshold. For example, if the comparison operator is '>=', finds tuple
    pairs whose overlap coefficient between the strings that are the values of
    the join attributes is greater than or equal to the input threshold, as
    specified in "threshold".

    Args:
        ltable (DataFrame): left input table.

        rtable (DataFrame): right input table.

        l_key_attr (string): key attribute in left table.

        r_key_attr (string): key attribute in right table.

        l_join_attr (string): join attribute in left table.

        r_join_attr (string): join attribute in right table.

        tokenizer (Tokenizer): tokenizer to be used to tokenize join
            attributes.

        threshold (float): overlap coefficient threshold to be satisfied.

        comp_op (string): comparison operator. Supported values are '>=', '>'
            and '=' (defaults to '>=').

        allow_empty (boolean): flag to indicate whether tuple pairs with empty
            set of tokens in both the join attributes should be included in the
            output (defaults to True).

        allow_missing (boolean): flag to indicate whether tuple pairs with
            missing value in at least one of the join attributes should be
            included in the output (defaults to False). If this flag is set to
            True, a tuple in ltable with missing value in the join attribute
            will be matched with every tuple in rtable and vice versa.

        l_out_attrs (list): list of attribute names from the left table to be
            included in the output table (defaults to None).

        r_out_attrs (list): list of attribute names from the right table to be
            included in the output table (defaults to None).

        l_out_prefix (string): prefix to be used for the attribute names coming
            from the left table, in the output table (defaults to 'l\_').

        r_out_prefix (string): prefix to be used for the attribute names coming
            from the right table, in the output table (defaults to 'r\_').

        out_sim_score (boolean): flag to indicate whether similarity score
            should be included in the output table (defaults to True). Setting
            this flag to True will add a column named '_sim_score' in the
            output table. This column will contain the similarity scores for the
            tuple pairs in the output.

        n_jobs (int): number of parallel jobs to use for the computation
            (defaults to 1). If -1 is given, all CPUs are used. If 1 is given,
            no parallel computing code is used at all, which is useful for
            debugging. For n_jobs below -1, (n_cpus + 1 + n_jobs) are used
            (where n_cpus is the total number of CPUs in the machine). Thus for
            n_jobs = -2, all CPUs but one are used. If (n_cpus + 1 + n_jobs)
            becomes less than 1, then no parallel computing code will be used
            (i.e., equivalent to the default).

        show_progress (boolean): flag to indicate whether task progress should
            be displayed to the user (defaults to True).

    Returns:
        An output table containing tuple pairs that satisfy the join
        condition (DataFrame).
    """

    # check if the input tables are dataframes
    validate_input_table(ltable, 'left table')
    validate_input_table(rtable, 'right table')

    # check if the key attributes and join attributes exist
    validate_attr(l_key_attr, ltable.columns,
                  'key attribute', 'left table')
    validate_attr(r_key_attr, rtable.columns,
                  'key attribute', 'right table')
    validate_attr(l_join_attr, ltable.columns,
                  'join attribute', 'left table')
    validate_attr(r_join_attr, rtable.columns,
                  'join attribute', 'right table')

    # check if the join attributes are not of numeric type
    validate_attr_type(l_join_attr, ltable[l_join_attr].dtype,
                       'join attribute', 'left table')
    validate_attr_type(r_join_attr, rtable[r_join_attr].dtype,
                       'join attribute', 'right table')

    # check if the input tokenizer is valid
    validate_tokenizer(tokenizer)

    # check if the input threshold is valid
    validate_threshold(threshold, 'OVERLAP_COEFFICIENT')

    # check if the comparison operator is valid
    validate_comp_op_for_sim_measure(comp_op, 'OVERLAP_COEFFICIENT')

    # check if the output attributes exist
    validate_output_attrs(l_out_attrs, ltable.columns,
                          r_out_attrs, rtable.columns)

    # check if the key attributes are unique and do not contain missing values
    validate_key_attr(l_key_attr, ltable, 'left table')
    validate_key_attr(r_key_attr, rtable, 'right table')

    # set return_set flag of tokenizer to be True, in case it is set to False
    revert_tokenizer_return_set_flag = False
    if not tokenizer.get_return_set():
        tokenizer.set_return_set(True)
        revert_tokenizer_return_set_flag = True

    # remove redundant attrs from output attrs.
    l_out_attrs = remove_redundant_attrs(l_out_attrs, l_key_attr)
    r_out_attrs = remove_redundant_attrs(r_out_attrs, r_key_attr)

    # get attributes to project.
    l_proj_attrs = get_attrs_to_project(l_out_attrs, l_key_attr, l_join_attr)
    r_proj_attrs = get_attrs_to_project(r_out_attrs, r_key_attr, r_join_attr)

    # Do a projection on the input dataframes to keep only the required
    # attributes. Then, remove rows with missing value in join attribute from
    # the input dataframes. Then, convert the resulting dataframes into ndarray.
    ltable_array = convert_dataframe_to_array(ltable, l_proj_attrs, l_join_attr)
    rtable_array = convert_dataframe_to_array(rtable, r_proj_attrs, r_join_attr)

    # computes the actual number of jobs to launch.
    n_jobs = min(get_num_processes_to_launch(n_jobs), len(rtable_array))

    if n_jobs <= 1:
        # if n_jobs is 1, do not use any parallel code.
        output_table = _overlap_coefficient_join_split(
                                           ltable_array, rtable_array,
                                           l_proj_attrs, r_proj_attrs,
                                           l_key_attr, r_key_attr,
                                           l_join_attr, r_join_attr,
                                           tokenizer, threshold, comp_op,
                                           allow_empty,
                                           l_out_attrs, r_out_attrs,
                                           l_out_prefix, r_out_prefix,
                                           out_sim_score, show_progress)
    else:
        # if n_jobs is above 1, split the right table into n_jobs splits and
        # join each right table split with the whole of left table in a separate
        # process.
        r_splits = split_table(rtable_array, n_jobs)
        results = Parallel(n_jobs=n_jobs)(
                                delayed(_overlap_coefficient_join_split)(
                                    ltable_array, r_splits[job_index],
                                    l_proj_attrs, r_proj_attrs,
                                    l_key_attr, r_key_attr,
                                    l_join_attr, r_join_attr,
                                    tokenizer, threshold, comp_op,
                                    allow_empty,
                                    l_out_attrs, r_out_attrs,
                                    l_out_prefix, r_out_prefix,
                                    out_sim_score,
                                    (show_progress and (job_index==n_jobs-1)))
                                for job_index in range(n_jobs))
        output_table = pd.concat(results)

    # If allow_missing flag is set, then compute all pairs with missing value in
    # at least one of the join attributes and then add it to the output
    # obtained from the join.
    if allow_missing:
        missing_pairs = get_pairs_with_missing_value(
                                            ltable, rtable,
                                            l_key_attr, r_key_attr,
                                            l_join_attr, r_join_attr,
                                            l_out_attrs, r_out_attrs,
                                            l_out_prefix, r_out_prefix,
                                            out_sim_score, show_progress)
        output_table = pd.concat([output_table, missing_pairs])

    # add an id column named '_id' to the output table.
    output_table.insert(0, '_id', range(0, len(output_table)))

    # revert the return_set flag of tokenizer, in case it was modified.
    if revert_tokenizer_return_set_flag:
        tokenizer.set_return_set(False)

    return output_table

def _overlap_coefficient_join_split(ltable_list, rtable_list,
                                    l_columns, r_columns,
                                    l_key_attr, r_key_attr,
                                    l_join_attr, r_join_attr,
                                    tokenizer, threshold, comp_op,
                                    allow_empty,
                                    l_out_attrs, r_out_attrs,
                                    l_out_prefix, r_out_prefix,
                                    out_sim_score, show_progress):
    """Perform overlap coefficient join for a split of ltable and rtable"""
    # find column indices of key attr, join attr and output attrs in ltable
    l_key_attr_index = l_columns.index(l_key_attr)
    l_join_attr_index = l_columns.index(l_join_attr)
    l_out_attrs_indices = find_output_attribute_indices(l_columns, l_out_attrs)

    # find column indices of key attr, join attr and output attrs in rtable
    r_key_attr_index = r_columns.index(r_key_attr)
    r_join_attr_index = r_columns.index(r_join_attr)
    r_out_attrs_indices = find_output_attribute_indices(r_columns, r_out_attrs)

    # Build inverted index over ltable
    inverted_index = InvertedIndex(ltable_list, l_join_attr_index,
                                   tokenizer, cache_size_flag=True)
    # While building the index, we cache the record ids with empty set of
    # tokens. This is needed to handle the allow_empty flag.
    cached_data = inverted_index.build(allow_empty)
    l_empty_records = cached_data['empty_records']

    overlap_filter = OverlapFilter(tokenizer, 1)
    comp_fn = COMP_OP_MAP[comp_op]

    output_rows = []
    has_output_attributes = (l_out_attrs is not None or
                             r_out_attrs is not None)

    if show_progress:
        prog_bar = pyprind.ProgBar(len(rtable_list))

    for r_row in rtable_list:
        r_string = r_row[r_join_attr_index]

        r_join_attr_tokens = tokenizer.tokenize(r_string)
        r_num_tokens = len(r_join_attr_tokens)

        # If allow_empty flag is set and the current rtable record has empty set
        # of tokens in the join attribute, then generate output pairs joining
        # the current rtable record with those records in ltable with empty set
        # of tokens in the join attribute. These ltable record ids are cached in
        # l_empty_records list which was constructed when building the inverted
        # index.
        if allow_empty and r_num_tokens == 0:
            for l_id in l_empty_records:
                if has_output_attributes:
                    output_row = get_output_row_from_tables(
                                     ltable_list[l_id], r_row,
                                     l_key_attr_index, r_key_attr_index,
                                     l_out_attrs_indices,
                                     r_out_attrs_indices)
                else:
                    output_row = [ltable_list[l_id][l_key_attr_index],
                                  r_row[r_key_attr_index]]

                if out_sim_score:
                    output_row.append(1.0)
                output_rows.append(output_row)
            continue

        # probe inverted index and find overlap of candidates
        candidate_overlap = overlap_filter.find_candidates(
                                r_join_attr_tokens, inverted_index)

        for cand, overlap in iteritems(candidate_overlap):
            # compute the actual similarity score
            sim_score = (float(overlap) /
                         float(min(r_num_tokens,
                                   inverted_index.size_cache[cand])))

            if comp_fn(sim_score, threshold):
                if has_output_attributes:
                    output_row = get_output_row_from_tables(
                                     ltable_list[cand], r_row,
                                     l_key_attr_index, r_key_attr_index,
                                     l_out_attrs_indices, r_out_attrs_indices)
                else:
                    output_row = [ltable_list[cand][l_key_attr_index],
                                  r_row[r_key_attr_index]]

                # if out_sim_score flag is set, append the overlap coefficient
                # score to the output record.
                if out_sim_score:
                    output_row.append(sim_score)

                output_rows.append(output_row)

        if show_progress:
            prog_bar.update()

    output_header = get_output_header_from_tables(l_key_attr, r_key_attr,
                                                  l_out_attrs, r_out_attrs,
                                                  l_out_prefix, r_out_prefix)
    if out_sim_score:
        output_header.append("_sim_score")

    output_table = pd.DataFrame(output_rows, columns=output_header)
    return output_table


import sys
import logging

import numpy as np
import matplotlib.pyplot as plt

import robo.models.neural_network as robo_net
import robo.models.bagged_networks as bn
from robo.initial_design.init_random_uniform import init_random_uniform

logging.basicConfig(stream=sys.stdout, level=logging.INFO)

def f(x):
    return np.sinc(x * 10 - 5).sum(axis=1)[:, None]

rng = np.random.RandomState(42)

X = init_random_uniform(np.zeros(1), np.ones(1), 20, rng).astype(np.float32)
Y = f(X)

x = np.linspace(0, 1, 512, dtype=np.float32)[:, None]
vals = f(x).astype(np.float32)

plt.grid()
plt.plot(x[:, 0], f(x)[:, 0], label="true", color="green")
plt.plot(X[:, 0], Y[:, 0], "ro")

model = bn.BaggedNets(robo_net.SGDNet, num_models=16, bootstrap_with_replacement=True,
                      n_epochs=16384, error_threshold=1e-3,
                      n_units=[32, 32, 32], dropout=0,
                      batch_size=10, learning_rate=1e-3,
                      shuffle_batches=True)

m = model.train(X, Y)

mean_pred, var_pred = model.predict(x)
std_pred = np.sqrt(var_pred)

plt.plot(x[:, 0], mean_pred[:, 0], label="bagged nets", color="blue")
plt.fill_between(x[:, 0], mean_pred[:, 0] + std_pred[:, 0], mean_pred[:, 0] - std_pred[:, 0], alpha=0.2, color="blue")

plt.legend()
plt.show()


"""
Testing for Clustering methods

"""

import numpy as np

from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_raises

from sklearn.cluster.affinity_propagation_ import AffinityPropagation
from sklearn.cluster.affinity_propagation_ import affinity_propagation
from sklearn.datasets.samples_generator import make_blobs
from sklearn.metrics import euclidean_distances

n_clusters = 3
centers = np.array([[1, 1], [-1, -1], [1, -1]]) + 10
X, _ = make_blobs(n_samples=60, n_features=2, centers=centers,
                  cluster_std=0.4, shuffle=True, random_state=0)

def test_affinity_propagation():
    # Affinity Propagation algorithm
    # Compute similarities
    S = -euclidean_distances(X, squared=True)
    preference = np.median(S) * 10
    # Compute Affinity Propagation
    cluster_centers_indices, labels = affinity_propagation(
        S, preference=preference)

    n_clusters_ = len(cluster_centers_indices)

    assert_equal(n_clusters, n_clusters_)

    af = AffinityPropagation(preference=preference, affinity="precomputed")
    labels_precomputed = af.fit(S).labels_

    af = AffinityPropagation(preference=preference, verbose=True)
    labels = af.fit(X).labels_

    assert_array_equal(labels, labels_precomputed)

    cluster_centers_indices = af.cluster_centers_indices_

    n_clusters_ = len(cluster_centers_indices)
    assert_equal(np.unique(labels).size, n_clusters_)
    assert_equal(n_clusters, n_clusters_)

    # Test also with no copy
    _, labels_no_copy = affinity_propagation(S, preference=preference,
                                             copy=False)
    assert_array_equal(labels, labels_no_copy)

    # Test input validation
    assert_raises(ValueError, affinity_propagation, S[:, :-1])
    assert_raises(ValueError, affinity_propagation, S, damping=0)
    af = AffinityPropagation(affinity="unknown")
    assert_raises(ValueError, af.fit, X)

def test_affinity_propagation_predict():
    # Test AffinityPropagation.predict
    af = AffinityPropagation(affinity="euclidean")
    labels = af.fit_predict(X)
    labels2 = af.predict(X)
    assert_array_equal(labels, labels2)

def test_affinity_propagation_predict_error():
    # Test exception in AffinityPropagation.predict
    # Not fitted.
    af = AffinityPropagation(affinity="euclidean")
    assert_raises(ValueError, af.predict, X)

    # Predict not supported when affinity="precomputed".
    S = np.dot(X, X.T)
    af = AffinityPropagation(affinity="precomputed")
    af.fit(S)
    assert_raises(ValueError, af.predict, X)


"""Common tests for metaestimators"""

import functools

import numpy as np

from sklearn.base import BaseEstimator
from sklearn.externals.six import iterkeys
from sklearn.datasets import make_classification
from sklearn.utils.testing import assert_true, assert_false, assert_raises
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
from sklearn.feature_selection import RFE, RFECV
from sklearn.ensemble import BaggingClassifier

class DelegatorData(object):
    def __init__(self, name, construct, skip_methods=(),
                 fit_args=make_classification()):
        self.name = name
        self.construct = construct
        self.fit_args = fit_args
        self.skip_methods = skip_methods

DELEGATING_METAESTIMATORS = [
    DelegatorData('Pipeline', lambda est: Pipeline([('est', est)])),
    DelegatorData('GridSearchCV',
                  lambda est: GridSearchCV(
                      est, param_grid={'param': [5]}, cv=2),
                  skip_methods=['score']),
    DelegatorData('RandomizedSearchCV',
                  lambda est: RandomizedSearchCV(
                      est, param_distributions={'param': [5]}, cv=2, n_iter=1),
                  skip_methods=['score']),
    DelegatorData('RFE', RFE,
                  skip_methods=['transform', 'inverse_transform', 'score']),
    DelegatorData('RFECV', RFECV,
                  skip_methods=['transform', 'inverse_transform', 'score']),
    DelegatorData('BaggingClassifier', BaggingClassifier,
                  skip_methods=['transform', 'inverse_transform', 'score',
                                'predict_proba', 'predict_log_proba', 'predict'])
]

def test_metaestimator_delegation():
    # Ensures specified metaestimators have methods iff subestimator does
    def hides(method):
        @property
        def wrapper(obj):
            if obj.hidden_method == method.__name__:
                raise AttributeError('%r is hidden' % obj.hidden_method)
            return functools.partial(method, obj)
        return wrapper

    class SubEstimator(BaseEstimator):
        def __init__(self, param=1, hidden_method=None):
            self.param = param
            self.hidden_method = hidden_method

        def fit(self, X, y=None, *args, **kwargs):
            self.coef_ = np.arange(X.shape[1])
            return True

        def _check_fit(self):
            if not hasattr(self, 'coef_'):
                raise RuntimeError('Estimator is not fit')

        @hides
        def inverse_transform(self, X, *args, **kwargs):
            self._check_fit()
            return X

        @hides
        def transform(self, X, *args, **kwargs):
            self._check_fit()
            return X

        @hides
        def predict(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def predict_proba(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def predict_log_proba(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def decision_function(self, X, *args, **kwargs):
            self._check_fit()
            return np.ones(X.shape[0])

        @hides
        def score(self, X, *args, **kwargs):
            self._check_fit()
            return 1.0

    methods = [k for k in iterkeys(SubEstimator.__dict__)
               if not k.startswith('_') and not k.startswith('fit')]
    methods.sort()

    for delegator_data in DELEGATING_METAESTIMATORS:
        delegate = SubEstimator()
        delegator = delegator_data.construct(delegate)
        for method in methods:
            if method in delegator_data.skip_methods:
                continue
            assert_true(hasattr(delegate, method))
            assert_true(hasattr(delegator, method),
                        msg="%s does not have method %r when its delegate does"
                            % (delegator_data.name, method))
            # delegation before fit raises an exception
            assert_raises(Exception, getattr(delegator, method),
                          delegator_data.fit_args[0])

        delegator.fit(*delegator_data.fit_args)
        for method in methods:
            if method in delegator_data.skip_methods:
                continue
            # smoke test delegation
            getattr(delegator, method)(delegator_data.fit_args[0])

        for method in methods:
            if method in delegator_data.skip_methods:
                continue
            delegate = SubEstimator(hidden_method=method)
            delegator = delegator_data.construct(delegate)
            assert_false(hasattr(delegate, method))
            assert_false(hasattr(delegator, method),
                         msg="%s has method %r when its delegate does not"
                             % (delegator_data.name, method))