# coding=utf-8
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

m4_changequote(`<!', `!>')

import sys
from os import path
import math
# Add convex module to the pythonpath. # TODO: ?
sys.path.append(path.dirname(path.dirname(path.dirname(path.dirname(path.abspath(__file__))))))
sys.path.append(path.dirname(path.dirname(path.dirname(path.abspath(__file__)))))

import unittest
from mock import *
import plpy_mock as plpy

class HyperbandScheduleTestCase(unittest.TestCase):
    def setUp(self):
        # The side effects of this class(writing to the output table) are not
        # tested here. They are tested in dev-check.
        self.plpy_mock = Mock(spec='error')
        patches = {
            'plpy': plpy
        }

        self.plpy_mock_execute = MagicMock()
        plpy.execute = self.plpy_mock_execute

        self.module_patcher = patch.dict('sys.modules', patches)
        self.module_patcher.start()
        import deep_learning.madlib_keras_automl
        self.module = deep_learning.madlib_keras_automl
        # self.module.MstLoaderInputValidator._validate_input_args = \
        #     MagicMock()

        self.subject = self.module.HyperbandSchedule
        self.schedule_table = 'schedule_table'
        self.R = 81
        self.eta = 3
        self.skip_last = 0

    def test_schedule_table_dimension(self):
        generate_schedule = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths, len(generate_schedule.schedule_vals))

    def test_max_skip_last(self):
        self.skip_last = int(math.floor(math.log(self.R, self.eta)))+1 # s_max+1
        with self.assertRaises(plpy.PLPYException):
            generate_schedule = self.subject(
                self.schedule_table,
                self.R,
                self.eta,
                self.skip_last
            )
    def test_negative_skip_last(self):
        self.skip_last = -3
        with self.assertRaises(plpy.PLPYException):
            generate_schedule = self.subject(
                self.schedule_table,
                self.R,
                self.eta,
                self.skip_last
            )
    def test_zero_resources(self):
        self.R = 0
        with self.assertRaises(plpy.PLPYException):
            generate_schedule = self.subject(
                self.schedule_table,
                self.R,
                self.eta,
                self.skip_last
            )
    def test_negative_resources(self):
        self.R = -3
        with self.assertRaises(plpy.PLPYException):
            generate_schedule = self.subject(
                self.schedule_table,
                self.R,
                self.eta,
                self.skip_last
            )
    def test_non_discarding_eta(self):
        self.eta = 1
        with self.assertRaises(plpy.PLPYException):
            generate_schedule = self.subject(
                self.schedule_table,
                self.R,
                self.eta,
                self.skip_last
            )
    def test_negative_eta(self):
        self.eta = -2
        with self.assertRaises(plpy.PLPYException):
            generate_schedule = self.subject(
                self.schedule_table,
                self.R,
                self.eta,
                self.skip_last
            )

    def test_different_R(self):
        self.R = 27
        generate_schedule1 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths, len(generate_schedule1.schedule_vals))

        self.R = 13
        generate_schedule2 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths, len(generate_schedule2.schedule_vals))

        self.R = 100
        generate_schedule3 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths, len(generate_schedule3.schedule_vals))

    def test_different_eta(self):
        self.eta = 4
        generate_schedule3 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths, len(generate_schedule3.schedule_vals))

        self.R = 91
        self.eta = 6
        generate_schedule3 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths, len(generate_schedule3.schedule_vals))

    def test_different_skip_last(self):
        self.skip_last = 2
        generate_schedule3 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths - (2*(s_max+1)-1), len(generate_schedule3.schedule_vals))

        self.skip_last = 3
        generate_schedule3 = self.subject(
            self.schedule_table,
            self.R,
            self.eta,
            self.skip_last
        )
        s_max = int(math.floor(math.log(self.R, self.eta)))
        num_depths = int((s_max+1) * (s_max+2) / 2.0)
        self.assertEqual(num_depths - (3*(s_max+1)-3), len(generate_schedule3.schedule_vals))

    def tearDown(self):
        self.module_patcher.stop()

if __name__ == '__main__':
    unittest.main()