# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for learn.dataframe.tensorflow_dataframe."""

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

import csv
import math
import tempfile

import numpy as np

from tensorflow.contrib.learn.python.learn.dataframe import tensorflow_dataframe as df
from tensorflow.contrib.learn.python.learn.dataframe.transforms import densify
from tensorflow.core.example import example_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.lib.io import tf_record
from tensorflow.python.ops import parsing_ops
from tensorflow.python.platform import test

# pylint: disable=g-import-not-at-top
try:
  import pandas as pd
  HAS_PANDAS = True
except ImportError:
  HAS_PANDAS = False


def _assert_df_equals_dict(expected_df, actual_dict):
  for col in expected_df:
    if expected_df[col].dtype in [np.float32, np.float64]:
      assertion = np.testing.assert_allclose
    else:
      assertion = np.testing.assert_array_equal

    if expected_df[col].dtype.kind in ["O", "S", "U"]:
      # Python 2/3 compatibility
      # TensorFlow always returns bytes, so we just convert the unicode
      # expectations to bytes also before comparing.
      expected_values = [x.encode("utf-8") for x in expected_df[col].values]
    else:
      expected_values = expected_df[col].values

    assertion(
        expected_values,
        actual_dict[col],
        err_msg="Expected {} in column '{}'; got {}.".format(expected_values,
                                                             col,
                                                             actual_dict[col]))


class TensorFlowDataFrameTestCase(test.TestCase):
  """Tests for `TensorFlowDataFrame`."""

  def _make_test_csv(self):
    f = tempfile.NamedTemporaryFile(
        dir=self.get_temp_dir(), delete=False, mode="w")
    w = csv.writer(f)
    w.writerow(["int", "float", "bool", "string"])
    for _ in range(100):
      intvalue = np.random.randint(-10, 10)
      floatvalue = np.random.rand()
      boolvalue = int(np.random.rand() > 0.3)
      stringvalue = "S: %.4f" % np.random.rand()

      row = [intvalue, floatvalue, boolvalue, stringvalue]
      w.writerow(row)
    f.close()
    return f.name

  def _make_test_csv_sparse(self):
    f = tempfile.NamedTemporaryFile(
        dir=self.get_temp_dir(), delete=False, mode="w")
    w = csv.writer(f)
    w.writerow(["int", "float", "bool", "string"])
    for _ in range(100):
      # leave columns empty; these will be read as default value (e.g. 0 or NaN)
      intvalue = np.random.randint(-10, 10) if np.random.rand() > 0.5 else ""
      floatvalue = np.random.rand() if np.random.rand() > 0.5 else ""
      boolvalue = int(np.random.rand() > 0.3) if np.random.rand() > 0.5 else ""
      stringvalue = (
          ("S: %.4f" % np.random.rand()) if np.random.rand() > 0.5 else "")

      row = [intvalue, floatvalue, boolvalue, stringvalue]
      w.writerow(row)
    f.close()
    return f.name

  def _make_test_tfrecord(self):
    f = tempfile.NamedTemporaryFile(dir=self.get_temp_dir(), delete=False)
    w = tf_record.TFRecordWriter(f.name)
    for i in range(100):
      ex = example_pb2.Example()
      ex.features.feature["var_len_int"].int64_list.value.extend(range((i % 3)))
      ex.features.feature["fixed_len_float"].float_list.value.extend(
          [float(i), 2 * float(i)])
      w.write(ex.SerializeToString())
    return f.name

  def _assert_pandas_equals_tensorflow(self, pandas_df, tensorflow_df,
                                       num_batches, batch_size):
    self.assertItemsEqual(
        list(pandas_df.columns) + ["index"], tensorflow_df.columns())

    for batch_num, batch in enumerate(tensorflow_df.run(num_batches)):
      row_numbers = [
          total_row_num % pandas_df.shape[0]
          for total_row_num in range(batch_size * batch_num, batch_size * (
              batch_num + 1))
      ]
      expected_df = pandas_df.iloc[row_numbers]
      _assert_df_equals_dict(expected_df, batch)

  def testInitFromPandas(self):
    """Test construction from Pandas DataFrame."""
    if not HAS_PANDAS:
      return
    pandas_df = pd.DataFrame({"sparrow": range(10), "ostrich": 1})
    tensorflow_df = df.TensorFlowDataFrame.from_pandas(
        pandas_df, batch_size=10, shuffle=False)

    batch = tensorflow_df.run_one_batch()

    np.testing.assert_array_equal(pandas_df.index.values, batch["index"],
                                  "Expected index {}; got {}".format(
                                      pandas_df.index.values, batch["index"]))
    _assert_df_equals_dict(pandas_df, batch)

  def testBatch(self):
    """Tests `batch` method.

    `DataFrame.batch()` should iterate through the rows of the
    `pandas.DataFrame`, and should "wrap around" when it reaches the last row.
    """
    if not HAS_PANDAS:
      return
    pandas_df = pd.DataFrame({
        "albatross": range(10),
        "bluejay": 1,
        "cockatoo": range(0, 20, 2),
        "penguin": list("abcdefghij")
    })