# coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilties for data generation for Wikipedia Revision problem.
"""


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import os
import random
import re
import subprocess

import six

from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import text_encoder

import tensorflow as tf


def to_unicode(s):
  return unicode(s, "utf-8") if six.PY2 else s.decode("utf-8")


def include_revision(revision_num, skip_factor=1.1):
  """Decide whether to include a revision.

  If the number of revisions is large, we exclude some revisions to avoid
  a quadratic blowup in runtime, since the article is likely also large.

  We make the ratio between consecutive included revision numbers
  appproximately equal to "factor".

  Args:
    revision_num: an integer
    skip_factor: a floating point number >= 1.0

  Returns:
    a boolean
  """
  if skip_factor <= 1.0:
    return True
  return (int(math.log(revision_num + 1.0) / math.log(skip_factor)) != int(
      math.log(revision_num + 2.0) / math.log(skip_factor)))


def file_page_generator(my_file, max_page_size=2**28):
  """Read wikipedia pages from a history dump.

  Since some pages can be terabytes in size (with all the revisions),
  we limit page size to max_page_size bytes.

  Args:
    my_file: an open file object.
    max_page_size: an integer

  Yields:
    strings
  """
  page_start = "  <page>\n"
  page_end = "  </page>\n"
  chunk_size = max_page_size
  page_start = "  <page>\n"
  page_end = "  </page>\n"
  leftovers = ""
  while True:
    chunk = my_file.read(chunk_size)
    if not chunk:
      break
    chunk = leftovers + chunk
    current_pos = 0
    while True:
      start_pos = chunk.find(page_start, current_pos)
      if start_pos == -1:
        break
      end_pos = chunk.find(page_end, start_pos)
      if end_pos == -1:
        if len(chunk) - start_pos > max_page_size:
          leftovers = ""
        else:
          leftovers = chunk[start_pos:]
        break
      raw_page = chunk[start_pos + len(page_start):end_pos]
      if len(raw_page) < max_page_size:
        ret = parse_page(raw_page)
        if ret:
          yield ret
      current_pos = end_pos + len(page_end)


def get_title(page):
  """Extract the title from a page.

  Args:
    page: a string
  Returns:
    a string
  """
  start_pos = page.find("<title>")
  end_pos = page.find("</title>")
  assert start_pos != -1
  assert end_pos != -1
  start_pos += len("<title>")
  return to_unicode(page[start_pos:end_pos])


def get_id(page):
  """Extract the id from a page.

  Args:
    page: a string
  Returns:
    an integer
  """
  start_pos = page.find("<id>")
  end_pos = page.find("</id>")
  assert start_pos != -1
  assert end_pos != -1
  start_pos += len("<id>")
  return int(page[start_pos:end_pos])


def get_revisions(page):
  """Extract the revisions of a page.

  Args:
    page: a string
  Returns:
    a list of strings
  """
  start_string = "    <revision>\n"
  end_string = "    </revision>\n"
  ret = []
  current_pos = 0
  while True:
    start_pos = page.find(start_string, current_pos)
    if start_pos == -1:
      break
    end_pos = page.find(end_string, start_pos)
    assert end_pos != -1
    ret.append(page[start_pos + len(start_string):end_pos])
    current_pos = end_pos + len(end_string)
  return ret


def parse_page(raw_page):
  """Create a dictionary with title, id, and list of revisions.

  The dictionary contains:
  "title": a string
  "id": an integer
  "revisions": a list of strings

  Args:
    raw_page: a string

  Returns:
    a dictionary, or None in the case of an error.
  """
  ret = {"title": get_title(raw_page), "id": get_id(raw_page)}
  if ":" in ret["title"]:
    return None
  ret["revisions"] = get_revisions(raw_page)
  return ret


def maybe_copy_file_to_directory(source_filepath, target_directory):
  """Copy a file to a directory if it is not already there.

  Returns the target filepath.

  Args:
    source_filepath: a string
    target_directory: a string

  Returns:
    a string
  """
  if not tf.gfile.Exists(target_directory):
    tf.logging.info("Creating directory %s" % target_directory)
    os.mkdir(target_directory)
  target_filepath = os.path.join(target_directory,
                                 os.path.basename(source_filepath))
  if not tf.gfile.Exists(target_filepath):
    tf.logging.info("Copying %s to %s" % (source_filepath, target_filepath))
    tf.gfile.Copy(source_filepath, target_filepath)
    statinfo = os.stat(target_filepath)
    tf.logging.info("Successfully copied %s, %s bytes." % (target_filepath,
                                                           statinfo.st_size))
  else:
    tf.logging.info("Not copying, file already found: %s" % target_filepath)
  return target_filepath


def corpus_page_generator(corpus_files, tmp_dir, max_page_size_exp):
  """Generate pages from a list of .7z encoded history dumps.

  Args:
    corpus_files: a list of strings
    tmp_dir: a string
    max_page_size_exp: an integer

  Yields:
    strings
  """
  for remote_filepath in corpus_files:

    filepath = maybe_copy_file_to_directory(remote_filepath, tmp_dir)
    tf.logging.info("Reading from " + filepath)

    command = ["7z", "x", "-so", filepath]
    tf.logging.info("Running command: %s", command)

    p = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=-1)

    for page in file_page_generator(p.stdout, 2**max_page_size_exp):
      yield page


def get_text(revision, strip=True):
  """Extract the text from a revision.

  Args:
    revision: a string
    strip: a boolean

  Returns:
    a string
  """
  # text start tag looks like "<text ..otherstuff>"
  start_pos = revision.find("<text")
  assert start_pos != -1
  end_tag_pos = revision.find(">", start_pos)
  assert end_tag_pos != -1
  end_tag_pos += len(">")
  end_pos = revision.find("</text>")
  if end_pos == -1:
    ret = ""
  else:
    ret = revision[end_tag_pos:end_pos]
  if strip:
    ret = strip_text(ret)
  ret = to_unicode(ret)
  return ret


def strip_text(text):
  """Strip wikipedia-stuff out of text, making it mostly prose.

  The reason for this is to learn a model that is good at editing prose.

  Args:
    text: a string

  Returns:
    a string
  """
  return _remove_boring_lines(
      _remove_triple_quotes(
          _remove_double_brackets(
              _remove_references(_remove_curly_braces(text)))))


def _find_and_replace(text, start_string, end_string, replace_fn):
  """Remove everything found between instances of start_string and end_string.

  Replace each such instance with replace_fn(removed_text)

  e.g. _find_and_replace("the [[fat]] cat [[sat]]", "[[", "]]", lambda x: x)
    = "the fat cat sat"

  Args:
    text: a string
    start_string: a string
    end_string: a string
    replace_fn: a unary function from string to string

  Returns:
    a string
  """
  ret = ""
  current_pos = 0
  while True:
    start_pos = text.find(start_string, current_pos)
    if start_pos == -1:
      ret += text[current_pos:]
      break
    ret += text[current_pos:start_pos]
    end_pos = text.find(end_string, start_pos + len(start_string))
    if end_pos == -1:
      break
    ret += replace_fn(text[start_pos + len(start_string):end_pos])
    current_pos = end_pos + len(end_string)
  return ret


def _remove_references(text):
  return _find_and_replace(text, "&lt;ref", "&lt;/ref&gt;", lambda s: "")


def _remove_triple_quotes(text):
  return _find_and_replace(text, "'''", "'''", lambda s: s)


def _remove_curly_braces(text):
  """Remove everything in curly braces.

  Curly braces may be nested, so we keep track of depth.

  Args:
    text: a string
  Returns:
    a string
  """
  current_pos = 0
  depth = 0
  ret = ""
  for match in re.finditer("[{}]", text):
    if depth == 0: