# Copyright 2020 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 input_pipeline_proto_to_gviz."""

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

import csv
import io
import enum

import gviz_api
import tensorflow as tf
# pylint: disable=g-importing-member
from google3.google.protobuf.any_pb2 import Any
# pylint: enable=g-importing-member

from tensorboard_plugin_profile.convert import input_pipeline_proto_to_gviz
from tensorboard_plugin_profile.protobuf import input_pipeline_pb2


class StrEnum(str, enum.Enum):
  pass


class MockValues(StrEnum):
  STEP_NUMBER = 1
  STEP_TIME_MS = 2
  UNKNOWN_TIME_MS = 3
  HOST_WAIT_INPUT_MS = 11
  HOST_TO_DEVICE_MS = 12
  OUTPUT_MS = 5
  DEVICE_COMPUTE_MS = 6
  DEVICE_TO_DEVICE_MS = 7
  DEVICE_COLLECTIVES_MS = 13
  HOST_COMPUTE_MS = 8
  HOST_PREPARE_MS = 9
  HOST_COMPILE_MS = 10
  TOOLTIP = """Step 1, duration: 2.00 ms
-All others: 3.00 ms
-Compilation: 10.00 ms
-Output: 5.00 ms
-Input: 23.00 ms
-Kernel launch: 9.00 ms
-Host compute: 8.00 ms
-Device collectives: 13.00 ms
-Device to device: 7.00 ms
-Device compute: 6.00 ms"""


