# Copyright 2019 The Google Research 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.

"""Unit tests for scripts/run_evaluation.py."""

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

import json
import os
import tempfile

from absl.testing import flagsaver
from nasbench.scripts import run_evaluation
import tensorflow as tf


class RunEvaluationTest(tf.test.TestCase):
    def setUp(self):
        """Set up files and directories that are expected by run_evaluation."""
        # Create temp directory for output files
        self.output_dir = tempfile.mkdtemp()
        self.models_file = os.path.join(self.output_dir, "models_file.json")

        self.toy_data = {
            "abc": ([[0, 1, 1], [0, 0, 1], [0, 0, 0]], [-1, 0, -2]),
            "abd": ([[0, 1, 0], [0, 0, 1], [0, 0, 0]], [-1, 0, -2]),
            "abe": ([[0, 0, 1], [0, 0, 0], [0, 0, 0]], [-1, 0, -2]),
        }

        with tf.io.gfile.GFile(self.models_file, "w") as f:
            json.dump(self.toy_data, f)

        # Create files & directories which are normally created by
        # evaluate.train_and_evaluate but have been mocked out.
        for model_id in self.toy_data:
            eval_dir = os.path.join(self.output_dir, "ab", model_id, "repeat_1")
            tf.io.gfile.makedirs(eval_dir)
        run_evaluation.FLAGS.train_data_files = "unused"
        run_evaluation.FLAGS.valid_data_file = "unused"
        run_evaluation.FLAGS.test_data_file = "unused"
        run_evaluation.FLAGS.num_repeats = 1

    @tf.compat.v1.test.mock.patch.object(run_evaluation, "evaluate")
    def test_evaluate_single_worker(self, mock_eval):
        """Tests single worker code path."""
        mock_eval.train_and_evaluate.return_value = "unused_output"
        evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir)
        evaluator.run_evaluation()

        expected_dir = os.path.join(self.output_dir, "ab")
        mock_eval.train_and_evaluate.assert_has_calls(
            [
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abc", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abd", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abe", "repeat_1"),
                ),
            ]
        )

        for model_id in self.toy_data:
            self.assertTrue(tf.io.gfile.exists(os.path.join(expected_dir, model_id, "repeat_1", "results.json")))

    @tf.compat.v1.test.mock.patch.object(run_evaluation, "evaluate")
    def test_evaluate_multi_worker_0(self, mock_eval):
        """Tests multi worker code path for worker 0."""
        mock_eval.train_and_evaluate.return_value = "unused_output"
        evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir, worker_id=0, total_workers=2)
        evaluator.run_evaluation()

        expected_dir = os.path.join(self.output_dir, "ab")
        mock_eval.train_and_evaluate.assert_has_calls(
            [
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abc", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abe", "repeat_1"),
                ),
            ]
        )

        for model_id in ["abc", "abe"]:
            self.assertTrue(tf.io.gfile.exists(os.path.join(expected_dir, model_id, "repeat_1", "results.json")))

    @tf.compat.v1.test.mock.patch.object(run_evaluation, "evaluate")
    def test_evaluate_multi_worker_1(self, mock_eval):
        """Tests multi worker code path for worker 1."""
        mock_eval.train_and_evaluate.return_value = "unused_output"
        evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir, worker_id=1, total_workers=2)
        evaluator.run_evaluation()

        expected_dir = os.path.join(self.output_dir, "ab")
        mock_eval.train_and_evaluate.assert_has_calls(
            [
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abd", "repeat_1"),
                )
            ]
        )

        self.assertTrue(tf.io.gfile.exists(os.path.join(expected_dir, "abd", "repeat_1", "results.json")))

    @tf.compat.v1.test.mock.patch.object(run_evaluation, "evaluate")
    def test_evaluate_regex(self, mock_eval):
        """Tests regex filters models."""
        mock_eval.train_and_evaluate.return_value = "unused_output"
        evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir, model_id_regex="^ab(d|e)")
        evaluator.run_evaluation()

        expected_dir = os.path.join(self.output_dir, "ab")
        mock_eval.train_and_evaluate.assert_has_calls(
            [
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abd", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abe", "repeat_1"),
                ),
            ]
        )

        for model_id in ["abd", "abe"]:
            self.assertTrue(tf.io.gfile.exists(os.path.join(expected_dir, model_id, "repeat_1", "results.json")))

    @tf.compat.v1.test.mock.patch.object(run_evaluation, "evaluate")
    def test_evaluate_repeat(self, mock_eval):
        """Tests evaluate with repeats."""
        mock_eval.train_and_evaluate.return_value = "unused_output"

        # Create extra directories not created in setUp for repeat_2
        for model_id in self.toy_data:
            eval_dir = os.path.join(self.output_dir, "ab", model_id, "repeat_2")
            tf.io.gfile.makedirs(eval_dir)

        with flagsaver.flagsaver(num_repeats=2):
            evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir)
            evaluator.run_evaluation()

        expected_dir = os.path.join(self.output_dir, "ab")
        mock_eval.train_and_evaluate.assert_has_calls(
            [
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abc", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abd", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abe", "repeat_1"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abc", "repeat_2"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abd", "repeat_2"),
                ),
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abe", "repeat_2"),
                ),
            ]
        )

        for model_id in self.toy_data:
            for repeat in range(2):
                self.assertTrue(
                    tf.io.gfile.exists(os.path.join(expected_dir, model_id, "repeat_%d" % (repeat + 1), "results.json"))
                )

    def test_clean_model_dir(self):
        """Tests clean-up of model directory keeps only intended files."""
        model_dir = os.path.join(self.output_dir, "ab", "abcde", "repeat_1")
        tf.io.gfile.makedirs(model_dir)

        # Write files which will be preserved
        preserved_files = ["model.ckpt-0.index", "model.ckpt-100.index", "results.json"]
        for filename in preserved_files:
            with tf.io.gfile.GFile(os.path.join(model_dir, filename), "w") as f:
                f.write("unused")

        # Write files which will be deleted
        for filename in ["checkpoint", "events.out.tfevents"]:
            with tf.io.gfile.GFile(os.path.join(model_dir, filename), "w") as f:
                f.write("unused")

        # Create subdirectory which will be deleted
        eval_dir = os.path.join(model_dir, "eval_dir")
        tf.io.gfile.makedirs(eval_dir)
        with tf.io.gfile.GFile(os.path.join(eval_dir, "events.out.tfevents"), "w") as f:
            f.write("unused")

        evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir)
        evaluator._clean_model_dir(model_dir)

        # Check only intended files are preserved
        remaining_files = tf.io.gfile.listdir(model_dir)
        self.assertItemsEqual(remaining_files, preserved_files)

    @tf.compat.v1.test.mock.patch.object(run_evaluation, "evaluate")
    def test_recovery_file(self, mock_eval):
        """Tests that evaluation recovers from restart."""
        mock_eval.train_and_evaluate.return_value = "unused_output"

        # Write recovery file
        recovery_dir = os.path.join(self.output_dir, "_recovery")
        tf.io.gfile.makedirs(recovery_dir)
        with tf.io.gfile.GFile(os.path.join(recovery_dir, "0"), "w") as f:
            f.write("2")  # Resume at 3rd entry

        evaluator = run_evaluation.Evaluator(self.models_file, self.output_dir)
        evaluator.run_evaluation()

        expected_dir = os.path.join(self.output_dir, "ab")
        mock_eval.train_and_evaluate.assert_has_calls(
            [
                tf.compat.v1.test.mock.call(
                    tf.compat.v1.test.mock.ANY,
                    tf.compat.v1.test.mock.ANY,
                    os.path.join(expected_dir, "abe", "repeat_1"),
                )
            ]
        )

        # Check that only 'abe' was evaluated, 'abc' and 'abe' are skipped due to
        # recovery.
        call_args = mock_eval.train_and_evaluate.call_args_list
        self.assertEqual(len(call_args), 1)

        # Check that recovery file is updated after run
        with tf.io.gfile.GFile(evaluator.recovery_file) as f:
            new_idx = int(f.read())
        self.assertEqual(new_idx, 3)


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