"""Tests for run_report_accessor.py.

Copyright 2018 The AdaNet 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

    https://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.
"""

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

from absl.testing import parameterized
from adanet import subnetwork
from adanet.core.report_accessor import _ReportAccessor
import tensorflow.compat.v2 as tf


class ReportAccessorTest(parameterized.TestCase, tf.test.TestCase):

  def test_read_from_empty_file(self):
    report_accessor = _ReportAccessor(self.get_temp_dir())
    self.assertEqual([], list(report_accessor.read_iteration_reports()))

  def test_add_to_empty_file(self):
    report_accessor = _ReportAccessor(self.get_temp_dir())
    materialized_reports = [
        subnetwork.MaterializedReport(
            iteration_number=0,
            name="foo",
            hparams={
                "p1": 1,
                "p2": "default_hparam",
                "p3": b"binary_hparam",
                "p4": u"unicode_hparam",
                "p5": True,
            },
            attributes={
                "a1": 1,
                "a2": "default_attribute",
                "a3": b"binary_attribute",
                "a4": u"unicode_attribute",
                "a5": True,
            },
            metrics={
                "m1": 1,
                "m2": "default_metric",
                "m3": b"binary_metric",
                "m4": u"unicode_metric",
                "m5": True,
            },
            included_in_final_ensemble=True,
        ),
    ]

    report_accessor.write_iteration_report(
        iteration_number=0,
        materialized_reports=materialized_reports,
    )
    actual_iteration_reports = list(report_accessor.read_iteration_reports())

    self.assertLen(actual_iteration_reports, 1)
    self.assertEqual(materialized_reports, actual_iteration_reports[0])

  def test_add_to_existing_file(self):
    materialized_reports = [
        [
            subnetwork.MaterializedReport(
                iteration_number=0,
                name="foo1",
                hparams={
                    "p1": 11,
                    "p2": "default_hparam",
                    "p3": b"binary_hparam",
                    "p4": u"unicode_hparam",
                    "p5": True,
                },
                attributes={
                    "a1": 11,
                    "a2": "default_attribute",
                    "a3": b"binary_attribute",
                    "a4": u"unicode_attribute",
                    "a5": True,
                },
                metrics={
                    "m1": 11,
                    "m2": "default_metric",
                    "m3": b"binary_metric",
                    "m4": u"unicode_metric",
                    "m5": True,
                },
                included_in_final_ensemble=False,
            ),
            subnetwork.MaterializedReport(
                iteration_number=0,
                name="foo2",
                hparams={
                    "p1": 12,
                    "p2": "default_hparam",
                    "p3": b"binary_hparam",
                    "p4": u"unicode_hparam",
                    "p5": True,
                },
                attributes={
                    "a1": 12,
                    "a2": "default_attribute",
                    "a3": b"binary_attribute",
                    "a4": u"unicode_attribute",
                    "a5": True,
                },
                metrics={
                    "m1": 12,
                    "m2": "default_metric",
                    "m3": b"binary_metric",
                    "m4": u"unicode_metric",
                    "m5": True,
                },
                included_in_final_ensemble=True,
            ),
        ],
        [
            subnetwork.MaterializedReport(
                iteration_number=1,
                name="foo1",
                hparams={
                    "p1": 21,
                    "p2": "default_hparam",
                    "p3": b"binary_hparam",
                    "p4": u"unicode_hparam",
                    "p5": True,
                },
                attributes={
                    "a1": 21,
                    "a2": "default_attribute",
                    "a3": b"binary_attribute",
                    "a4": u"unicode_attribute",
                    "a5": True,
                },
                metrics={
                    "m1": 21,
                    "m2": "default_metric",
                    "m3": b"binary_metric",
                    "m4": u"unicode_metric",
                    "m5": True,
                },
                included_in_final_ensemble=True,
            ),
            subnetwork.MaterializedReport(
                iteration_number=1,
                name="foo2",
                hparams={
                    "p1": 22,
                    "p2": "default_hparam",
                    "p3": b"binary_hparam",
                    "p4": u"unicode_hparam",
                    "p5": True,
                },
                attributes={
                    "a1": 22,
                    "a2": "default_attribute",
                    "a3": b"binary_attribute",
                    "a4": u"unicode_attribute",
                    "a5": True,
                },
                metrics={
                    "m1": 22,
                    "m2": "default_metric",
                    "m3": b"binary_metric",
                    "m4": u"unicode_metric",
                    "m5": True,
                },
                included_in_final_ensemble=False,
            ),
        ],
        [
            subnetwork.MaterializedReport(
                iteration_number=2,
                name="foo1",
                hparams={
                    "p1": 31,
                    "p2": "default_hparam",
                    "p3": b"binary_hparam",
                    "p4": u"unicode_hparam",
                    "p5": True,
                },
                attributes={
                    "a1": 31,
                    "a2": "default_attribute",
                    "a3": b"binary_attribute",
                    "a4": u"unicode_attribute",
                    "a5": True,
                },
                metrics={
                    "m1": 31,
                    "m2": "default_metric",
                    "m3": b"binary_metric",
                    "m4": u"unicode_metric",
                    "m5": True,
                },
                included_in_final_ensemble=False,
            ),
            subnetwork.MaterializedReport(
                iteration_number=2,
                name="foo2",
                hparams={
                    "p1": 32,
                    "p2": "default_hparam",
                    "p3": b"binary_hparam",
                    "p4": u"unicode_hparam",
                    "p5": True,
                },
                attributes={
                    "a1": 32,
                    "a2": "default_attribute",
                    "a3": b"binary_attribute",
                    "a4": u"unicode_attribute",
                    "a5": True,
                },
                metrics={
                    "m1": 32,
                    "m2": "default_metric",
                    "m3": b"binary_metric",
                    "m4": u"unicode_metric",
                    "m5": True,
                },
                included_in_final_ensemble=True,
            ),
        ],
    ]

    report_accessor = _ReportAccessor(self.get_temp_dir())

    report_accessor.write_iteration_report(0, materialized_reports[0])
    report_accessor.write_iteration_report(1, materialized_reports[1])
    report_accessor.write_iteration_report(2, materialized_reports[2])
    actual_reports = list(report_accessor.read_iteration_reports())
    self.assertEqual(materialized_reports, actual_reports)

  def test_write_iteration_report_encoding(self):
    """Tests GitHub issue #4."""

    report_accessor = _ReportAccessor(self.get_temp_dir())
    binary_type_value = b"\n\x83\x01\n;adanet/iteration_2/ensemble_2_layer_dnn/"
    text_type_value = u"\U0001f937"
    materialized_reports = [
        subnetwork.MaterializedReport(
            iteration_number=0,
            name="foo",
            hparams={
                "p1": binary_type_value,
                "p2": text_type_value,
            },
            attributes={
                "a1": binary_type_value,
                "a2": text_type_value,
            },
            metrics={
                "m1": binary_type_value,
                "m2": text_type_value,
            },
            included_in_final_ensemble=True,
        ),
    ]

    report_accessor.write_iteration_report(
        iteration_number=0,
        materialized_reports=materialized_reports,
    )
    actual_iteration_reports = list(report_accessor.read_iteration_reports())
    self.assertLen(actual_iteration_reports, 1)

  @parameterized.named_parameters({
      "testcase_name": "hparams_invalid_type",
      "hparams": {
          "h1": None
      },
  }, {
      "testcase_name": "attributes_invalid_type",
      "attributes": {
          "a1": None
      },
  }, {
      "testcase_name": "metrics_invalid_type",
      "metrics": {
          "m1": None
      },
  })
  def test_value_error(self, hparams=None, attributes=None, metrics=None):
    if hparams is None:
      hparams = {}
    if attributes is None:
      attributes = {}
    if metrics is None:
      metrics = {}
    report_accessor = _ReportAccessor(self.get_temp_dir())
    materialized_reports = [
        subnetwork.MaterializedReport(
            iteration_number=0,
            name="foo",
            hparams=hparams,
            attributes=attributes,
            metrics=metrics,
            included_in_final_ensemble=True,
        ),
    ]
    with self.assertRaises(ValueError):
      report_accessor.write_iteration_report(
          iteration_number=0,
          materialized_reports=materialized_reports,
      )


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