"""Tests for quantization"""
import numpy as np
import unittest
import os
import shutil
import yaml
import tensorflow as tf


def build_fake_yaml():
    fake_yaml = '''
        model:
          name: fake_yaml
          framework: tensorflow
          inputs: x
          outputs: op_to_store
        device: cpu
        evaluation:
          accuracy:
            metric:
              topk: 1
        tuning:
            strategy:
              name: random
            accuracy_criterion:
              relative: 0.01
            workspace:
              path: saved
        '''
    y = yaml.load(fake_yaml, Loader=yaml.SafeLoader)
    with open('fake_yaml.yaml', "w", encoding="utf-8") as f:
        yaml.dump(y, f)
    f.close()


def build_fake_yaml2():
    fake_yaml = '''
        model:
          name: fake_yaml
          framework: tensorflow
          inputs: x
          outputs: op_to_store
        device: cpu
        evaluation:
          accuracy:
            metric:
              topk: 1
        tuning:
          strategy:
            name: random
          exit_policy:
            max_trials: 5
          accuracy_criterion:
            relative: -0.01
          workspace:
            path: saved
        '''
    y = yaml.load(fake_yaml, Loader=yaml.SafeLoader)
    with open('fake_yaml2.yaml', "w", encoding="utf-8") as f:
        yaml.dump(y, f)
    f.close()


def build_fake_model():
    try:
        graph = tf.Graph()
        graph_def = tf.GraphDef()
        with tf.Session() as sess:
            x = tf.placeholder(tf.float64, shape=(1, 3, 3, 1), name='x')
            y = tf.constant(np.random.random((2, 2, 1, 1)), name='y')
            op = tf.nn.conv2d(input=x, filter=y, strides=[
                              1, 1, 1, 1], padding='VALID', name='op_to_store')

            sess.run(tf.global_variables_initializer())
            constant_graph = tf.graph_util.convert_variables_to_constants(
                sess, sess.graph_def, ['op_to_store'])

        graph_def.ParseFromString(constant_graph.SerializeToString())
        with graph.as_default():
            tf.import_graph_def(graph_def, name='')
    except:
        graph = tf.Graph()
        graph_def = tf.compat.v1.GraphDef()
        with tf.compat.v1.Session() as sess:
            x = tf.compat.v1.placeholder(tf.float64, shape=(1, 3, 3, 1), name='x')
            y = tf.compat.v1.constant(np.random.random((2, 2, 1, 1)), name='y')
            op = tf.nn.conv2d(input=x, filters=y, strides=[
                              1, 1, 1, 1], padding='VALID', name='op_to_store')

            sess.run(tf.compat.v1.global_variables_initializer())
            constant_graph = tf.compat.v1.graph_util.convert_variables_to_constants(sess, sess.graph_def, [
                                                                                    'op_to_store'])

        graph_def.ParseFromString(constant_graph.SerializeToString())
        with graph.as_default():
            tf.import_graph_def(graph_def, name='')
    return graph


class TestQuantization(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        self.constant_graph = build_fake_model()
        build_fake_yaml()
        build_fake_yaml2()

    @classmethod
    def tearDownClass(self):
        os.remove('fake_yaml.yaml')
        os.remove('fake_yaml2.yaml')

        shutil.rmtree("saved", ignore_errors=True)

    def test_ru_random_one_trial(self):
        from neural_compressor.experimental import Quantization, common

        quantizer = Quantization('fake_yaml.yaml')
        dataset = quantizer.dataset('dummy', (100, 3, 3, 1), label=True)
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        quantizer.model = self.constant_graph
        quantizer()

    def test_ru_random_max_trials(self):
        from neural_compressor.experimental import Quantization, common

        quantizer = Quantization('fake_yaml2.yaml')
        dataset = quantizer.dataset('dummy', (100, 3, 3, 1), label=True)
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        quantizer.model = self.constant_graph
        quantizer()


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

[(30, 'yaml.load', 'yaml.load', False, 'import yaml\n'), (58, 'yaml.load', 'yaml.load', False, 'import yaml\n'), (137, 'unittest.main', 'unittest.main', False, 'import unittest\n'), (32, 'yaml.dump', 'yaml.dump', False, 'import yaml\n'), (60, 'yaml.dump', 'yaml.dump', False, 'import yaml\n'), (66, 'tensorflow.Graph', 'tf.Graph', True, 'import tensorflow as tf\n'), (67, 'tensorflow.GraphDef', 'tf.GraphDef', True, 'import tensorflow as tf\n'), (110, 'os.remove', 'os.remove', False, 'import os\n'), (111, 'os.remove', 'os.remove', False, 'import os\n'), (113, 'shutil.rmtree', 'shutil.rmtree', False, 'import shutil\n'), (118, 'neural_compressor.experimental.Quantization', 'Quantization', False, 'from neural_compressor.experimental import Quantization, common\n'), (120, 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', False, 'from neural_compressor.experimental import Quantization, common\n'), (121, 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', False, 'from neural_compressor.experimental import Quantization, common\n'), (128, 'neural_compressor.experimental.Quantization', 'Quantization', False, 'from neural_compressor.experimental import Quantization, common\n'), (130, 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', False, 'from neural_compressor.experimental import Quantization, common\n'), (131, 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', False, 'from neural_compressor.experimental import Quantization, common\n'), (68, 'tensorflow.Session', 'tf.Session', True, 'import tensorflow as tf\n'), (69, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (71, 'tensorflow.nn.conv2d', 'tf.nn.conv2d', True, 'import tensorflow as tf\n'), (75, 'tensorflow.graph_util.convert_variables_to_constants', 'tf.graph_util.convert_variables_to_constants', True, 'import tensorflow as tf\n'), (80, 'tensorflow.import_graph_def', 'tf.import_graph_def', True, 'import tensorflow as tf\n'), (82, 'tensorflow.Graph', 'tf.Graph', True, 'import tensorflow as tf\n'), (83, 'tensorflow.compat.v1.GraphDef', 'tf.compat.v1.GraphDef', True, 'import tensorflow as tf\n'), (70, 'numpy.random.random', 'np.random.random', True, 'import numpy as np\n'), (74, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', True, 'import tensorflow as tf\n'), (84, 'tensorflow.compat.v1.Session', 'tf.compat.v1.Session', True, 'import tensorflow as tf\n'), (85, 'tensorflow.compat.v1.placeholder', 'tf.compat.v1.placeholder', True, 'import tensorflow as tf\n'), (87, 'tensorflow.nn.conv2d', 'tf.nn.conv2d', True, 'import tensorflow as tf\n'), (91, 'tensorflow.compat.v1.graph_util.convert_variables_to_constants', 'tf.compat.v1.graph_util.convert_variables_to_constants', True, 'import tensorflow as tf\n'), (96, 'tensorflow.import_graph_def', 'tf.import_graph_def', True, 'import tensorflow as tf\n'), (86, 'numpy.random.random', 'np.random.random', True, 'import numpy as np\n'), (90, 'tensorflow.compat.v1.global_variables_initializer', 'tf.compat.v1.global_variables_initializer', True, 'import tensorflow as tf\n')]