#!/usr/bin/env python3

"""
Dump translation dataset to HDF.
"""

from __future__ import print_function

import sys
import argparse
import gzip
import pickle
import typing

import numpy
import h5py
from operator import itemgetter
from itertools import islice, zip_longest

UNKNOWN_LABEL = "<UNK>"
POSTFIX = "</S>"


class HDFTranslationDatasetCreator(object):
  """
  Creates the same HDF file as hdf_dump.py called on a TranslationDataset, but does so much faster
  and using much less memory.

  Input parameters are basically the files in a TranslationDataset folder. Additionally, the number of lines has to
  be given.
  """

  def __init__(self, hdf_file, source_file, target_file, source_vocabularies, target_vocabularies,
               source_factors, target_factors, number_of_lines, factor_separator="|",
               compression_method=None, line_buffer_size=100000, data_buffer_size=5000000):
    """
    :param str hdf_file: filename for the hdf file being created
    :param str source_file: filename of the source text file
    :param str target_file: filename of the target text file
    :param list[str] source_vocabularies: Filenames of the source vocabularies (in pickle format).
      Usually there is just one. In case of source factors, provide one per factor.
    :param list[str] target_vocabularies: Filenames of the target vocabularies (in pickle format).
      Usually there is just one. In case of target factors, provide one per factor.
    :param list[str] source_factors: Data keys for the source factors. First data key is always 'data'
      and must not be contained in the list.
    :param list[str] target_factors: Data keys for the target factors. First data key is always 'classes'
      and must not be contained in the list.
    :param str factor_separator: "String used to separate factors of the words. E.g. if "|", words are expected to be
      of format "<lemma>|<factor>|...".
    :param int number_of_lines: line count in source_file and target_file
    :param str compression_method: Optional compression method as supported by h5py.File.create_dataset().
        Applied to the main data only ('inputs' and 'target/data/classes').
    :param int line_buffer_size: number of corpus lines to read and process at once
    :param int data_buffer_size: space to reserve inside the hdf file at once, in numbers of integers
    """

    self.hdf_file_name = hdf_file
    self.hdf_file = None
    self.compression_method = compression_method

    self.source_file_handle = self._open_file(source_file)
    self.target_file_handle = self._open_file(target_file)

    source_vocabularies = [self._read_vocabulary(v) for v in source_vocabularies]
    target_vocabularies = [self._read_vocabulary(v) for v in target_vocabularies]

    self.source_data_keys = ["data"] + source_factors
    self.target_data_keys = ["classes"] + target_factors

    self._vocabularies = {"source": source_vocabularies, "target": target_vocabularies}
    self._vocabulary_sizes = {
      "source": [len(v) for v in source_vocabularies],
      "target": [len(v) for v in target_vocabularies]}
    self._unknown_ids = {
      "source": [v.get(UNKNOWN_LABEL) for v in source_vocabularies],
      "target": [v.get(UNKNOWN_LABEL) for v in target_vocabularies]}

    self.number_of_lines = number_of_lines
    self.line_buffer_size = line_buffer_size
    self.data_buffer_size = data_buffer_size

    self._number_of_processed_lines = 0
    self._write_offsets = {data_key: 0 for data_key in self.source_data_keys + self.target_data_keys}

    self.factor_separator = factor_separator

  def create(self):
    """
    Main function writing the HDF file.
    """
    self._init_hdf_file()

    print("Setting attributes...", file=sys.stderr)
    sys.stderr.flush()
    self._write_attributes()
    print("Done.", file=sys.stderr)

    print("Writing labels (vocabulary)...", file=sys.stderr)
    sys.stderr.flush()
    self._write_labels()
    print("Done.", file=sys.stderr)

    print("Writing source, target, sequence lengths and tags for all lines...", file=sys.stderr)
    sys.stderr.flush()
    end_of_file = False
    while not end_of_file:
      end_of_file = self._write_data()
      if not end_of_file:
        print("> Processed {} lines.".format(self._number_of_processed_lines), file=sys.stderr)
        sys.stderr.flush()
    print("Done.", file=sys.stderr)
    sys.stderr.flush()

    self.hdf_file.close()

  def _init_hdf_file(self):
    """
    Sets up the HDF file and initializes the datasets that will be filled.
    """
    self.hdf_file = h5py.File(self.hdf_file_name, "w")

    self.hdf_file.create_group("targets/data")
    self.hdf_file.create_group("targets/size")
    self.hdf_file.create_group("targets/labels")

    num_data_keys = len(self.source_data_keys) + len(self.target_data_keys)
    self.hdf_file.create_dataset("seqLengths", (self.number_of_lines, num_data_keys), dtype="int32")
    max_tag_length = len("line-") + len(str(self.number_of_lines))
    self._tag_dtype = "S{}".format(max_tag_length)
    self.hdf_file.create_dataset("seqTags", (self.number_of_lines,), dtype=self._tag_dtype)

    self.hdf_file.create_dataset("inputs", (self.data_buffer_size,), maxshape=(None,),
                                 dtype="int32", compression=self.compression_method)

    # HDF format expects one input only, so store source factors as target too.
    for data_key in self.source_data_keys[1:] + self.target_data_keys:
      self.hdf_file["targets/data"].create_dataset(data_key, (self.data_buffer_size,), maxshape=(None,),
                                                   dtype="int32", compression=self.compression_method)

  def _write_attributes(self):
    """
    Writes several attributes to the HDF file.
    """
    for index, data_key in enumerate(self.source_data_keys[1:], start=1):
      self.hdf_file["targets/size"].attrs[data_key] = (self._vocabulary_sizes["source"][index], 1)
    for index, data_key in enumerate(self.target_data_keys):
      self.hdf_file["targets/size"].attrs[data_key] = (self._vocabulary_sizes["target"][index], 1)

    # Those should be deprecated, but include nevertheless to exactly reproduce hdf_dump.
    self.hdf_file.attrs["inputPattSize"] = self._vocabulary_sizes["source"][0]
    self.hdf_file.attrs["numLabels"] = self._vocabulary_sizes["target"][0]

  def _write_labels(self):
    """
    Writes the labels (i.e. target vocabulary) to the HDF file.
    """
    for side in ["source", "target"]:
      # We have to write it for the source factors too, because they are treated like targets.

      data_keys = self.source_data_keys if side == "source" else self.target_data_keys
      for index, data_key in enumerate(data_keys):
        if data_key == "data":
          continue

        sorted_vocabulary_tuples = sorted(self._vocabularies[side][index].items(), key=itemgetter(1))
        labels = [word.encode("utf8") for (word, _) in sorted_vocabulary_tuples]

        assert len(labels) == self._vocabulary_sizes[side][index], "Word ids were not unique."

        max_label_length = max([len(label) for label in labels])
        dtype = "S{}".format(max_label_length + 1)

        labels = [numpy.array(label, dtype=dtype, ndmin=1) for label in labels]
        labels = numpy.concatenate(labels)

        self.hdf_file["targets/labels"].create_dataset(
          data_key, (self._vocabulary_sizes[side][index],),
          data=labels, dtype=dtype)

  def _write_data(self):
    """
    Loads a chunk of lines from the corpus and writes all corresponding data to the HDF file.

    :return: whether the end of the corpus is reached
    :rtype: bool
    """
    data_chunks = self._get_chunk()

    if not data_chunks["data"]:
      self._finalize_data()
      return True

    source_lengths = [len(line) for line in data_chunks["data"]]
    target_lengths = [len(line) for line in data_chunks["classes"]]

    chunk_size = len(data_chunks["data"])

    self._write_lengths(source_lengths, target_lengths)
    self._write_tags(chunk_size)

    for data_key, data_chunk in data_chunks.items():
      self._write_data_indices(data_chunk, data_key)

    self._number_of_processed_lines += chunk_size

    return False

  def _get_chunk(self):
    """
    Reads in the next chunk of lines from the corpus files.

    :return: a dict 'data_key' -> word indices for 'data_key' (int32, shape [Lines]) for all source and target data
    :rtype: dict[str,list[numpy.ndarray]]
    """
    source_lines = islice(self.source_file_handle, self.line_buffer_size)
    target_lines = islice(self.target_file_handle, self.line_buffer_size)

    data_chunks = {data_key: [] for data_key in self.source_data_keys + self.target_data_keys}
    for source_line, target_line in zip_longest(source_lines, target_lines):
      assert source_line is not None and target_line is not None, "Number of source and target lines differ."

      source_indices = self._line_to_indices(source_line, "source")
      for data_key, indices in zip(self.source_data_keys, source_indices):
        data_chunks[data_key].append(indices)

      target_indices = self._line_to_indices(target_line, "target")
      for data_key, indices in zip(self.target_data_keys, target_indices):
        data_chunks[data_key].append(indices)

    return data_chunks

  def _write_lengths(self, source_lengths, target_lengths):
    """
    Writes the sequence lengths to the HDF file.

    :param list[int] source_lengths: lengths of all source lines in current chunk
    :param list[int] target_lengths: lengths of all target lines in current chunk
    """
    # We treat source factors as targets internally because HDF format does not support multiple inputs.

    # For each sequence, seqLengths is expected to contain the length of the input and the length of each of the
    # targets, in alphabetical order of the target data keys. As all source and all target factors share the
    # same lengths we just choose between the source and target lengths.
    target_sequence_lengths = {}
    for data_key in self.source_data_keys[1:]:
      target_sequence_lengths[data_key] = source_lengths
    for data_key in self.target_data_keys:
      target_sequence_lengths[data_key] = target_lengths

    # Now sort by key.
    key_lengths_tuples_sorted = sorted(target_sequence_lengths.items(), key=lambda x: x[0])  # type: typing.List[typing.Tuple[str,typing.List[int]]]  # nopep8
    target_lengths_sorted = [key_length_tuple[1] for key_length_tuple in key_lengths_tuples_sorted]

    # Finally, add one time the source lengths for the input ("data") and convert to numpy.
    lengths = numpy.array([source_lengths] + target_lengths_sorted, dtype="int32").transpose()

    offset = self._number_of_processed_lines
    assert len(lengths) + offset <= self.number_of_lines, "More lines in the corpus files than specified."

    self.hdf_file["seqLengths"][offset:offset + len(lengths), :] = lengths

  def _write_tags(self, chunk_size):
    """
    Writes the sequence tags to the HDF file.

    :param int chunk_size: number of lines in the current chunk
    """
    offset = self._number_of_processed_lines

    tags = [numpy.array("line-" + str(offset + i), dtype=self._tag_dtype, ndmin=1) for i in range(chunk_size)]
    tags = numpy.concatenate(tags)

    self.hdf_file["seqTags"][offset: offset + chunk_size] = tags

  def _write_data_indices(self, chunk, data_key):
    """
    Writes the main data (word indices for the source or target corpus) to the HDF file.

    :param list[numpy.ndarray] chunk: word indices for all lines in the current chunk
    :param str data_key: "data", "classes" or a name of a factor
    """
    indices = numpy.concatenate(chunk)

    if data_key == "data":
      dataset = self.hdf_file["inputs/"]
    else:
      dataset = self.hdf_file["targets/data/{}".format(data_key)]

    offset = self._write_offsets[data_key]
    length = len(indices)
    if offset + length > len(dataset):
      buffer_size = max(self.data_buffer_size, length)
      dataset.resize((offset + buffer_size,))

    dataset[offset:offset + length] = indices

    self._write_offsets[data_key] += length

  def _finalize_data(self):
    """
    Called after all data is written. Checks number of lines and resizes datasets down to actual data size.
    """
    # Make sure the number of lines given by the user was correct.
    # Otherwise lengths and labels would have trailing zeros.
    assert self.number_of_lines == self._number_of_processed_lines, (
      "Fewer lines ({}) in the corpus files "
      "than specified ({}).".format(self._number_of_processed_lines, self.number_of_lines))

    # Trim datasets to actually occupied length, i.e. remove unused reserved space.
    self.hdf_file["inputs"].resize((self._write_offsets["data"],))
    for data_key in self.source_data_keys[1:] + self.target_data_keys:
      self.hdf_file["targets/data/{}".format(data_key)].resize((self._write_offsets[data_key],))

  def _line_to_indices(self, line, side):
    """
    Converts a line of text to arrays of word indices.

    :param str line: input line
    :param str side: "source" or "target"
    :return: word indices (int32, shape [num_words]) for all source or target factors
    :rtype: list[numpy.ndarray]
    """
    data_keys = self.source_data_keys if side == "source" else self.target_data_keys

    words = line.strip().split()

    if len(data_keys) == 1:
      word_list_per_factor = [words + [POSTFIX]]
    else:
      if words:
          words_split_into_factors = [word.split(self.factor_separator) for word in words]
          assert all(len(factors) == len(data_keys) for factors in words_split_into_factors), (
            "All words must have all factors. Expected: " + self.factor_separator.join(data_keys))
          word_list_per_factor = zip(*words_split_into_factors)
      else:
          word_list_per_factor = [[]] * len(data_keys)
      word_list_per_factor = [list(words) + [POSTFIX] for words in word_list_per_factor]

    indices_list = []
    for index, (data_key, words) in enumerate(zip(data_keys, word_list_per_factor)):
      indices = [self._vocabularies[side][index].get(word, self._unknown_ids[side][index]) for word in words]

      indices_numpy = numpy.array(indices, dtype=numpy.int32)
      indices_list.append(indices_numpy)

    return indices_list

  @staticmethod
  def _open_file(file_name):
    """
    :param str file_name: filename of a plain text file, possibly zipped
    :return: file handle
    :rtype: io.TextIOWrapper|gzip.GzipFile
    """
    if file_name.endswith(".gz"):
      return gzip.open(file_name, "rt")
    else:
      return open(file_name, "r")

  @staticmethod
  def _read_vocabulary(file_name):
    """
    :param str file_name: filename of the vocabulary (in pickle format)
    :return: mapping from words to indices
    :rtype: dict[str,int]
    """
    file_handle = open(file_name, "rb")
    vocabulary = pickle.load(file_handle)
    return vocabulary


