
# -*- coding: utf8 -*-

import sys
import os

import _setup_test_env  # noqa
sys.path += ["tools"]

from hdf_dump import *
import os
from returnn.log import log
import tempfile
from returnn.datasets.generating import DummyDataset
from returnn.datasets.hdf import HDFDataset
from returnn.util.basic import DictAsObj
import unittest
from returnn.util import better_exchook


options = {
  "epoch": 1,
  "start_seq": 0,
  "end_seq": float("inf")
}


def test_hdf_dataset_init():
  hdf_filename = tempfile.mktemp(suffix=".hdf", prefix="nose-dataset-init")
  hdf_dataset_init(hdf_filename)
  assert os.path.exists(hdf_filename)
  os.remove(hdf_filename)


def test_hdf_create():
  hdf_filename = tempfile.mktemp(suffix=".hdf", prefix="nose-dataset-create")
  hdf_dataset = hdf_dataset_init(hdf_filename)
  assert os.path.exists(hdf_filename)

  dataset = DummyDataset(input_dim=2, output_dim=3, num_seqs=4)
  dataset.init_seq_order(epoch=1)

  hdf_dump_from_dataset(dataset, hdf_dataset, DictAsObj(options))
  hdf_close(hdf_dataset)

  os.remove(hdf_filename)


def test_hdf_create_and_load():
  hdf_filename = tempfile.mktemp(suffix=".hdf", prefix="nose-dataset-load")
  hdf_dataset = hdf_dataset_init(hdf_filename)
  assert os.path.exists(hdf_filename)

  dataset = DummyDataset(input_dim=2, output_dim=3, num_seqs=4)
  dataset.init_seq_order(epoch=1)

  hdf_dump_from_dataset(dataset, hdf_dataset, DictAsObj(options))
  hdf_close(hdf_dataset)

  loaded_dataset = HDFDataset()
  loaded_dataset.add_file(hdf_filename)

  os.remove(hdf_filename)


def test_hdf_create_unicode_labels():
  hdf_filename = tempfile.mktemp(suffix=".hdf", prefix="nose-dataset-create")
  hdf_dataset = hdf_dataset_init(hdf_filename)
  assert os.path.exists(hdf_filename)

  dataset = DummyDataset(input_dim=2, output_dim=3, num_seqs=4)
  assert "classes" in dataset.get_target_list()
  dataset.labels["classes"] = ['’', 'ä', 'x']  # have some Unicode chars here
  dataset.init_seq_order(epoch=1)

  hdf_dump_from_dataset(dataset, hdf_dataset, DictAsObj(options))
  hdf_close(hdf_dataset)

  os.remove(hdf_filename)


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
