# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# 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 unittest
import os
import time
import numpy as np
from paddle.dataset.common import download, DATA_HOME
from paddle.fluid.contrib.slim.quantization import WeightQuantization
import paddle

paddle.enable_static()


def _load_variable_data(scope, var_name):
    '''
    Load variable value from scope
    '''
    var_node = scope.find_var(var_name)
    assert var_node is not None, \
        "Cannot find " + var_name + " in scope."
    return np.array(var_node.get_tensor())


def _set_variable_data(scope, place, var_name, np_value):
    '''
    Set the value of var node by name, if the node exits,
    '''
    assert isinstance(np_value, np.ndarray), \
        'The type of value should be numpy array.'
    var_node = scope.find_var(var_name)
    if var_node != None:
        tensor = var_node.get_tensor()
        tensor.set(np_value, place)


class TestWeightQuantization(unittest.TestCase):
    def setUp(self):
        self.weight_quantization_dir = 'weight_quantization'
        self.cache_folder = os.path.join(DATA_HOME,
                                         self.weight_quantization_dir)

    def download_model(self, model_name, data_url, data_md5):
        download(data_url, self.weight_quantization_dir, data_md5)
        file_name = data_url.split('/')[-1]
        file_path = os.path.join(self.cache_folder, file_name)
        print(model_name + ' is downloaded at ' + file_path)

        unziped_path = os.path.join(self.cache_folder, model_name)
        self.cache_unzipping(unziped_path, file_path)
        print(model_name + ' is unziped at ' + unziped_path)
        return unziped_path

    def cache_unzipping(self, target_folder, zip_path):
        if not os.path.exists(target_folder):
            cmd = 'mkdir {0} && tar xf {1} -C {0}'.format(target_folder,
                                                          zip_path)
            os.system(cmd)

    def quantize_to_int(self, model_name, model_data_url, model_data_md5,
                        weight_bits, quantizable_op_type, weight_quantize_type,
                        generate_test_model, threshold_rate):

        model_dir = self.download_model(model_name, model_data_url,
                                        model_data_md5)
        load_model_dir = os.path.join(model_dir, model_name)

        timestamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
        save_model_dir = os.path.join(
            os.getcwd(),
            model_name + "_wq_" + str(weight_bits) + "_" + timestamp)

        weight_quant = WeightQuantization(model_dir=load_model_dir)
        weight_quant.quantize_weight_to_int(
            save_model_dir=save_model_dir,
            weight_bits=weight_bits,
            quantizable_op_type=quantizable_op_type,
            weight_quantize_type=weight_quantize_type,
            generate_test_model=generate_test_model,
            threshold_rate=threshold_rate)
        print("finish weight quantization for " + model_name + "\n")

        try:
            os.system("rm -rf {}".format(save_model_dir))
        except Exception as e:
            print("Failed to delete {} due to {}".format(save_model_dir, str(
                e)))

    def convert_to_fp16(self, model_name, model_data_url, model_data_md5,
                        model_filename, params_filename):
        model_dir = self.download_model(model_name, model_data_url,
                                        model_data_md5)
        load_model_dir = os.path.join(model_dir, model_name)

        timestamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
        save_model_dir = os.path.join(os.getcwd(),
                                      model_name + "_wq_fp16_" + timestamp)

        weight_quant = WeightQuantization(load_model_dir, model_filename,
                                          params_filename)

        weight_quant.convert_weight_to_fp16(save_model_dir)

        print("finish converting the data type of weights to fp16 for " +
              model_name)
        print("fp16 model saved in " + save_model_dir + "\n")

        input_data = np.ones([1, 3, 224, 224], dtype=np.float32)
        res_fp32 = self.run_models(load_model_dir, model_filename,
                                   params_filename, input_data, False)
        res_fp16 = self.run_models(save_model_dir, model_filename,
                                   params_filename, input_data, True)

        self.assertTrue(
            np.allclose(
                res_fp32, res_fp16, rtol=1e-5, atol=1e-08, equal_nan=True),
            msg='Failed to test the accuracy of the fp32 and fp16 model.')

        try:
            os.system("rm -rf {}".format(save_model_dir))
        except Exception as e:
            print("Failed to delete {} due to {}".format(save_model_dir, str(
                e)))

    def run_models(self, model_dir, model_filename, params_filename, input_data,
                   is_fp16_model):
        print(model_dir)

        place = paddle.CPUPlace()
        exe = paddle.static.Executor(place)
        scope = paddle.static.Scope()
        with paddle.static.scope_guard(scope):
            [inference_program, feed_target_names, fetch_targets] = \
                paddle.fluid.io.load_inference_model(model_dir, exe,
                    model_filename=model_filename,
                    params_filename=params_filename)

        if is_fp16_model:
            for var in inference_program.list_vars():
                if (var.type == paddle.fluid.core.VarDesc.VarType.RAW) or \
                    (not var.persistable) or (var.name in ['feed', 'fetch']) \
                    or (var.dtype != paddle.fluid.core.VarDesc.VarType.FP16):
                    continue
                tensor = _load_variable_data(scope, var.name)
                _set_variable_data(scope, place, var.name,
                                   tensor.astype(np.float32))

        results = exe.run(inference_program,
                          feed={feed_target_names[0]: input_data},
                          fetch_list=fetch_targets)
        return np.array(results[0])


