# 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.
# ==============================================================================
"""Tests for the `SnapshotDataset` transformation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import multiprocessing
import os
import shutil
import time

from absl.testing import parameterized
import numpy as np

from tensorflow.python.data.experimental.kernel_tests import reader_dataset_ops_test_base
from tensorflow.python.data.experimental.ops import snapshot
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.framework import combinations
from tensorflow.python.framework import errors
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.platform import test


class SnapshotDatasetTest(reader_dataset_ops_test_base.TFRecordDatasetTestBase,
                          parameterized.TestCase):

  def setUp(self):
    super(SnapshotDatasetTest, self).setUp()
    tmpdir = self.get_temp_dir()
    tmpdir = os.path.join(tmpdir, "snapshot")
    os.mkdir(tmpdir)
    self._snapshot_dir = tmpdir

  def tearDown(self):
    super(SnapshotDatasetTest, self).tearDown()
    shutil.rmtree(self._snapshot_dir)

  def createTFRecords(self, num_files=10, num_records=100):
    self._num_files = num_files
    self._num_records = num_records
    self._test_filenames = self._createFiles()

  def removeTFRecords(self):
    for filename in self._test_filenames:
      os.remove(filename)
    self._test_filenames = []
    self._num_files = None
    self._num_records = None

  def assertDatasetProducesSet(self, dataset, expected):
    actual = []
    next_fn = self.getNext(dataset)
    for _ in range(len(expected)):
      elem = self.evaluate(next_fn())
      actual.append(elem)
    self.assertCountEqual(actual, expected)
    with self.assertRaises(errors.OutOfRangeError):
      self.evaluate(next_fn())

  def assertSnapshotDirectoryContains(self, directory, num_fingerprints,
                                      num_runs_per_fingerprint,
                                      num_snapshot_shards_per_run):
    dirlist_raw = os.listdir(directory)
    dirlist = []

    # Ignore the graphdef pbtxts we write for debugging purposes.
    for i in range(len(dirlist_raw)):
      if not dirlist_raw[i].endswith("-graph.pbtxt"):
        dirlist.append(dirlist_raw[i])

    self.assertLen(dirlist, num_fingerprints)

    for i in range(num_fingerprints):
      fingerprint_dir = os.path.join(directory, dirlist[i])
      fingerprint_dir_list = sorted(os.listdir(fingerprint_dir))
      self.assertLen(fingerprint_dir_list, num_runs_per_fingerprint + 1)
      self.assertEqual(fingerprint_dir_list[num_runs_per_fingerprint],
                       "snapshot.metadata")

      for j in range(num_runs_per_fingerprint):
        run_dir = os.path.join(fingerprint_dir, fingerprint_dir_list[j])
        run_dirlist = sorted(os.listdir(run_dir))
        self.assertLen(run_dirlist, num_snapshot_shards_per_run)

        file_counter = 0
        for filename in run_dirlist:
          self.assertEqual(filename, "%08d.shard" % file_counter)
          file_counter += 1

  @combinations.generate(test_base.default_test_combinations())
  def testCreateSnapshotDataset(self):
    dataset = dataset_ops.Dataset.from_tensors([1, 2, 3])
    dataset.apply(snapshot.snapshot(self._snapshot_dir))

  @combinations.generate(test_base.default_test_combinations())
  def testReadSnapshotDatasetDefault(self):
    self.createTFRecords()
    filenames = self._test_filenames
    expected = [
        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
        for f in range(0, 10)
        for r in range(0, 100)
    ]

    dataset = core_readers._TFRecordDataset(filenames)
    dataset = dataset.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset, expected)
    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

    self.removeTFRecords()
    dataset2 = core_readers._TFRecordDataset(filenames)
    dataset2 = dataset2.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset2, expected)

  @combinations.generate(test_base.default_test_combinations())
  def testReadSnapshotDatasetAutoWriteSnappyRead(self):
    self.createTFRecords()
    filenames = self._test_filenames
    expected = [
        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
        for f in range(0, 10)
        for r in range(0, 100)
    ]

    dataset = core_readers._TFRecordDataset(filenames)
    dataset = dataset.apply(
        snapshot.snapshot(self._snapshot_dir, compression="AUTO"))
    self.assertDatasetProduces(dataset, expected)

    self.removeTFRecords()
    dataset2 = core_readers._TFRecordDataset(filenames)
    dataset2 = dataset2.apply(
        snapshot.snapshot(self._snapshot_dir, compression="SNAPPY"))
    self.assertDatasetProduces(dataset2, expected)

  @combinations.generate(test_base.default_test_combinations())
  def testReadSnapshotDatasetCustomShardFn(self):
    self.createTFRecords()
    filenames = self._test_filenames
    expected = [
        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
        for f in range(0, 10)
        for r in range(0, 100)
    ]

    dataset = core_readers._TFRecordDataset(filenames)
    dataset = dataset.apply(
        snapshot.snapshot(self._snapshot_dir, shard_func=lambda _: np.int64(0)))
    self.assertDatasetProduces(dataset, expected)
    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=1)

    self.removeTFRecords()
    dataset2 = core_readers._TFRecordDataset(filenames)
    dataset2 = dataset2.apply(
        snapshot.snapshot(self._snapshot_dir, shard_func=lambda _: 0))
    self.assertDatasetProduces(dataset2, expected)

  @combinations.generate(test_base.default_test_combinations())
  def testReadSnapshotDatasetCustomReaderFn(self):
    self.createTFRecords()
    filenames = self._test_filenames
    expected = [
        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
        for f in range(0, 10)
        for r in range(0, 100)
    ]

    dataset = core_readers._TFRecordDataset(filenames)
    dataset = dataset.apply(
        snapshot.snapshot(
            self._snapshot_dir,
            reader_func=(
                lambda ds: ds.interleave(  # pylint:disable=g-long-lambda
                    lambda x: x,
                    cycle_length=4,
                    num_parallel_calls=4))))
    self.assertDatasetProduces(dataset, expected)
    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

    self.removeTFRecords()
    dataset2 = core_readers._TFRecordDataset(filenames)
    dataset2 = dataset2.apply(
        snapshot.snapshot(
            self._snapshot_dir,
            reader_func=(
                lambda ds: ds.interleave(  # pylint:disable=g-long-lambda
                    lambda x: x,
                    cycle_length=4,
                    num_parallel_calls=4))))
    self.assertDatasetProducesSet(dataset2, expected)

  @combinations.generate(test_base.default_test_combinations())
  def testSnapshotDatasetInvalidShardFn(self):
    dataset = dataset_ops.Dataset.range(1000)
    with self.assertRaises(TypeError):
      dataset = dataset.apply(
          snapshot.snapshot(
              self._snapshot_dir, shard_func=lambda _: "invalid_fn"))
      next_fn = self.getNext(dataset)
      self.evaluate(next_fn())

  @combinations.generate(test_base.default_test_combinations())
  def testSnapshotDatasetInvalidReaderFn(self):
    dataset = dataset_ops.Dataset.range(1000)
    with self.assertRaises(TypeError):
      dataset = dataset.apply(
          snapshot.snapshot(self._snapshot_dir, reader_func=lambda x: x + 1))
      next_fn = self.getNext(dataset)
      self.evaluate(next_fn())

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotDatasetSimple(self):
    dataset = dataset_ops.Dataset.range(1000)
    dataset = dataset.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset, list(range(1000)))
    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotDatasetMultipleFingerprints(self):
    dataset1 = dataset_ops.Dataset.range(1000)
    dataset1 = dataset1.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset1, list(range(1000)))

    dataset2 = dataset_ops.Dataset.range(2000)
    dataset2 = dataset2.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset2, list(range(2000)))

    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=2,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotDatasetSameFingerprintMultipleCompleteRuns(self):
    dataset1 = dataset_ops.Dataset.range(1000)
    dataset1 = dataset1.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset1, list(range(1000)))
    dataset2 = dataset_ops.Dataset.range(1000)
    dataset2 = dataset2.apply(snapshot.snapshot(self._snapshot_dir))
    self.assertDatasetProduces(dataset2, list(range(1000)))

    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotDatasetSameFingerprintIncompleteRunRestart(self):
    dataset1 = dataset_ops.Dataset.range(1000)
    dataset1 = dataset1.apply(snapshot.snapshot(self._snapshot_dir))
    next1 = self.getNext(dataset1)
    for i in range(500):
      self.assertEqual(i, self.evaluate(next1()))

    dataset2 = dataset_ops.Dataset.range(1000)
    dataset2 = dataset2.apply(snapshot.snapshot(self._snapshot_dir))
    next2 = self.getNext(dataset2)
    for i in range(500):
      self.assertEqual(i, self.evaluate(next2()))

    for i in range(500, 1000):
      self.assertEqual(i, self.evaluate(next1()))
      self.assertEqual(i, self.evaluate(next2()))

    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=2,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotCustomShardFunction(self):
    dataset = dataset_ops.Dataset.range(1000)
    dataset = dataset.enumerate()
    dataset = dataset.apply(
        snapshot.snapshot(self._snapshot_dir, shard_func=lambda i, _: i % 2))
    dataset = dataset.map(lambda _, elem: elem)
    self.assertDatasetProduces(dataset, list(range(1000)))
    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=2)

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotDatasetWithTuples(self):
    dataset1 = dataset_ops.Dataset.range(0, 1000)
    dataset2 = dataset_ops.Dataset.range(1000, 2000)
    dataset3 = dataset_ops.Dataset.range(2000, 3000)
    dataset4 = dataset_ops.Dataset.range(3000, 4000)

    dataset = dataset_ops.Dataset.zip((dataset1, dataset2, dataset3, dataset4))
    dataset = dataset.apply(snapshot.snapshot(self._snapshot_dir))
    next1 = self.getNext(dataset)
    for i in range(0, 1000):
      self.assertEqual((i, i + 1000, i + 2000, i + 3000),
                       self.evaluate(next1()))
    self.assertSnapshotDirectoryContains(
        self._snapshot_dir,
        num_fingerprints=1,
        num_runs_per_fingerprint=1,
        num_snapshot_shards_per_run=multiprocessing.cpu_count())

  @combinations.generate(test_base.default_test_combinations())
  def testWriteSnapshotShuffleSameFingerprint(self):

    def make_dataset():
      dataset = dataset_ops.Dataset.range(1000)
      dataset = dataset.shuffle(1000)
      dataset = dataset.apply(snapshot.snapshot(self._snapshot_dir))
      return dataset

    dataset1 = make_dataset()
    self.assertDatasetProducesSet(dataset1, list(range(1000)))
    dataset2 = make_dataset()