# coding=utf-8
# 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
#
#      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 compiler_opt.rl.compilation_runner."""

import os
import string
import subprocess
import threading
import time
from unittest import mock

from absl import flags
import tensorflow as tf

# This is https://github.com/google/pytype/issues/764
from google.protobuf import text_format  # pytype: disable=pyi-error

from compiler_opt.rl import compilation_runner
from compiler_opt.rl import constant
from compiler_opt.rl import corpus
from compiler_opt.rl import policy_saver

_DEFAULT_FEATURE_VALUE = 12
_POLICY_FEATURE_VALUE = 34

_DEFAULT_REWARD = 10
_POLICY_REWARD = 8

_MOVING_AVERAGE_DECAY_RATE = 0.8


def _get_sequence_example_with_reward(feature_value, reward):
  sequence_example_text = string.Template("""
    feature_lists {
      feature_list {
        key: "feature_0"
        value {
          feature { int64_list { value: $feature_value } }
          feature { int64_list { value: $feature_value } }
        }
      }
      feature_list {
        key: "reward"
        value {
          feature { float_list { value: $reward } }
          feature { float_list { value: $reward } }
        }
      }
    }""").substitute(
      feature_value=feature_value, reward=reward)
  return text_format.Parse(sequence_example_text, tf.train.SequenceExample())


def _get_sequence_example(feature_value):
  sequence_example_text = string.Template("""
    feature_lists {
      feature_list {
        key: "feature_0"
        value {
          feature { int64_list { value: $feature_value } }
          feature { int64_list { value: $feature_value } }
        }
      }
    }""").substitute(feature_value=feature_value)
  return text_format.Parse(sequence_example_text, tf.train.SequenceExample())


def _mock_compile_fn(file_paths, tf_policy_path, reward_only):  # pylint: disable=unused-argument
  del file_paths
  if tf_policy_path:
    sequence_example = _get_sequence_example(_POLICY_FEATURE_VALUE)
    native_size = _POLICY_REWARD
  else:
    sequence_example = _get_sequence_example(_DEFAULT_FEATURE_VALUE)
    native_size = _DEFAULT_REWARD

  if reward_only:
    return {'default': (None, native_size)}
  else:
    return {'default': (sequence_example, native_size)}


_mock_policy = policy_saver.Policy(bytes(), bytes())

_mock_loaded_module_spec = corpus.LoadedModuleSpec(
    name='dummy', loaded_ir=bytes())


