# coding=utf-8
# Copyright 2022 The Fiddle-Config Authors.
#
# 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 print."""

import dataclasses
import re
import textwrap
from typing import Union

from absl.testing import absltest
import fiddle as fdl
from fiddle import printing
from fiddle import tagging


class DummyTag(tagging.Tag):
  """Dummy tag for testing."""


class OtherTag(tagging.Tag):
  """Second tag, for fun & profit!"""


def fn_x_y(x, y):  # pylint: disable=unused-argument
  pass


def fn_with_kwargs(x, y, **kwargs):  # pylint: disable=unused-argument
  pass


class TestHelper:

  def __init__(self, a, b):  # pylint: disable=unused-argument
    pass


def fn_with_type_annotations(x: int, y: str, z: float):  # pylint: disable=unused-argument
  pass


@dataclasses.dataclass
class DataclassHelper:
  x: int
  y: str


def advanced_annotations_helper(x: TestHelper, **kwargs: DataclassHelper):  # pylint: disable=unused-argument
  pass


def annotated_kwargs_helper(**kwargs: int):  # pylint: disable=unused-argument
  pass


# Sometimes the module of local symbols shows up differently, depending on how
# this test is imported/run.
_local_module_regex = r'(__main__|fiddle\.printing_test)'


class AsStrFlattenedTests(absltest.TestCase):

  def check_result(self, actual, expected):
    expected = re.escape(textwrap.dedent(expected)).replace(
        '__main__', _local_module_regex)
    self.assertRegex(actual, expected)

  def test_path_to_str(self):
    self.assertEqual(printing._path_to_str(()), '')
    self.assertEqual(
        printing._path_to_str((printing._ParamName('x'), 0)), 'x[0]')
    self.assertEqual(
        printing._path_to_str((printing._ParamName('a'), 2, 'x')), "a[2]['x']")

  def test_simple_printing(self):
    cfg = fdl.Config(fn_x_y, 1, 'abc')
    output = printing.as_str_flattened(cfg)

    self.check_result(output, """\
        x = 1
        y = 'abc'""")

  def test_unset_argument(self):
    cfg = fdl.Config(fn_x_y, 3.14)
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x = 3.14
        y = <[unset]>""")
    self.assertEqual(output, expected)

  def test_nested(self):
    cfg = fdl.Config(fn_x_y, 'x', fdl.Config(fn_x_y, 'nest_x', 123))
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x = 'x'
        y.x = 'nest_x'
        y.y = 123""")
    self.assertEqual(output, expected)

  def test_class(self):
    cfg = fdl.Config(TestHelper, 'a_param', b=123)
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        a = 'a_param'
        b = 123""")
    self.assertEqual(output, expected)

  def test_kwargs(self):
    cfg = fdl.Config(fn_with_kwargs, 1, abc='extra kwarg value')
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x = 1
        y = <[unset]>
        abc = 'extra kwarg value'""")
    self.assertEqual(output, expected)

  def test_nested_kwargs(self):
    cfg = fdl.Config(
        fn_with_kwargs,
        extra=fdl.Config(fn_with_kwargs, 1, nested_extra='whee'))
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x = <[unset]>
        y = <[unset]>
        extra.x = 1
        extra.y = <[unset]>
        extra.nested_extra = 'whee'""")
    self.assertEqual(output, expected)

  def test_nested_collections(self):
    cfg = fdl.Config(fn_x_y,
                     [fdl.Config(fn_x_y, 1, '1'),
                      fdl.Config(TestHelper, 2)])
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x[0].x = 1
        x[0].y = '1'
        x[1].a = 2
        x[1].b = <[unset]>
        y = <[unset]>""")
    self.assertEqual(output, expected)

  def test_multiple_nested_collections(self):
    cfg = fdl.Config(fn_x_y, {
        'a': fdl.Config(fn_with_kwargs, abc=[1, 2, 3]),
        'b': [3, 2, 1]
    }, [fdl.Config(fn_x_y, [fdl.Config(fn_x_y, 1, 2)])])
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x['a'].x = <[unset]>
        x['a'].y = <[unset]>
        x['a'].abc = [1, 2, 3]
        x['b'][0] = 3
        x['b'][1] = 2
        x['b'][2] = 1
        y[0].x[0].x = 1
        y[0].x[0].y = 2
        y[0].y = <[unset]>""")
    self.assertEqual(output, expected)

  def test_default_values(self):

    def test_fn(w, x, y=3, z='abc'):  # pylint: disable=unused-argument
      pass

    cfg = fdl.Config(test_fn, 1)
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        w = 1
        x = <[unset]>
        y = <[unset; default: 3]>
        z = <[unset; default: 'abc']>""")
    self.assertEqual(output, expected)

  def test_tagged_values(self):
    cfg = fdl.Config(fn_x_y, x=DummyTag.new(), y=DummyTag.new(default='abc'))
    output = printing.as_str_flattened(cfg)

    self.check_result(
        output, """\
        x = <[unset]> #__main__.DummyTag
        y = 'abc' #__main__.DummyTag""")

    tagging.set_tagged(cfg, tag=DummyTag, value='cba')
    output = printing.as_str_flattened(cfg)

    self.check_result(
        output, """\
        x = 'cba' #__main__.DummyTag
        y = 'cba' #__main__.DummyTag""")

  def test_tagged_values_multiple_tags(self):
    cfg = fdl.Config(
        fn_x_y,
        x=tagging.TaggedValue(tags=(DummyTag, OtherTag)),
        y=tagging.TaggedValue(tags=(DummyTag, OtherTag), default='abc'))
    output = printing.as_str_flattened(cfg)

    self.check_result(
        output, """\
        x = <[unset]> #__main__.DummyTag #__main__.OtherTag
        y = 'abc' #__main__.DummyTag #__main__.OtherTag""")

    tagging.set_tagged(cfg, tag=DummyTag, value='cba')
    output = printing.as_str_flattened(cfg)

    self.check_result(
        output, """\
        x = 'cba' #__main__.DummyTag #__main__.OtherTag
        y = 'cba' #__main__.DummyTag #__main__.OtherTag""")

  def test_partial(self):
    partial = fdl.Partial(fn_x_y)
    partial.x = 'abc'
    output = printing.as_str_flattened(partial)

    expected = textwrap.dedent("""\
        x = 'abc'
        y = <[unset]>""")
    self.assertEqual(output, expected)

  def test_builtin_types_annotations(self):
    cfg = fdl.Config(fn_with_type_annotations, 1)
    cfg.y = 'abc'
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x: int = 1
        y: str = 'abc'
        z: float = <[unset]>""")
    self.assertEqual(output, expected)

  def test_advanced_type_annotations(self):
    cfg = fdl.Config(advanced_annotations_helper)
    cfg.abc = fdl.Config(DataclassHelper)
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x: TestHelper = <[unset]>
        abc.x: int = <[unset]>
        abc.y: str = <[unset]>""")
    self.assertEqual(output, expected)

  def test_annotated_kwargs(self):
    cfg = fdl.Config(annotated_kwargs_helper, x=1, y='oops')
    output = printing.as_str_flattened(cfg)

    expected = textwrap.dedent("""\
        x: int = 1
        y: int = 'oops'""")
    self.assertEqual(output, expected)

  def test_disabling_type_annotations(self):
    cfg = fdl.Config(fn_with_type_annotations, 1)
    cfg.y = 'abc'
    output = printing.as_str_flattened(cfg, include_types=False)

    expected = textwrap.dedent("""\
        x = 1
        y = 'abc'
        z = <[unset]>""")
    self.assertEqual(output, expected)

  def test_can_print_union_type(self):

    def to_integer(x: Union[int, str]):
      return int(x)

    cfg = fdl.Config(to_integer, 1)
    output = printing.as_str_flattened(cfg, include_types=True)

    expected = textwrap.dedent("""\
        x: typing.Union[int, str] = 1""")
    self.assertEqual(output, expected)

  def test_materialized_default_values(self):

    def test_fn(w, x, y=3, z='abc'):
      del w, x, y, z  # Unused.

    cfg = fdl.Config(test_fn, 1)
    fdl.materialize_defaults(cfg)
    output = printing.as_str_flattened(cfg)
    expected = textwrap.dedent("""\
        w = 1
        x = <[unset]>
        y = 3
        z = 'abc'""")
    self.assertEqual(output, expected)


