import sys
import _setup_test_env  # noqa
import unittest
from nose.tools import assert_equal, assert_is_instance, assert_in, assert_greater, assert_true, assert_false
from pprint import pprint
from returnn.config import Config
from returnn.util.basic import PY3
from returnn.util import better_exchook

better_exchook.replace_traceback_format_tb()
if PY3:
    from io import StringIO
else:
    # noinspection PyUnresolvedReferences,PyCompatibility
    from StringIO import StringIO


def test_old_format():
    config = Config()
    config.load_file(
        StringIO(
            """
  # comment
  num_inputs 3
  hidden_type forward,lstm
  """
        )
    )

    assert_true(config.has("num_inputs"))
    assert_true(config.has("hidden_type"))
    assert_equal(config.int("num_inputs", -1), 3)
    assert_equal(config.value("hidden_type", "x"), "forward,lstm")
    assert_equal(config.value("hidden_type", "x", index=0), "forward")
    assert_equal(config.value("hidden_type", "x", index=1), "lstm")
    assert_equal(config.list("hidden_type", ["x"]), ["forward", "lstm"])

    assert_false(config.is_typed("num_inputs"))
    assert_false(config.is_typed("hidden_type"))


def test_json_format():
    config = Config()
    config.load_file(
        StringIO(
            """
  {
  // comment
  "num_inputs": 3,
  "hidden_type": ["forward", "lstm"]
  }
  """
        )
    )

    assert_true(config.has("num_inputs"))
    assert_true(config.has("hidden_type"))
    assert_equal(config.int("num_inputs", -1), 3)
    assert_equal(config.value("hidden_type", "x"), "forward,lstm")
    assert_equal(config.value("hidden_type", "x", index=0), "forward")
    assert_equal(config.value("hidden_type", "x", index=1), "lstm")
    assert_equal(config.list("hidden_type", ["x"]), ["forward", "lstm"])

    assert_true(config.is_typed("num_inputs"))
    assert_true(config.is_typed("hidden_type"))
    assert_is_instance(config.typed_value("num_inputs"), int)
    assert_is_instance(config.typed_value("hidden_type"), list)
    assert_equal(config.typed_value("hidden_type"), ["forward", "lstm"])


def test_py_config():
    config = Config()
    config.load_file(
        StringIO(
            """#!rnn.py
# comment
num_inputs = 3
hidden_type = ["forward", "lstm"]
  """
        )
    )

    assert_true(config.has("num_inputs"))
    assert_true(config.has("hidden_type"))
    assert_equal(config.int("num_inputs", -1), 3)
    assert_equal(config.value("hidden_type", "x"), "forward,lstm")
    assert_equal(config.value("hidden_type", "x", index=0), "forward")
    assert_equal(config.value("hidden_type", "x", index=1), "lstm")
    assert_equal(config.list("hidden_type", ["x"]), ["forward", "lstm"])

    assert_true(config.is_typed("num_inputs"))
    assert_true(config.is_typed("hidden_type"))
    assert_is_instance(config.typed_value("num_inputs"), int)
    assert_is_instance(config.typed_value("hidden_type"), list)
    assert_equal(config.typed_value("hidden_type"), ["forward", "lstm"])


def test_rnn_init_config_py_global_var():
    import returnn.__main__ as rnn
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", suffix=".config", prefix="test_rnn_initConfig") as cfgfile:
        cfgfile.write(
            """#!rnn.py
task = config.value("task", "train")

test_value = 42

def test_func():
  return task

    """
        )
        cfgfile.flush()
        rnn.init_config(command_line_options=[cfgfile.name, "--task", "search"])

    assert isinstance(rnn.config, Config)
    rnn.config.typed_dict.pop("__builtins__", None)  # not needed, too verbose for pprint
    pprint(rnn.config.dict)
    pprint(rnn.config.typed_dict)
    assert rnn.config.has("task")
    assert rnn.config.has("test_value")
    assert rnn.config.has("test_func")
    assert_equal(rnn.config.value("task", None), "search")
    assert rnn.config.is_typed("test_value")
    assert_equal(rnn.config.typed_value("test_value"), 42)
    assert rnn.config.is_typed("test_func")
    # So far it's fine.
    # Now something a bit strange.
    # Earlier, this failed, because the command-line overwrote this.
    assert rnn.config.is_typed("task")
    test_func = rnn.config.typed_dict["test_func"]
    assert callable(test_func)
    assert_equal(test_func(), "search")


