#
# Copyright 2016 The BigDL 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.
#


import pytest
from unittest import TestCase
import tensorflow as tf
from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2
import numpy as np
from bigdl.nano.tf.keras import Model


class TestModelQuantize(TestCase):
    def test_model_quantize_ptq(self):
        model = MobileNetV2(weights=None, input_shape=[40, 40, 3], classes=10)
        model = Model(inputs=model.inputs, outputs=model.outputs)
        model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
                      loss=tf.keras.losses.BinaryCrossentropy(),
                      metrics=[tf.keras.metrics.CategoricalAccuracy()],)
        train_examples = np.random.random((100, 40, 40, 3))
        train_labels = np.random.randint(0, 10, size=(100,))
        train_dataset = tf.data.Dataset.from_tensor_slices((train_examples, train_labels))
        # Case 1: Default
        q_model = model.quantize(calib_dataset=train_dataset)
        assert q_model
        output = q_model(train_examples[0:10])
        assert output.shape == (10, 10)

        # Case 2: Override by arguments
        q_model = model.quantize(calib_dataset=train_dataset,
                                 metric=tf.keras.metrics.CategoricalAccuracy(),
                                 tuning_strategy='basic',
                                 accuracy_criterion={'relative': 0.99,
                                                     'higher_is_better': True})
        assert q_model
        output = q_model(train_examples[0:10])
        assert output.shape == (10, 10)

        # Case 3: Invalid approach, dynamic or qat is not supported
        invalid_approach = 'dynamic'
        with pytest.raises(RuntimeError, match="post_training_dynamic_quant is invalid."):
            model.quantize(approach=invalid_approach)

numpy.random.random tensorflow.data.Dataset.from_tensor_slices tensorflow.keras.applications.mobilenet_v2.MobileNetV2 tensorflow.keras.losses.BinaryCrossentropy tensorflow.keras.optimizers.Adam tensorflow.keras.metrics.CategoricalAccuracy numpy.random.randint
[(28, 'tensorflow.keras.applications.mobilenet_v2.MobileNetV2', 'MobileNetV2', ([], {'weights': 'None', 'input_shape': '[40, 40, 3]', 'classes': '(10)'}), False, 'from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2\n'), (29, 'bigdl.nano.tf.keras.Model', 'Model', ([], {'inputs': 'model.inputs', 'outputs': 'model.outputs'}), False, 'from bigdl.nano.tf.keras import Model\n'), (33, 'numpy.random.random', 'np.random.random', (['(100, 40, 40, 3)'], {}), True, 'import numpy as np\n'), (34, 'numpy.random.randint', 'np.random.randint', (['(0)', '(10)'], {'size': '(100,)'}), True, 'import numpy as np\n'), (35, 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', (['(train_examples, train_labels)'], {}), True, 'import tensorflow as tf\n'), (54, 'pytest.raises', 'pytest.raises', (['RuntimeError'], {'match': '"""post_training_dynamic_quant is invalid."""'}), False, 'import pytest\n'), (30, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': '(0.001)'}), True, 'import tensorflow as tf\n'), (31, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (44, 'tensorflow.keras.metrics.CategoricalAccuracy', 'tf.keras.metrics.CategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (32, 'tensorflow.keras.metrics.CategoricalAccuracy', 'tf.keras.metrics.CategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n')]