class CompilationRunnerTest(tf.test.TestCase):

  def assertListProtoEqual(self, a, b):
    self.assertEqual(len(a), len(b))
    a.sort()
    b.sort()
    for (x, y) in zip(a, b):
      self.assertProtoEquals(x, y)

  @mock.patch(constant.BASE_MODULE_DIR +
              '.compilation_runner.CompilationRunner.compile_fn')
  def test_policy(self, mock_compile_fn):
    mock_compile_fn.side_effect = _mock_compile_fn
    runner = compilation_runner.CompilationRunner(
        moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
    data = runner.collect_data(
        loaded_module_spec=_mock_loaded_module_spec, policy=_mock_policy)
    self.assertEqual(2, mock_compile_fn.call_count)

    expected_example = _get_sequence_example_with_reward(
        _POLICY_FEATURE_VALUE, 0.1998)
    self.assertListProtoEqual([expected_example], [
        tf.train.SequenceExample.FromString(x)
        for x in data.serialized_sequence_examples
    ])
    self.assertEqual(2, data.length)
    self.assertCountEqual(
        {
            'default':
                compilation_runner.RewardStat(
                    default_reward=_DEFAULT_REWARD,
                    moving_average_reward=_DEFAULT_REWARD *
                    _MOVING_AVERAGE_DECAY_RATE + _POLICY_REWARD *
                    (1 - _MOVING_AVERAGE_DECAY_RATE))
        }, data.reward_stats)
    self.assertAllClose([0.1998002], data.rewards)
    self.assertAllClose([8], data.policy_rewards)

  @mock.patch(constant.BASE_MODULE_DIR +
              '.compilation_runner.CompilationRunner.compile_fn')
  def test_default(self, mock_compile_fn):
    mock_compile_fn.side_effect = _mock_compile_fn
    runner = compilation_runner.CompilationRunner(
        moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)

    data = runner.collect_data(loaded_module_spec=_mock_loaded_module_spec)
    # One call when we ask for the default policy, because it can provide both
    # trace and default size.
    self.assertEqual(1, mock_compile_fn.call_count)

    expected_example = _get_sequence_example_with_reward(
        _DEFAULT_FEATURE_VALUE, 0)
    self.assertListProtoEqual([expected_example], [
        tf.train.SequenceExample.FromString(x)
        for x in data.serialized_sequence_examples
    ])
    self.assertEqual(2, data.length)
    self.assertCountEqual(
        {
            'default':
                compilation_runner.RewardStat(
                    default_reward=_DEFAULT_REWARD,
                    moving_average_reward=_DEFAULT_REWARD)
        }, data.reward_stats)
    self.assertAllClose([0], data.rewards)
    self.assertAllClose([10], data.policy_rewards)

  @mock.patch(constant.BASE_MODULE_DIR +
              '.compilation_runner.CompilationRunner.compile_fn')
  def test_given_default_size(self, mock_compile_fn):
    mock_compile_fn.side_effect = _mock_compile_fn
    runner = compilation_runner.CompilationRunner(
        moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)

    data = runner.collect_data(
        loaded_module_spec=_mock_loaded_module_spec,
        policy=_mock_policy,
        reward_stat={
            'default':
                compilation_runner.RewardStat(
                    default_reward=_DEFAULT_REWARD, moving_average_reward=7)
        })
    self.assertEqual(1, mock_compile_fn.call_count)

    expected_example = _get_sequence_example_with_reward(
        _POLICY_FEATURE_VALUE,
        1 - (_POLICY_REWARD + constant.DELTA) / (7 + constant.DELTA))
    self.assertListProtoEqual([expected_example], [
        tf.train.SequenceExample.FromString(x)
        for x in data.serialized_sequence_examples
    ])
    self.assertEqual(2, data.length)
    self.assertCountEqual(
        {
            'default':
                compilation_runner.RewardStat(
                    default_reward=_DEFAULT_REWARD,
                    moving_average_reward=7 * _MOVING_AVERAGE_DECAY_RATE +
                    _POLICY_REWARD * (1 - _MOVING_AVERAGE_DECAY_RATE))
        }, data.reward_stats)
    self.assertAllClose([0.199800], data.rewards)
    self.assertAllClose([8], data.policy_rewards)

  @mock.patch(constant.BASE_MODULE_DIR +
              '.compilation_runner.CompilationRunner.compile_fn')
  def test_exception_handling(self, mock_compile_fn):
    mock_compile_fn.side_effect = subprocess.CalledProcessError(
        returncode=1, cmd='error')
    runner = compilation_runner.CompilationRunner(
        moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)

    with self.assertRaisesRegex(subprocess.CalledProcessError, 'error'):
      _ = runner.collect_data(
          loaded_module_spec=_mock_loaded_module_spec,
          policy=_mock_policy,
          reward_stat=None)
    self.assertEqual(1, mock_compile_fn.call_count)

  def test_start_subprocess_output(self):
    cm = compilation_runner.WorkerCancellationManager()
    output = compilation_runner.start_cancellable_process(
        ['ls', '-l'], timeout=100, cancellation_manager=cm, want_output=True)
    if output:
      output_str = output.decode('utf-8')
    else:
      self.fail('output should have been non-empty')
    self.assertNotEmpty(output_str)

  def test_timeout_kills_process(self):
    sentinel_file = os.path.join(flags.FLAGS.test_tmpdir,
                                 'test_timeout_kills_test_file')
    if os.path.exists(sentinel_file):
      os.remove(sentinel_file)
    with self.assertRaises(subprocess.TimeoutExpired):
      compilation_runner.start_cancellable_process(
          ['bash', '-c', 'sleep 1s ; touch ' + sentinel_file],
          timeout=0.5,
          cancellation_manager=None)
    time.sleep(2)
    self.assertFalse(os.path.exists(sentinel_file))

  def test_pause_resume(self):
    cm = compilation_runner.WorkerCancellationManager()
    start_time = time.time()

    def stop_and_start():
      time.sleep(0.25)
      cm.pause_all_processes()
      time.sleep(1)
      cm.resume_all_processes()

    threading.Thread(target=stop_and_start).start()
    compilation_runner.start_cancellable_process(['sleep', '0.5'],
                                                 30,
                                                 cancellation_manager=cm)
    # should be at least 1 second due to the pause.
    self.assertGreater(time.time() - start_time, 1)


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