def test_rnn_init_config_py_cmd_type():
    import returnn.__main__ as rnn
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", suffix=".config", prefix="test_rnn_initConfig") as cfgfile:
        cfgfile.write(
            """#!rnn.py
max_seq_length = {'bpe': 75}

def test_func():
  return max_seq_length

    """
        )
        cfgfile.flush()
        rnn.init_config(command_line_options=[cfgfile.name, "++max_seq_length", "0"])

    assert isinstance(rnn.config, Config)
    assert rnn.config.has("max_seq_length")
    assert rnn.config.has("test_func")
    assert rnn.config.is_typed("max_seq_length")
    assert rnn.config.is_typed("test_func")
    test_func = rnn.config.typed_dict["test_func"]
    assert callable(test_func)
    assert_equal(test_func(), 0)


def test_config_py_ext():
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", prefix="test_rnn_initConfig") as cfgfile:
        cfgfile.write(
            """
def test_func():
  return config.value("task", "train")
    """
        )
        cfgfile.flush()
        config = Config()
        config.load_file(cfgfile.name)  # should determine format by suffix ".py"

    config.typed_dict.pop("__builtins__", None)  # not needed, too verbose for pprint
    pprint(config.dict)
    pprint(config.typed_dict)
    assert config.has("test_func")
    assert config.is_typed("test_func")
    test_func = config.typed_dict["test_func"]
    assert callable(test_func)
    assert_equal(test_func(), "train")
    config.set("task", "search")
    assert_equal(test_func(), "search")


def test_config_py_old_returnn_imports():
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", prefix="test_rnn_initConfig") as cfgfile:
        cfgfile.write(
            """
# These are some common imports found in older config files.
from Pretrain import WrapEpochValue
from TFUtil import where_bc
import TFUtil
import returnn.TFUtil

import TFHorovod
hvd = TFHorovod.get_ctx(config=config)  # should return None if no Horovod context

from returnn.Util import describe_returnn_version
returnn_version = describe_returnn_version()

tf_version_tuple = returnn.TFUtil.tf_version_tuple()

    """
        )
        cfgfile.flush()
        config = Config()
        config.load_file(cfgfile.name)  # should determine format by suffix ".py"

    config.typed_dict.pop("__builtins__", None)  # not needed, too verbose for pprint
    pprint(config.typed_dict)

    import returnn.util.basic as util
    import returnn.tf.util.basic as tf_util

    assert config.typed_dict["where_bc"] is tf_util.where_bc
    assert config.typed_dict["TFUtil"].where_bc is tf_util.where_bc
    assert config.typed_dict["hvd"] is None
    assert config.typed_dict["tf_version_tuple"] == tf_util.tf_version_tuple()
    assert config.typed_dict["returnn_version"] == util.describe_returnn_version()
    assert config.typed_dict["returnn"].TFUtil is tf_util


if __name__ == "__main__":
    better_exchook.install()
    if len(sys.argv) <= 1:
        for k, v in sorted(globals().items()):
            if k.startswith("test_"):
                print("-" * 40)
                print("Executing: %s" % k)
                try:
                    v()
                except unittest.SkipTest as exc:
                    print("SkipTest:", exc)
                print("-" * 40)
        print("Finished all tests.")
    else:
        assert len(sys.argv) >= 2
        for arg in sys.argv[1:]:
            print("Executing: %s" % arg)
            if arg in globals():
                globals()[arg]()  # assume function and execute
            else:
                eval(arg)  # assume Python code and execute