class TestWeightQuantizationMobilenetv1(TestWeightQuantization):
    nocomb_model_name = "mobilenetv1_fp32_nocombined"
    nocomb_model_data_url = "https://paddle-inference-dist.cdn.bcebos.com/Paddle-Inference-Demo/mobilenetv1_fp32_nocombined.tar.gz"
    nocomb_model_data_md5 = "c9aae3b04d9d535c84590ae557be0a0b"

    comb_model_name = "mobilenetv1_fp32_combined"
    comb_model_data_url = "https://paddle-inference-dist.cdn.bcebos.com/Paddle-Inference-Demo/mobilenetv1_fp32_combined.tar.gz"
    comb_model_data_md5 = "087c67e2b2b0a8b689fcc570a56c005f"

    def test_weight_quantization_mobilenetv1_8bit_abs_max(self):
        weight_bits = 8
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul']
        weight_quantize_type = "abs_max"
        generate_test_model = True
        threshold_rate = 0.0
        self.quantize_to_int(self.nocomb_model_name, self.nocomb_model_data_url,
                             self.nocomb_model_data_md5, weight_bits,
                             quantizable_op_type, weight_quantize_type,
                             generate_test_model, threshold_rate)

    def test_weight_quantization_mobilenetv1_8bit_channel_wise_abs_max(self):
        weight_bits = 8
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul']
        weight_quantize_type = "channel_wise_abs_max"
        generate_test_model = True
        threshold_rate = 0.0
        self.quantize_to_int(self.nocomb_model_name, self.nocomb_model_data_url,
                             self.nocomb_model_data_md5, weight_bits,
                             quantizable_op_type, weight_quantize_type,
                             generate_test_model, threshold_rate)

    def test_weight_quantization_mobilenetv1_16bit_abs_max(self):
        weight_bits = 16
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul']
        weight_quantize_type = "abs_max"
        generate_test_model = False
        threshold_rate = 0
        self.quantize_to_int(self.nocomb_model_name, self.nocomb_model_data_url,
                             self.nocomb_model_data_md5, weight_bits,
                             quantizable_op_type, weight_quantize_type,
                             generate_test_model, threshold_rate)

    def test_weight_quantization_mobilenetv1_16bit_channel_wise_abs_max(self):
        weight_bits = 16
        quantizable_op_type = ['conv2d', 'depthwise_conv2d', 'mul']
        weight_quantize_type = "channel_wise_abs_max"
        generate_test_model = False
        threshold_rate = 1e-9
        self.quantize_to_int(self.nocomb_model_name, self.nocomb_model_data_url,
                             self.nocomb_model_data_md5, weight_bits,
                             quantizable_op_type, weight_quantize_type,
                             generate_test_model, threshold_rate)

    def test_mobilenetv1_fp16_combined(self):
        model_filename = '__model__'
        params_filename = '__params__'
        self.convert_to_fp16(self.comb_model_name, self.comb_model_data_url,
                             self.comb_model_data_md5, model_filename,
                             params_filename)

    def test_mobilenetv1_fp16_nocombined(self):
        model_filename = None
        params_filename = None
        self.convert_to_fp16(self.nocomb_model_name, self.nocomb_model_data_url,
                             self.nocomb_model_data_md5, model_filename,
                             params_filename)


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