# Copyright 2018 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.
# ==============================================================================
"""Input generator for image data."""

import os
import lingvo.compat as tf
from lingvo.core import base_input_generator

from tensorflow.python.ops import io_ops


class _MnistInputBase(base_input_generator.BaseTinyDatasetInput):
  """Base input params for MNIST."""

  @classmethod
  def Params(cls):
    """Defaults params."""
    p = super().Params()
    p.data_dtype = tf.uint8
    p.data_shape = (28, 28, 1)
    p.label_dtype = tf.uint8
    return p

  def _Preprocess(self, raw):
    data = tf.stack([
        tf.image.per_image_standardization(img) for img in tf.unstack(raw)
    ])
    data.set_shape(raw.shape)
    return data


class MnistTrainInput(_MnistInputBase):
  """MNist training set."""

  @classmethod
  def Params(cls):
    """Defaults params."""
    p = super().Params()
    p.data = 'x_train'
    p.label = 'y_train'
    p.num_samples = 60000
    p.batch_size = 256
    p.repeat = True
    return p


class MnistTestInput(_MnistInputBase):
  """MNist test set."""

  @classmethod
  def Params(cls):
    """Defaults params."""
    p = super().Params()
    p.data = 'x_test'
    p.label = 'y_test'
    p.num_samples = 10000
    p.batch_size = 256
    p.repeat = False
    return p


def _GetRandomImages(batch_size):
  images = tf.random.uniform((batch_size, 28, 28, 1), 0, 255, tf.int32)
  return tf.cast(images, tf.uint8)


def _GetRandomLabels(batch_size):
  labels = tf.random.categorical(0.1 * tf.ones((1, 10)), batch_size)
  return tf.cast(labels, tf.uint8)


def FakeMnistData(tmpdir, train_size=60000, test_size=10000):
  """Fake Mnist data for unit tests."""
  data_path = os.path.join(tmpdir, 'ckpt')
  with tf.Graph().as_default():
    tf.random.set_seed(91)
    with tf.Session() as sess:
      sess.run(
          io_ops.save_v2(
              data_path,
              tensor_names=['x_train', 'y_train', 'x_test', 'y_test'],
              shape_and_slices=['', '', '', ''],
              tensors=[
                  _GetRandomImages(train_size),
                  _GetRandomLabels(train_size),
                  _GetRandomImages(test_size),
                  _GetRandomLabels(test_size)
              ]))
  return data_path
