"""ARK data file to TFRecords converter.

The scripts takes the ARK data file and optionally the indexed target text
to write aligned source and target data.
"""

from __future__ import print_function

import argparse
import io
import numpy as np
import tensorflow as tf

from opennmt.inputters.record_inputter import write_sequence_record


def consume_next_vector(ark_file):
  """Consumes the next vector.

  Args:
    ark_file: The ARK data file.

  Returns:
    The next vector as a 2D Numpy array.
  """
  idx = None
  vector = []

  for line in ark_file:
    line = line.strip()
    fields = line.split()

    if not idx:
      idx = fields[0]
      fields.pop(0)
      fields.pop(0)

    end = fields and fields[-1] == "]"

    if end:
      fields.pop()

    if fields:
      vector.append(fields)

    if end:
      break

  return idx, np.asarray(vector, dtype=tf.float32)

def consume_next_text(text_file):
  """Consumes the next text line from `text_file`."""
  idx = None
  text = text_file.readline()

  if text:
    tokens = text.strip().split()
    idx = tokens[0]
    tokens.pop(0)
    text = " ".join(tokens)

  return idx, text

def write_text(text, writer):
  """Serializes a line of text."""
  writer.write(text)
  writer.write("\n")

def ark_to_records_aligned(ark_filename, text_filename, out_prefix, compression_type=None):
  """Converts ARK and text datasets to aligned TFRecords and text datasets."""
  record_filename = "%s.records" % out_prefix
  if compression_type == "GZIP":
    record_filename = "%s.gz" % record_filename
  record_writer = tf.io.TFRecordWriter(record_filename, options=compression_type)
  text_writer = io.open(out_prefix + ".txt", encoding="utf-8", mode="w")

  ark_buffer = {}
  text_buffer = {}
  count = 0

  def _write_example(vector, text):
    write_sequence_record(vector, record_writer)
    write_text(text, text_writer)

  def _search_aligned():
    for idx in ark_buffer:
      if idx in text_buffer:
        vector = ark_buffer[idx]
        text = text_buffer[idx]

        del ark_buffer[idx]
        del text_buffer[idx]

        return vector, text

    return None, None

  with io.open(ark_filename, encoding="utf-8") as ark_file, open(text_filename, encoding="utf-8") as text_file: #pylint: disable=line-too-long
    while True:
      ark_idx, vector = consume_next_vector(ark_file)
      text_idx, text = consume_next_text(text_file)

      if not ark_idx and not text_idx:
        # Both files are empty.
        break

      if ark_idx == text_idx:
        # If the indices match, write the example.
        _write_example(vector, text)
        count += 1
      else:
        # Otherwise store the entries.
        if ark_idx:
          ark_buffer[ark_idx] = vector
        if text_idx:
          text_buffer[text_idx] = text

        # Look if we can now find aligned entries.
        vector, text = _search_aligned()

        if vector is not None:
          _write_example(vector, text)
          count += 1

  # Search alignments in stored entries.
  while True:
    vector, text = _search_aligned()
    if vector is None:
      break
    _write_example(vector, text)
    count += 1

  record_writer.close()
  text_writer.close()

  print("Saved {} aligned records.".format(count))

def ark_to_records(ark_filename, out_prefix, compression_type=None):
  """Converts ARK dataset to TFRecords."""
  record_writer = tf.io.TFRecordWriter(out_prefix + ".records", options=compression_type)
  count = 0

  with io.open(ark_filename, encoding="utf-8") as ark_file:
    while True:
      ark_idx, vector = consume_next_vector(ark_file)
      if not ark_idx:
        break
      write_sequence_record(vector, record_writer)
      count += 1

  record_writer.close()
  print("Saved {} records.".format(count))


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument("--ark", required=True,
                      help="Indexed ARK data file.")
  parser.add_argument("--txt",
                      help=("Indexed target text data file "
                            "(must set it to align source and target files)."))
  parser.add_argument("--out", required=True,
                      help="Output files prefix (will be suffixed by .records and .txt).")
  parser.add_argument("--compression_type", default=None, choices=["GZIP"],
                      help="Optional compression type.")
  args = parser.parse_args()

  if args.txt:
    ark_to_records_aligned(args.ark, args.txt, args.out, compression_type=args.compression_type)
  else:
    ark_to_records(args.ark, args.out, compression_type=args.compression_type)

if __name__ == "__main__":
  main()