class ProtoToGvizTest(tf.test.TestCase):

  def create_empty_input_pipeline(self):
    return input_pipeline_pb2.InputPipelineAnalysisResult()

  def create_mock_step_summary(self, base):
    step_summary = input_pipeline_pb2.StepSummary()
    step_summary.average = 1 + base
    step_summary.standard_deviation = 2 + base
    step_summary.minimum = 3 + base
    step_summary.maximum = 4 + base
    return step_summary

  def create_mock_input_pipeline(self):
    ipa = input_pipeline_pb2.InputPipelineAnalysisResult()
    ipa.hardware_type = "CPU_ONLY"
    ipa.step_time_summary.CopyFrom(self.create_mock_step_summary(10))
    ipa.input_percent_summary.CopyFrom(self.create_mock_step_summary(20))

    # Add 3 rows
    for _ in range(0, 3):
      step_details = input_pipeline_pb2.PerGenericStepDetails()
      step_details.step_number = int(MockValues.STEP_NUMBER)
      step_details.step_name = MockValues.STEP_NUMBER
      step_details.step_time_ms = int(MockValues.STEP_TIME_MS)
      step_details.unknown_time_ms = int(MockValues.UNKNOWN_TIME_MS)
      step_details.host_wait_input_ms = int(MockValues.HOST_WAIT_INPUT_MS)
      step_details.host_to_device_ms = int(MockValues.HOST_TO_DEVICE_MS)
      step_details.output_ms = int(MockValues.OUTPUT_MS)
      step_details.device_compute_ms = int(MockValues.DEVICE_COMPUTE_MS)
      step_details.device_to_device_ms = int(MockValues.DEVICE_TO_DEVICE_MS)
      step_details.device_collectives_ms = int(MockValues.DEVICE_COLLECTIVES_MS)
      step_details.host_compute_ms = int(MockValues.HOST_COMPUTE_MS)
      step_details.host_prepare_ms = int(MockValues.HOST_PREPARE_MS)
      step_details.host_compile_ms = int(MockValues.HOST_COMPILE_MS)

      step_details_any = Any()
      step_details_any.Pack(step_details)
      ipa.step_details.append(step_details_any)

    input_time_breakdown = input_pipeline_pb2.InputTimeBreakdown()
    input_time_breakdown.demanded_file_read_us = 1
    input_time_breakdown.advanced_file_read_us = 2
    input_time_breakdown.preprocessing_us = 3
    input_time_breakdown.enqueue_us = 4
    input_time_breakdown.unclassified_non_enqueue_us = 5
    ipa.input_time_breakdown.CopyFrom(input_time_breakdown)

    for _ in range(0, 3):
      input_op_details = input_pipeline_pb2.InputOpDetails()
      input_op_details.op_name = str(1)
      input_op_details.count = 2
      input_op_details.time_in_ms = 3
      input_op_details.time_in_percent = 4
      input_op_details.self_time_in_ms = 5
      input_op_details.self_time_in_percent = 6
      input_op_details.category = str(7)
      ipa.input_op_details.append(input_op_details)

    recommendation = input_pipeline_pb2.InputPipelineAnalysisRecommendation()
    for ss in ["a", "b", "c", "d", "e"]:
      recommendation.details.append(ss)
    ipa.recommendation.CopyFrom(recommendation)

    step_time_breakdown = input_pipeline_pb2.GenericStepTimeBreakdown()
    step_time_breakdown.unknown_time_ms_summary.CopyFrom(
        self.create_mock_step_summary(1))
    step_time_breakdown.host_wait_input_ms_summary.CopyFrom(
        self.create_mock_step_summary(9))
    step_time_breakdown.host_to_device_ms_summary.CopyFrom(
        self.create_mock_step_summary(10))
    step_time_breakdown.input_ms_summary.CopyFrom(
        self.create_mock_step_summary(11))
    step_time_breakdown.output_ms_summary.CopyFrom(
        self.create_mock_step_summary(3))
    step_time_breakdown.device_compute_ms_summary.CopyFrom(
        self.create_mock_step_summary(4))
    step_time_breakdown.device_to_device_ms_summary.CopyFrom(
        self.create_mock_step_summary(5))
    step_time_breakdown.device_collectives_ms_summary.CopyFrom(
        self.create_mock_step_summary(12))
    step_time_breakdown.host_compute_ms_summary.CopyFrom(
        self.create_mock_step_summary(6))
    step_time_breakdown.host_prepare_ms_summary.CopyFrom(
        self.create_mock_step_summary(7))
    step_time_breakdown.host_compile_ms_summary.CopyFrom(
        self.create_mock_step_summary(8))

    step_time_breakdown_any = Any()
    step_time_breakdown_any.Pack(step_time_breakdown)
    ipa.step_time_breakdown.CopyFrom(step_time_breakdown_any)

    return ipa

  def test_input_pipeline_empty(self):
    ipa = self.create_empty_input_pipeline()
    data_table = input_pipeline_proto_to_gviz.generate_step_breakdown_table(ipa)

    self.assertEqual(0, data_table.NumberOfRows(),
                     "Empty table should have 0 rows.")
    # Input pipeline chart data table has 11 columns.
    self.assertLen(data_table.columns, 11)

  def test_input_pipeline_step_breakdown(self):
    ipa = self.create_mock_input_pipeline()
    (table_description, data, custom_properties
    ) = input_pipeline_proto_to_gviz.get_step_breakdown_table_args(ipa)
    data_table = gviz_api.DataTable(table_description, data, custom_properties)

    # Data is a list of 3 rows.
    self.assertLen(data, 3)
    self.assertEqual(3, data_table.NumberOfRows(), "Simple table has 3 rows.")
    # Table descriptor is a list of 11 columns.
    self.assertLen(table_description, 11)
    # DataTable also has 10 columns.
    self.assertLen(data_table.columns, 11)

    csv_file = io.StringIO(data_table.ToCsv())
    reader = csv.reader(csv_file)

    expected = [
        str(int(MockValues.STEP_NUMBER)),
        int(MockValues.DEVICE_COMPUTE_MS),
        int(MockValues.DEVICE_TO_DEVICE_MS),
        int(MockValues.DEVICE_COLLECTIVES_MS),
        int(MockValues.HOST_COMPUTE_MS),
        int(MockValues.HOST_PREPARE_MS),
        int(MockValues.HOST_WAIT_INPUT_MS) + int(MockValues.HOST_TO_DEVICE_MS),
        int(MockValues.OUTPUT_MS),
        int(MockValues.HOST_COMPILE_MS),
        int(MockValues.UNKNOWN_TIME_MS),
        MockValues.TOOLTIP.value,
    ]

    for (rr, row_values) in enumerate(reader):
      if rr == 0:
        # DataTable columns match schema defined in table_description.
        for (cc, column_header) in enumerate(row_values):
          self.assertEqual(table_description[cc][2], column_header)
      else:
        for (cc, cell_str) in enumerate(row_values):
          raw_value = data[rr - 1][cc]
          value_type = table_description[cc][1]

          # Only number and strings are used in our DataTable schema.
          self.assertIn(value_type, ["number", "string"])

          # Encode in similar fashion as DataTable.ToCsv().
          expected_value = gviz_api.DataTable.CoerceValue(raw_value, value_type)
          self.assertNotIsInstance(expected_value, tuple)
          self.assertEqual(expected_value, raw_value)
          self.assertEqual(str(expected_value), cell_str)

          # Check against expected values we have set in our mock table.
          if isinstance(expected[cc], str):
            self.assertEqual(expected[cc], cell_str)
          else:
            self.assertEqual(str(float(expected[cc])), cell_str)


if __name__ == "__main__":
  tf.test.main()