def parse_args():
  """
  :rtype: argparse.Namespace
  """
  parser = argparse.ArgumentParser()
  parser.add_argument("-s", "--source_corpus", required=True, help="Source corpus file, possibly zipped.")
  parser.add_argument("-t", "--target_corpus", required=True, help="Target corpus file, possibly zipped.")
  parser.add_argument(
    "-v", "--source_vocabulary", required=True, help="Source vocabulary in pickle format."
    "In case of source factors provide a comma separated list containing vocabularies for each factor.")
  parser.add_argument(
    "-w", "--target_vocabulary", required=True, help="Target vocabulary in pickle format."
    "In case of target factors provide a comma separated list containing vocabularies for each factor.")
  parser.add_argument("-o", "--hdf_file", required=True, help="Output HDF file name.")
  parser.add_argument("--source_factors", help="Comma separated list of data keys for the source factors.")
  parser.add_argument("--target_factors", help="Comma separated list of data keys for the target factors.")
  parser.add_argument(
    "--factor_separator", default="|",
    help="String used to separate factors of the words, E.g. if '|', words are expected to be "
    "of format '<lemma>|<factor>|...'")
  parser.add_argument(
    "-n", "--number_of_lines", required=True, type=int,
    help="The number of total lines in the corpus files.")
  parser.add_argument("-c", "--compression", help="Type of compression (e.g. 'gzip', 'lzf'). Turned off if not given.")
  parser.add_argument("-l", "--line_buffer_size", type=int, help="How many lines to read at once.", default=100000)
  parser.add_argument(
    "-d", "--data_buffer_size", type=int, help="How much space to reserve in the HDF dataset "
    "at once (in number of integers).", default=5000000)

  return parser.parse_args()


def main():
  """
  Main entry.
  """
  args = parse_args()

  # In case of source or target factors we need a vocabularies for each.
  source_vocabularies = args.source_vocabulary.split(",")
  target_vocabularies = args.target_vocabulary.split(",")

  source_factors = args.source_factors.split(",") if args.source_factors else []
  target_factors = args.target_factors.split(",") if args.target_factors else []

  assert len(source_factors) + 1 == len(source_vocabularies), (
    "Number of source factors must be one less "
    "than number of source vocabularies (first factor is always called 'data')")
  assert len(target_factors) + 1 == len(target_vocabularies), (
    "Number of target factors must be one less "
    "than number of target vocabularies (first factor is always called 'classes')")

  HDFTranslationDatasetCreator(
    args.hdf_file, args.source_corpus, args.target_corpus,
    source_vocabularies, target_vocabularies, source_factors, target_factors, args.number_of_lines,
    args.factor_separator,
    args.compression, args.line_buffer_size, args.data_buffer_size).create()


if __name__ == "__main__":
  sys.exit(main())
