# 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 backend_context."""


import operator
from unittest import mock

import tensorflow as tf

from google.protobuf import text_format
from tensorboard import context
from tensorboard.data import provider
from tensorboard.plugins import base_plugin
from tensorboard.plugins.hparams import api_pb2
from tensorboard.plugins.hparams import backend_context
from tensorboard.plugins.hparams import metadata
from tensorboard.plugins.hparams import plugin_data_pb2

DATA_TYPE_EXPERIMENT = "experiment"
DATA_TYPE_SESSION_START_INFO = "session_start_info"
DATA_TYPE_SESSION_END_INFO = "session_end_info"


class BackendContextTest(tf.test.TestCase):
    # Make assertProtoEquals print all the diff.
    maxDiff = None  # pylint: disable=invalid-name

    def setUp(self):
        self._mock_tb_context = base_plugin.TBContext()

        self._mock_tb_context.data_provider = mock.create_autospec(
            provider.DataProvider
        )
        self._mock_tb_context.data_provider.list_tensors.side_effect = (
            self._mock_list_tensors
        )
        self._mock_tb_context.data_provider.list_scalars.side_effect = (
            self._mock_list_scalars
        )
        self._mock_tb_context.data_provider.list_hyperparameters.side_effect = (
            self._mock_list_hyperparameters
        )

        self.session_1_start_info_ = ""
        self.session_2_start_info_ = ""
        self.session_3_start_info_ = ""
        self._hyperparameters = []

    def _mock_list_tensors(
        self, ctx, *, experiment_id, plugin_name, run_tag_filter
    ):
        hparams_content = {
            "exp/session_1": {
                metadata.SESSION_START_INFO_TAG: self._serialized_plugin_data(
                    DATA_TYPE_SESSION_START_INFO, self.session_1_start_info_
                ),
            },
            "exp/session_2": {
                metadata.SESSION_START_INFO_TAG: self._serialized_plugin_data(
                    DATA_TYPE_SESSION_START_INFO, self.session_2_start_info_
                ),
            },
            "exp/session_3": {
                metadata.SESSION_START_INFO_TAG: self._serialized_plugin_data(
                    DATA_TYPE_SESSION_START_INFO, self.session_3_start_info_
                ),
            },
        }
        result = {}
        for run, tag_to_content in hparams_content.items():
            result.setdefault(run, {})
            for tag, content in tag_to_content.items():
                t = provider.TensorTimeSeries(
                    max_step=0,
                    max_wall_time=0,
                    plugin_content=content,
                    description="",
                    display_name="",
                )
                result[run][tag] = t
        return result

    def _mock_list_scalars(
        self,
        ctx,
        *,
        experiment_id,
        plugin_name,
        run_tag_filter=provider.RunTagFilter(),
    ):
        scalars_content = {
            "exp/session_1": {"loss": b"", "accuracy": b""},
            "exp/session_1/eval": {
                "loss": b"",
            },
            "exp/session_1/train": {
                "loss": b"",
            },
            "exp/session_2": {
                "loss": b"",
                "accuracy": b"",
            },
            "exp/session_2/eval": {
                "loss": b"",
            },
            "exp/session_2/train": {
                "loss": b"",
            },
            "exp/session_3": {
                "loss": b"",
                "accuracy": b"",
            },
            "exp/session_3/eval": {
                "loss": b"",
            },
            "exp/session_3xyz/": {
                "loss2": b"",
            },
        }
        result = {}
        for run, tag_to_content in scalars_content.items():
            result.setdefault(run, {})
            for tag, content in tag_to_content.items():
                t = provider.ScalarTimeSeries(
                    max_step=0,
                    max_wall_time=0,
                    plugin_content=content,
                    description="",
                    display_name="",
                )
                result[run][tag] = t
        return result

    def _mock_list_hyperparameters(
        self,
        ctx,
        *,
        experiment_ids,
    ):
        return self._hyperparameters

    def _experiment_from_metadata(self):
        """Calls the expected operations for generating an Experiment proto."""
        ctxt = backend_context.Context(self._mock_tb_context)
        request_ctx = context.RequestContext()
        return ctxt.experiment_from_metadata(
            request_ctx,
            "123",
            ctxt.hparams_metadata(request_ctx, "123"),
            ctxt.hparams_from_data_provider(request_ctx, "123"),
        )

    def test_experiment_with_experiment_tag(self):
        experiment = """
            description: 'Test experiment'
            metric_infos: [
              { name: { tag: 'current_temp' } }
            ]
        """
        run = "exp"
        tag = metadata.EXPERIMENT_TAG
        t = provider.TensorTimeSeries(
            max_step=0,
            max_wall_time=0,
            plugin_content=self._serialized_plugin_data(
                DATA_TYPE_EXPERIMENT, experiment
            ),
            description="",
            display_name="",
        )
        self._mock_tb_context.data_provider.list_tensors.side_effect = None
        self._mock_tb_context.data_provider.list_tensors.return_value = {
            run: {tag: t}
        }
        self.assertProtoEquals(experiment, self._experiment_from_metadata())

    def test_experiment_without_experiment_tag(self):
        self.session_1_start_info_ = """
            hparams: [
              {key: 'batch_size' value: {number_value: 100}},
              {key: 'lr' value: {number_value: 0.01}},
              {key: 'model_type' value: {string_value: 'CNN'}}
            ]
        """
        self.session_2_start_info_ = """
            hparams:[
              {key: 'batch_size' value: {number_value: 200}},
              {key: 'lr' value: {number_value: 0.02}},
              {key: 'model_type' value: {string_value: 'LATTICE'}}
            ]
        """
        self.session_3_start_info_ = """
            hparams:[
              {key: 'batch_size' value: {number_value: 300}},
              {key: 'lr' value: {number_value: 0.05}},
              {key: 'model_type' value: {string_value: 'CNN'}}
            ]
        """
        expected_exp = """
            hparam_infos: {
              name: 'batch_size'
              type: DATA_TYPE_FLOAT64
            },
            hparam_infos: {
              name: 'lr'
              type: DATA_TYPE_FLOAT64
            },
            hparam_infos: {
              name: 'model_type'
              type: DATA_TYPE_STRING
              domain_discrete: {
                values: [{string_value: 'CNN'},
                         {string_value: 'LATTICE'}]
              }
            }
            metric_infos: {
              name: {group: '', tag: 'accuracy'}
            }
            metric_infos: {
              name: {group: '', tag: 'loss'}
            }
            metric_infos: {
              name: {group: 'eval', tag: 'loss'}
            }
            metric_infos: {
              name: {group: 'train', tag: 'loss'}
            }
        """
        actual_exp = self._experiment_from_metadata()
        _canonicalize_experiment(actual_exp)
        self.assertProtoEquals(expected_exp, actual_exp)

    def test_experiment_without_experiment_tag_different_hparam_types(self):
        self.session_1_start_info_ = """
            hparams:[
              {key: 'batch_size' value: {number_value: 100}},
              {key: 'lr' value: {string_value: '0.01'}}
            ]
        """
        self.session_2_start_info_ = """
            hparams:[
              {key: 'lr' value: {number_value: 0.02}},
              {key: 'model_type' value: {string_value: 'LATTICE'}}
            ]
        """
        self.session_3_start_info_ = """
            hparams:[
              {key: 'batch_size' value: {bool_value: true}},
              {key: 'model_type' value: {string_value: 'CNN'}}
            ]
        """
        expected_exp = """
            hparam_infos: {
              name: 'batch_size'
              type: DATA_TYPE_STRING
              domain_discrete: {
                values: [{string_value: '100.0'},
                         {string_value: 'true'}]
              }
            }
            hparam_infos: {
              name: 'lr'
              type: DATA_TYPE_STRING
              domain_discrete: {
                values: [{string_value: '0.01'},
                         {string_value: '0.02'}]
              }
            }
            hparam_infos: {
              name: 'model_type'
              type: DATA_TYPE_STRING
              domain_discrete: {
                values: [{string_value: 'CNN'},
                         {string_value: 'LATTICE'}]
              }
            }
            metric_infos: {
              name: {group: '', tag: 'accuracy'}
            }
            metric_infos: {
              name: {group: '', tag: 'loss'}
            }
            metric_infos: {
              name: {group: 'eval', tag: 'loss'}
            }
            metric_infos: {
              name: {group: 'train', tag: 'loss'}
            }
        """
        actual_exp = self._experiment_from_metadata()
        _canonicalize_experiment(actual_exp)
        self.assertProtoEquals(expected_exp, actual_exp)

    def test_experiment_with_bool_types(self):
        self.session_1_start_info_ = """
            hparams:[
              {key: 'batch_size' value: {bool_value: true}}
            ]
        """
        self.session_2_start_info_ = """
            hparams:[
              {key: 'batch_size' value: {bool_value: true}}
            ]
        """
        self.session_3_start_info_ = """
            hparams:[
            ]
        """
        expected_exp = """
            hparam_infos: {
              name: 'batch_size'
              type: DATA_TYPE_BOOL
              domain_discrete: {
                values: [{bool_value: true}, {bool_value: false}]
              }
            }
            metric_infos: {
              name: {group: '', tag: 'accuracy'}
            }
            metric_infos: {
              name: {group: '', tag: 'loss'}
            }
            metric_infos: {
              name: {group: 'eval', tag: 'loss'}
            }
            metric_infos: {
              name: {group: 'train', tag: 'loss'}
            }
        """
        actual_exp = self._experiment_from_metadata()
        _canonicalize_experiment(actual_exp)
        self.assertProtoEquals(expected_exp, actual_exp)

    def test_experiment_without_any_hparams(self):
        request_ctx = context.RequestContext()
        actual_exp = self._experiment_from_metadata()
        self.assertIsInstance(actual_exp, api_pb2.Experiment)
        self.assertProtoEquals("", actual_exp)

    def test_experiment_from_data_provider_interval_hparam(self):
        self._hyperparameters = [
            provider.Hyperparameter(
                hyperparameter_name="hparam1_name",
                hyperparameter_display_name="hparam1_display_name",
                domain_type=provider.HyperparameterDomainType.INTERVAL,
                domain=(-10.0, 15),
            )
        ]
        self._mock_tb_context.data_provider.list_tensors.side_effect = None
        actual_exp = self._experiment_from_metadata()
        expected_exp = """
            hparam_infos: {
              name: 'hparam1_name'
              display_name: 'hparam1_display_name'
              type: DATA_TYPE_FLOAT64
              domain_interval: {
                min_value: -10.0
                max_value: 15
              }
            }
        """
        self.assertProtoEquals(expected_exp, actual_exp)

    def test_experiment_from_data_provider_discrete_bool_hparam(self):
        self._hyperparameters = [
            provider.Hyperparameter(
                hyperparameter_name="hparam1_name",
                hyperparameter_display_name="hparam1_display_name",
                domain_type=provider.HyperparameterDomainType.DISCRETE_BOOL,
                domain=[True],
            ),
            provider.Hyperparameter(
                hyperparameter_name="hparam2_name",
                hyperparameter_display_name="hparam2_display_name",
                domain_type=provider.HyperparameterDomainType.DISCRETE_BOOL,
                domain=[True, False],
            ),
            provider.Hyperparameter(
                hyperparameter_name="hparam3_name",
                hyperparameter_display_name="hparam3_display_name",
                domain_type=provider.HyperparameterDomainType.DISCRETE_BOOL,
                domain=[False],
            ),
            provider.Hyperparameter(
                hyperparameter_name="hparam4_name",
                hyperparameter_display_name="hparam4_display_name",
                domain_type=provider.HyperparameterDomainType.DISCRETE_BOOL,
                domain=[],
            ),
        ]
        self._mock_tb_context.data_provider.list_tensors.side_effect = None
        actual_exp = self._experiment_from_metadata()
        expected_exp = """
            hparam_infos: {
              name: 'hparam1_name'
              display_name: 'hparam1_display_name'
              type: DATA_TYPE_BOOL
              domain_discrete: {
                values: [{bool_value: true}]
              }
            }
            hparam_infos: {
              name: 'hparam2_name'
              display_name: 'hparam2_display_name'
              type: DATA_TYPE_BOOL
              domain_discrete: {
                values: [{bool_value: true}, {bool_value: false}]
              }
            }
            hparam_infos: {
              name: 'hparam3_name'
              display_name: 'hparam3_display_name'
              type: DATA_TYPE_BOOL
              domain_discrete: {
                values: [{bool_value: false}]
              }
            }
            hparam_infos: {
              name: 'hparam4_name'
              display_name: 'hparam4_display_name'
              type: DATA_TYPE_BOOL
            }
        """
        self.assertProtoEquals(expected_exp, actual_exp)

    def test_experiment_from_data_provider_discrete_float_hparam(self):
        self._hyperparameters = [
            provider.Hyperparameter(
                hyperparameter_name="hparam1_name",
                hyperparameter_display_name="hparam1_display_name",
                domain_type=provider.HyperparameterDomainType.DISCRETE_FLOAT,
                domain=[-1.0, 1.5, 0.0],
            ),
        ]
        self._mock_tb_context.data_provider.list_tensors.side_effect = None
        actual_exp = self._experiment_from_metadata()
        expected_exp = """
            hparam_infos: {
              name: 'hparam1_name'
              display_name: 'hparam1_display_name'
              type: DATA_TYPE_FLOAT64
              domain_discrete: {
                values: [
                  {number_value: -1.0},
                  {number_value: 1.5},
                  {number_value: 0.0}
                ]
              }
            }
        """
        self.assertProtoEquals(expected_exp, actual_exp)

    def test_experiment_from_data_provider_discrete_string_hparam(self):
        self._hyperparameters = [
            provider.Hyperparameter(
                hyperparameter_name="hparam1_name",
                hyperparameter_display_name="hparam1_display_name",
                domain_type=provider.HyperparameterDomainType.DISCRETE_STRING,
                domain=["one", "two", "aaaa"],
            ),
        ]
        self._mock_tb_context.data_provider.list_tensors.side_effect = None
        actual_exp = self._experiment_from_metadata()
        expected_exp = """
            hparam_infos: {
              name: 'hparam1_name'
              display_name: 'hparam1_display_name'
              type: DATA_TYPE_STRING
              domain_discrete: {
                values: [
                  {string_value: 'one'},
                  {string_value: 'two'},
                  {string_value: 'aaaa'}
                ]
              }
            }
        """
        self.assertProtoEquals(expected_exp, actual_exp)

    def _serialized_plugin_data(self, data_oneof_field, text_protobuffer):
        oneof_type_dict = {
            DATA_TYPE_EXPERIMENT: api_pb2.Experiment,
            DATA_TYPE_SESSION_START_INFO: plugin_data_pb2.SessionStartInfo,
            DATA_TYPE_SESSION_END_INFO: plugin_data_pb2.SessionEndInfo,
        }
        protobuffer = text_format.Merge(
            text_protobuffer, oneof_type_dict[data_oneof_field]()
        )
        plugin_data = plugin_data_pb2.HParamsPluginData()
        getattr(plugin_data, data_oneof_field).CopyFrom(protobuffer)
        return metadata.create_summary_metadata(plugin_data).plugin_data.content


def _canonicalize_experiment(exp):
    """Sorts the repeated fields of an Experiment message."""
    exp.hparam_infos.sort(key=operator.attrgetter("name"))
    exp.metric_infos.sort(key=operator.attrgetter("name.group", "name.tag"))
    for hparam_info in exp.hparam_infos:
        if hparam_info.HasField("domain_discrete"):
            hparam_info.domain_discrete.values.sort(
                key=operator.attrgetter("string_value")
            )


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