# Copyright 2020 Google LLC
#
# 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
#
#     https://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.
""" Testing for preprocess_utils.py"""

import csv
from io import StringIO
import os
import subprocess
import unittest
from unittest import TestCase
from unittest.mock import patch

import preprocess_utils


class PreprocessUtilsTestCase(TestCase):
  def test_text_strip(self):
    # test remove escape characters
    test_case = ["\t", "\r", "\n", "\t" * 3, "\r" * 3, "\n" * 7]
    correct_ans = [" "] * 6
    result = preprocess_utils.text_strip(test_case)
    self.assertEqual(result, correct_ans)

    # test removing redundant special characters
    special_chars = ["_", "+", "-", "~", ":", "."]
    test_case = []
    correct_ans = []
    for char in special_chars:
      test_case.append("text" + char + "text")
      test_case.append("text" + char * 2 + "text")
      correct_ans += (["text" + char + "text"]) * 2
    result = preprocess_utils.text_strip(test_case)
    self.assertEqual(result, correct_ans)

    # test removing -, :, and _ at end or beginning of string (not in the middle)
    special_chars = ["-", ":", "_"]
    test_case = []
    correct_ans = []
    for char in special_chars:
      test_case.append("text" + char + "text")
      test_case.append(char + "text")
      test_case.append("text" + char)
      correct_ans += (["text" + char + "text"])
      correct_ans += (["text"]) * 2
    result = preprocess_utils.text_strip(test_case)
    self.assertEqual(result, correct_ans)

    # test removing ~ at end of string (not in the middle or beginning)
    test_case = ["text~", "text ~ text", "text~text", "~text"]
    correct_ans = ['text', 'text ~ text', 'text~text', '~text']
    result = preprocess_utils.text_strip(test_case)
    self.assertEqual(result, correct_ans)

    # test removing . at beginning of string (not in the middle or end)
    test_case = ["text.", "text . text", "text.text", ".text"]
    correct_ans = ['text.', 'text . text', 'text.text', 'text']
    result = preprocess_utils.text_strip(test_case)
    self.assertEqual(result, correct_ans)

  def test_validate_dataset_unequal(self):
    sentences = ["text"] * 5
    summaries = ["text"] * 7
    with self.assertRaises(Exception):
      preprocess_utils.validate_dataset(sentences, summaries)

    sentences = ["text"] * 7
    summaries = ["text"] * 7
    try:
      preprocess_utils.validate_dataset(sentences, summaries)
    except:
      self.fail("validate_dataset raised Exception unexpectedly!")

  def test_validate_dataset_empty(self):
    sentences = [" "] * 5
    summaries = ["text"] * 5
    with self.assertRaises(Exception):
      preprocess_utils.validate_dataset(sentences, summaries)

    summaries = [" "] * 5
    sentences = ["text"] * 5
    with self.assertRaises(Exception):
      preprocess_utils.validate_dataset(sentences, summaries)

    summaries = ["text"] * 3 + ["   "] * 2
    sentences = ["text"] * 5
    with self.assertRaises(Exception):
      preprocess_utils.validate_dataset(sentences, summaries)

  def test_calculate_stats(self):
    output_format = "-------Calculating statistics-------\n" + \
                    "Average word count of original sentence is {:.2f} ( std: {:.2f} )\n" + \
                    "Max word count is {}\n" + \
                    "Min word count is {}\n" + \
                    "Average word count of shortened sentence is {:.2f} ( std: {:.2f} )\n" + \
                    "Max Length is {}\n" + \
                    "Min Length is {}\n" + \
                    "On average, there are {:.2f} sentences in each original text ( std: {:.2f} )\n" + \
                    "On average, there are {:.2f} words in each shortened sentence that are not in the original " \
                    "sentence. ( std: {:.2f} )\n" + \
                    "The average compression ratio is {:.2f} ( std: {:.2f} )\n"

    sentences = ["sentences"]
    summaries = ["sum"]

    expected_output = output_format.format(1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 0,
                                           1, 0)
    with patch('sys.stdout', new=StringIO()) as func_output:
      preprocess_utils.calculate_stats(sentences, summaries)
      self.assertEqual(func_output.getvalue(), expected_output)

    sentences = ["sentences 1", "sentences number 2"]
    summaries = ["sum 1", "2"]

    expected_output = output_format.format(2.5, 0.5, 3, 2, 1.5, 0.5, 2, 1, 1,
                                           0, 0.5, 0.5, 0.67, 0.33)
    with patch('sys.stdout', new=StringIO()) as func_output:
      preprocess_utils.calculate_stats(sentences, summaries)
      self.assertEqual(func_output.getvalue(), expected_output)

    summaries = ["text"] * 3 + ["   "] * 2
    sentences = ["text"] * 5
    with self.assertRaises(Exception):
      preprocess_utils.calculate_stats(sentences, summaries)

  def test_tokenize_with_space(self):
    test_case = ["I've eaten.", "Have you?"]
    correct_ans = ["I 've eaten .", "Have you ?"]
    result = preprocess_utils.tokenize_with_space(test_case)
    self.assertEqual(result, correct_ans)