# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Tests for preprocessing."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools
import tempfile

import tensorflow as tf  # pylint: disable=g-bad-import-order

from research.minigo import coords
from research.minigo import features
from research.minigo import go
from research.minigo import model_params
import numpy as np
from research.minigo import preprocessing
from research.minigo import utils_test

tf.logging.set_verbosity(tf.logging.ERROR)

TEST_SGF = '''(;CA[UTF-8]SZ[9]PB[Murakawa Daisuke]PW[Iyama Yuta]KM[6.5]
           HA[0]RE[W+1.5]GM[1];B[fd];W[cf])'''


def create_random_data(num_examples):
    raw_data = []
    for _ in range(num_examples):
        feature = np.random.random([
            utils_test.BOARD_SIZE, utils_test.BOARD_SIZE,
            features.NEW_FEATURES_PLANES]).astype(np.uint8)
        pi = np.random.random([utils_test.BOARD_SIZE * utils_test.BOARD_SIZE
                               + 1]).astype(np.float32)
        value = np.random.random()
        raw_data.append((feature, pi, value))
    return raw_data


class TestPreprocessing(utils_test.MiniGoUnitTest):

    def extract_data(self, tf_record, filter_amount=1):
        pos_tensor, label_tensors = preprocessing.get_input_tensors(
            model_params.DummyMiniGoParams(), 1, [tf_record], num_repeats=1,
            shuffle_records=False, shuffle_examples=False,
            filter_amount=filter_amount)
        recovered_data = []
        with tf.Session() as sess:
            while True:
                try:
                    pos_value, label_values = sess.run([pos_tensor, label_tensors])
                    recovered_data.append((
                        pos_value,
                        label_values['pi_tensor'],
                        label_values['value_tensor']))
                except tf.errors.OutOfRangeError:
                    break
        return recovered_data

    def assertEqualData(self, data1, data2):
        # Assert that two data are equal, where both are of form:
        # data = List<Tuple<feature_array, pi_array, value>>
        self.assertEqual(len(data1), len(data2))
        for datum1, datum2 in zip(data1, data2):
            # feature
            self.assertEqualNPArray(datum1[0], datum2[0])
            # pi
            self.assertEqualNPArray(datum1[1], datum2[1])
            # value
            self.assertEqual(datum1[2], datum2[2])

    def test_serialize_round_trip(self):
        np.random.seed(1)
        raw_data = create_random_data(10)
        tfexamples = list(map(preprocessing.make_tf_example, *zip(*raw_data)))

        with tempfile.NamedTemporaryFile() as f:
            preprocessing.write_tf_examples(f.name, tfexamples)
            recovered_data = self.extract_data(f.name)

        self.assertEqualData(raw_data, recovered_data)

    def test_filter(self):
        raw_data = create_random_data(100)
        tfexamples = list(map(preprocessing.make_tf_example, *zip(*raw_data)))

        with tempfile.NamedTemporaryFile() as f:
            preprocessing.write_tf_examples(f.name, tfexamples)
            recovered_data = self.extract_data(f.name, filter_amount=.05)

        self.assertLess(len(recovered_data), 50)

    def test_serialize_round_trip_no_parse(self):
        np.random.seed(1)
        raw_data = create_random_data(10)
        tfexamples = list(map(preprocessing.make_tf_example, *zip(*raw_data)))

        with tempfile.NamedTemporaryFile() as start_file, \
                tempfile.NamedTemporaryFile() as rewritten_file:
            preprocessing.write_tf_examples(start_file.name, tfexamples)
            # We want to test that the rewritten, shuffled file contains correctly
            # serialized tf.Examples.
            batch_size = 4
            batches = list(preprocessing.shuffle_tf_examples(
                1000, batch_size, [start_file.name]))
            # 2 batches of 4, 1 incomplete batch of 2.
            self.assertEqual(len(batches), 3)

            # concatenate list of lists into one list
            all_batches = list(itertools.chain.from_iterable(batches))

            for _ in batches:
                preprocessing.write_tf_examples(
                    rewritten_file.name, all_batches, serialize=False)

            original_data = self.extract_data(start_file.name)
            recovered_data = self.extract_data(rewritten_file.name)

        # stuff is shuffled, so sort before checking equality
        def sort_key(nparray_tuple):
            return nparray_tuple[2]

        original_data = sorted(original_data, key=sort_key)
        recovered_data = sorted(recovered_data, key=sort_key)

        self.assertEqualData(original_data, recovered_data)

    def test_make_dataset_from_sgf(self):
        with tempfile.NamedTemporaryFile() as sgf_file, \
                tempfile.NamedTemporaryFile() as record_file:
            sgf_file.write(TEST_SGF.encode('utf8'))
            sgf_file.seek(0)
            preprocessing.make_dataset_from_sgf(
                utils_test.BOARD_SIZE, sgf_file.name, record_file.name)
            recovered_data = self.extract_data(record_file.name)
        start_pos = go.Position(utils_test.BOARD_SIZE)
        first_move = coords.from_sgf('fd')
        next_pos = start_pos.play_move(first_move)
        second_move = coords.from_sgf('cf')
        expected_data = [
            (
                features.extract_features(utils_test.BOARD_SIZE, start_pos),
                preprocessing._one_hot(utils_test.BOARD_SIZE, coords.to_flat(
                    utils_test.BOARD_SIZE, first_move)), -1
            ),
            (
                features.extract_features(utils_test.BOARD_SIZE, next_pos),
                preprocessing._one_hot(utils_test.BOARD_SIZE, coords.to_flat(
                    utils_test.BOARD_SIZE, second_move)), -1
            )
        ]
        self.assertEqualData(expected_data, recovered_data)


if __name__ == '__main__':
    tf.test.main()
