# coding: utf8
from __future__ import unicode_literals

import sys
import os
import unittest
import tempfile
import shutil
import gzip
import pickle

from nose.tools import assert_equal
from nose.tools import assert_not_equal
from nose.tools import assert_raises
from nose.tools import raises

import _setup_test_env  # noqa
from returnn.util import better_exchook
from returnn.datasets.lm import TranslationDataset, TranslationFactorsDataset
from returnn.util.basic import init_thread_join_hack

better_exchook.install()
better_exchook.replace_traceback_format_tb()
init_thread_join_hack()


dummy_source_text = ("This is some example text.\n"
                     "It is used to test the translation dataset.\n"
                     "We will write it into a temporary file.\n")

dummy_target_text = ("Das ist ein Beispieltext.\n"
                     "Er wird zum Testen des TranslationDatasets benutzt.\n"
                     "Wir werden ihn in eine temporäre Datei schreiben.\n")


def create_vocabulary(text):
  """
  :param str text: any natural text
  :return: mapping of words in the text to ids, as well as the inverse mapping
  :rtype: (dict[str, int], dict[int, str])
  """
  vocabulary = {word: index for index, word in enumerate(set(text.strip().split()))}
  inverse_vocabulary = {index: word for word, index in vocabulary.items()}

  return vocabulary, inverse_vocabulary


def word_ids_to_sentence(word_ids, vocabulary):
  """
  :param list[int] word_ids:
  :param dict[int, str] vocabulary: mapping from word ids to words
  :return: concatenation of all words
  :rtype: str
  """
  words = [vocabulary[word_id] for word_id in word_ids]
  return ' '.join(words)


def test_translation_dataset():
  """
  Checks whether a dummy translation dataset can be read and whether the returned word indices are correct.
  We create the necessary corpus and vocabulary files on the fly.
  """

  dummy_dataset = tempfile.mkdtemp()
  source_file_name = os.path.join(dummy_dataset, "source.test.gz")  # testing both zipped and unzipped
  target_file_name = os.path.join(dummy_dataset, "target.test")

  with gzip.open(source_file_name, "wb") as source_file:  # writing in binary format works both in Python 2 and 3
    source_file.write(dummy_source_text.encode("utf-8"))

  with open(target_file_name, "wb") as target_file:
    target_file.write(dummy_target_text.encode("utf-8"))

  for postfix in ["", " </S>"]:  # test with and without postfix

    # Replace one word by <UNK>. This way it will not appear in the vocabulary (and <UNK> is added to the vocabulary).
    # We will test below whether this word is assigned the unknown id by checking whether the reconstruction also
    # contains <UNK>. Note, that the input file is already written and contains the original word.
    dummy_target_text_with_unk = dummy_target_text.replace("TranslationDatasets", "<UNK>")

    # Append postfix just to have it in the vocabulary
    source_vocabulary, inverse_source_vocabulary = create_vocabulary(dummy_source_text + postfix)
    target_vocabulary, inverse_target_vocabulary = create_vocabulary(dummy_target_text_with_unk + postfix)

    source_vocabulary_file_name = os.path.join(dummy_dataset, "source.vocab.pkl")
    target_vocabulary_file_name = os.path.join(dummy_dataset, "target.vocab.pkl")

    with open(source_vocabulary_file_name, "wb") as source_vocabulary_file:
      pickle.dump(source_vocabulary, source_vocabulary_file)

    with open(target_vocabulary_file_name, "wb") as target_vocabulary_file:
      pickle.dump(target_vocabulary, target_vocabulary_file)

    translation_dataset = TranslationDataset(
      path=dummy_dataset, file_postfix="test", source_postfix=postfix, target_postfix=postfix,
      unknown_label={"classes": "<UNK>"})
    translation_dataset.init_seq_order()
    translation_dataset.load_seqs(0, 10)

    num_seqs = len(dummy_source_text.splitlines())
    assert_equal(translation_dataset.num_seqs, num_seqs)

    # Reconstruct the sentences from the word ids and compare with input.
    for sequence_index in range(num_seqs):
      source_word_ids = translation_dataset.get_data(sequence_index, "data")
      source_sentence = word_ids_to_sentence(source_word_ids, inverse_source_vocabulary)
      assert_equal(source_sentence, dummy_source_text.splitlines()[sequence_index] + postfix)

      target_word_ids = translation_dataset.get_data(sequence_index, "classes")
      target_sentence = word_ids_to_sentence(target_word_ids, inverse_target_vocabulary)
      assert_equal(target_sentence, dummy_target_text_with_unk.splitlines()[sequence_index] + postfix)

  shutil.rmtree(dummy_dataset)