class HistoryPerLeafParamTests(absltest.TestCase):

  def test_simple_history(self):
    cfg = fdl.Config(fn_x_y, 1, 'abc')
    cfg.x = 2
    output = printing.history_per_leaf_parameter(cfg)
    expected = textwrap.dedent(r"""
        x = 2 @ .*/printing_test.py:\d+:test_simple_history
          - previously: 1 @ .*/printing_test.py:\d+:test_simple_history
        y = 'abc' @ .*/printing_test.py:\d+:test_simple_history""".strip('\n'))
    self.assertRegex(output, expected)

  def test_nested_in_collections(self):
    cfg = fdl.Config(fn_x_y,
                     [fdl.Config(fn_x_y, 1, '1'),
                      fdl.Config(TestHelper, 2)])
    cfg.x[0].x = 3
    cfg.x[1].a = 2  # Reset to same value.
    cfg.x[0].y = 'abc'
    cfg.x[0].x = 4
    output = printing.history_per_leaf_parameter(cfg)
    expected = textwrap.dedent(rf"""
        __fn_or_cls__ = .*fn_x_y .+/printing_test.py:\d+:test_nested_in_collections
        x\[0\].__fn_or_cls__ = .*fn_x_y .+/printing_test.py:\d+:test_nested_in_collections
        x\[0\].x = 4 @ .*/printing_test.py:\d+:test_nested_in_collections
          - previously: 3 @ .*/printing_test.py:\d+:test_nested_in_collections
          - previously: 1 @ .*/printing_test.py:\d+:test_nested_in_collections
        x\[0\].y = 'abc' @ .*/printing_test.py:\d+:test_nested_in_collections
          - previously: '1' @ .*/printing_test.py:\d+:test_nested_in_collections
        x\[1\].__fn_or_cls__ = .*{_local_module_regex}.TestHelper.*/printing_test.py:\d+:test_nested_in_collections
        x\[1\].a = 2 @ .*/printing_test.py:\d+:test_nested_in_collections
          - previously: 2 @ .*/printing_test.py:\d+:test_nested_in_collections
        x\[1\].b = <\[unset\]>
        y = <\[unset\]>""".strip('\n'))
    self.assertRegex(output, expected)

  def test_update_callable_history(self):
    cfg = fdl.Config(fn_x_y, x=1, y=2)
    fdl.update_callable(cfg, fn_with_kwargs)
    cfg.abc = '123'
    output = printing.history_per_leaf_parameter(cfg)
    expected = textwrap.dedent(r"""
        __fn_or_cls__ = .*fn_with_kwargs .+/printing_test.py:\d+:test_update_callable_history
          - previously: .*fn_x_y .+/printing_test.py:\d+:test_update_callable_history
        x = 1 @ .+/printing_test.py:\d+:test_update_callable_history
        y = 2 @ .+/printing_test.py:\d+:test_update_callable_history
        abc = '123' @ .+/printing_test.py:\d+:test_update_callable_history
        kwargs = <\[unset\]>""".strip('\n'))
    self.assertRegex(output, expected)

  def test_materialize_defaults_history(self):

    def test_defaulting_helper(w, x, y=1, z=2):
      del w, x, y, z  # Unused.

    cfg = fdl.Config(test_defaulting_helper, w=0)
    cfg.y = 5
    cfg.y = 6
    fdl.materialize_defaults(cfg)
    output = printing.history_per_leaf_parameter(cfg)
    expected = textwrap.dedent(r"""
        __fn_or_cls__ = .*test_defaulting_helper .+/printing_test.py:\d+:test_materialize_defaults_history
        w = 0 @ .*/printing_test.py:\d+:test_materialize_defaults_history
        y = 6 @ .*/printing_test.py:\d+:test_materialize_defaults_history
          - previously: 5 .*
        z = 2 @ .*/printing_test.py:\d+:test_materialize_defaults_history
        x = <\[unset\]>""".strip('\n'))
    self.assertRegex(output, expected)


if __name__ == '__main__':
  absltest.main()