num_source_factors = 2
dummy_source_text_factor_0 = ("This is some example text.\n"
                              "The factors here have no meaning\n")
dummy_source_text_factor_1 = ("a b c d e\n"
                              "a b c d e f\n")
dummy_source_text_factored_format = ("This|a is|b some|c example|d text.|e\n"
                                     "The|a factors|b here|c have|d no|e meaning|f\n")

num_target_factors = 3
dummy_target_text_factor_0 = ("Das ist ein Beispieltext.\n"
                              "Die Factors hier haben keinen Sinn.\n")
dummy_target_text_factor_1 = ("a b c d\n"
                              "a b c d e f\n")
dummy_target_text_factor_2 = ("1 2 3 4\n"
                              "1 2 3 4 5 6\n")
dummy_target_text_factored_format = ("Das|a|1 ist|b|2 ein|c|3 Beispieltext.|d|4\n"
                                     "Die|a|1 Factors|b|2 hier|c|3 haben|d|4 keinen|e|5 Sinn.|f|6\n")


def test_translation_factors_dataset():
  """
  Similar to test_translation_dataset(), but using translation factors.
  """
  source_text_per_factor = [dummy_source_text_factor_0, dummy_source_text_factor_1]
  target_text_per_factor = [dummy_target_text_factor_0, dummy_target_text_factor_1, dummy_target_text_factor_2]

  source_vocabulary_names = ["source.vocab.pkl", "source_factor1.vocab.pkl"]
  target_vocabulary_names = ["target.vocab.pkl", "target_factor1.vocab.pkl", "target_factor2.vocab.pkl"]

  source_data_keys = ["data", "source_factor1"]
  target_data_keys = ["classes", "target_factor1", "target_factor2"]

  dummy_dataset = tempfile.mkdtemp()
  source_file_name = os.path.join(dummy_dataset, "source.test.gz")  # testing both zipped and unzipped
  target_file_name = os.path.join(dummy_dataset, "target.test")

  with gzip.open(source_file_name, "wb") as source_file:  # writing in binary format works both in Python 2 and 3
    source_file.write(dummy_source_text_factored_format.encode("utf8"))

  with open(target_file_name, "wb") as target_file:
    target_file.write(dummy_target_text_factored_format.encode("utf8"))

  for postfix in ["", " </S>"]:  # test with and without postfix
    vocabularies, inverse_vocabularies = [], []
    for dummy_text in source_text_per_factor + target_text_per_factor:
      # Append postfix just to have it in the vocabulary
      vocabulary, inverse_vocabulary = create_vocabulary(dummy_text + postfix)
      vocabularies.append(vocabulary)
      inverse_vocabularies.append(inverse_vocabulary)

    vocabulary_names = source_vocabulary_names + target_vocabulary_names
    for index, vocabulary in enumerate(vocabularies):
      with open(os.path.join(dummy_dataset, vocabulary_names[index]), "wb") as vocabulary_file:
        pickle.dump(vocabulary, vocabulary_file)

    translation_dataset = TranslationFactorsDataset(
      path=dummy_dataset, file_postfix="test", factor_separator="|",
      source_factors=source_data_keys[1:], target_factors=target_data_keys[1:],
      source_postfix=postfix, target_postfix=postfix)

    translation_dataset.init_seq_order()
    translation_dataset.load_seqs(0, 10)

    num_seqs = len(dummy_target_text_factored_format.splitlines())
    assert_equal(translation_dataset.num_seqs, num_seqs)

    # Reconstruct the sentences from the word ids for all factors and compare with input.
    data_keys = source_data_keys + target_data_keys
    texts_per_factor = source_text_per_factor + target_text_per_factor
    for index, text in enumerate(texts_per_factor):
      for sequence_index in range(num_seqs):
        word_ids = translation_dataset.get_data(sequence_index, data_keys[index])
        sentence = word_ids_to_sentence(word_ids, inverse_vocabularies[index])
        assert_equal(sentence, text.splitlines()[sequence_index] + postfix)

  shutil.rmtree(dummy_dataset)


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
