# 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.

import argparse
import logging
import os
import re
import sys
from inspect import currentframe, getframeinfo

import mxnet as mx
from mxnet import gluon
from mxnet.contrib.quantization import quantize_net
from mxnet.gluon.data import DataLoader
from mxnet.gluon.data.vision import transforms
from mxnet.gluon.model_zoo.vision import get_model

SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, '..', '..')))
from tools.rec2idx import IndexCreator


def download_calib_dataset(dataset_url, calib_dataset, logger=None):
    if logger is not None:
        logger.info(f'Downloading calibration dataset from {dataset_url} to {calib_dataset}')
    mx.test_utils.download(dataset_url, calib_dataset)


def get_from_gluon(model_name, classes=1000, logger=None):
    dir_path = os.path.dirname(os.path.realpath(__file__))
    model_path = os.path.join(dir_path, 'model')
    if logger is not None:
        logger.info(f'Converting model from Gluon-CV ModelZoo {model_name}... into path {model_path}')
    net = get_model(name=model_name, classes=classes, pretrained=True)
    prefix = os.path.join(model_path, model_name)
    return net, prefix


def regex_find_excluded_symbols(patterns_dict, model_name):
    for key, value in patterns_dict.items():
        if re.search(key, model_name) is not None:
            return value
    return None


def get_exclude_symbols(model_name, exclude_first_conv):
    """Grouped supported models at the time of commit:
    - alexnet
    - densenet121, densenet161
    - densenet169, densenet201
    - inceptionv3
    - mobilenet0.25, mobilenet0.5, mobilenet0.75, mobilenet1.0,
    - mobilenetv2_0.25, mobilenetv2_0.5, mobilenetv2_0.75, mobilenetv2_1.0
    - resnet101_v1, resnet152_v1, resnet18_v1, resnet34_v1, resnet50_v1
    - resnet101_v2, resnet152_v2, resnet18_v2, resnet34_v2, resnet50_v2
    - squeezenet1.0, squeezenet1.1
    - vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn
    """
    exclude_symbol_regex = {
        'mobilenet[^v]': ['mobilenet_hybridsequential0_flatten0_flatten0', 'mobilenet_hybridsequential0_globalavgpool2d0_fwd'],
        'mobilenetv2': ['mobilenetv2_hybridsequential1_flatten0_flatten0'],
        # resnetv2_hybridsequential0_hybridsequential0_bottleneckv20_batchnorm0_fwd is excluded for the sake of accuracy
        'resnet.*v2': ['resnetv2_hybridsequential0_flatten0_flatten0', 'resnetv2_hybridsequential0_hybridsequential0_bottleneckv20_batchnorm0_fwd'],
        'squeezenet1': ['squeezenet_hybridsequential1_flatten0_flatten0'],
    }
    excluded_sym_names = regex_find_excluded_symbols(exclude_symbol_regex, model_name)
    if excluded_sym_names is None:
        excluded_sym_names = []
    if exclude_first_conv:
        first_conv_regex = {
            'alexnet': ['alexnet_hybridsequential0_conv2d0_fwd'],
            'densenet': ['densenet_hybridsequential0_conv2d0_fwd'],
            'inceptionv3': ['inception3_hybridsequential0_hybridsequential0_conv2d0_fwd'],
            'mobilenet[^v]': ['mobilenet_hybridsequential0_conv2d0_fwd'],
            'mobilenetv2': ['mobilenetv2_hybridsequential0_conv2d0_fwd'],
            'resnet.*v1': ['resnetv1_hybridsequential0_conv2d0_fwd'],
            'resnet.*v2': ['resnetv2_hybridsequential0_conv2d0_fwd'],
            'squeezenet1': ['squeezenet_hybridsequential0_conv2d0_fwd'],
            'vgg': ['vgg_hybridsequential0_conv2d0_fwd'],
        }
        excluded_first_conv_sym_names = regex_find_excluded_symbols(first_conv_regex, model_name)
        if excluded_first_conv_sym_names is None:
            raise ValueError(f'Currently, model {model_name} is not supported in this script')
        excluded_sym_names += excluded_first_conv_sym_names
    return excluded_sym_names


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate a calibrated quantized model from a FP32 model with oneDNN support')
    parser.add_argument('--model', type=str, default='resnet50_v1',
                        help='model to be quantized. If no-pretrained is set then'
                             'model must be provided to `model` directory in the same path'
                             'as this python script')
    parser.add_argument('--epoch', type=int, default=0,
                        help='number of epochs, default is 0')
    parser.add_argument('--no-pretrained', action='store_true', default=False,
                        help='If enabled, will not download pretrained model from MXNet or Gluon-CV modelzoo.')
    parser.add_argument('--batch-size', type=int, default=32)
    parser.add_argument('--calib-dataset', type=str, default='data/val_256_q90.rec',
                        help='path of the calibration dataset')
    parser.add_argument('--image-shape', type=str, default='3,224,224',
                        help='number of channels, height and width of input image separated by comma')
    parser.add_argument('--data-nthreads', type=int, default=0,
                        help='number of threads for data loading')
    parser.add_argument('--num-calib-batches', type=int, default=10,
                        help='number of batches for calibration')
    parser.add_argument('--exclude-first-conv', action='store_true', default=False,
                        help='excluding quantizing the first conv layer since the'
                             ' input data may have negative value which doesn\'t support at moment')
    parser.add_argument('--shuffle-dataset', action='store_true',
                        help='shuffle the calibration dataset')
    parser.add_argument('--calib-mode', type=str, default='entropy',
                        help='calibration mode used for generating calibration table for the quantized symbol; supports'
                             ' 1. none: no calibration will be used. The thresholds for quantization will be calculated'
                             ' on the fly. This will result in inference speed slowdown and loss of accuracy'
                             ' in general.'
                             ' 2. naive: simply take min and max values of layer outputs as thresholds for'
                             ' quantization. In general, the inference accuracy worsens with more examples used in'
                             ' calibration. It is recommended to use `entropy` mode as it produces more accurate'
                             ' inference results.'
                             ' 3. entropy: calculate KL divergence of the fp32 output and quantized output for optimal'
                             ' thresholds. This mode is expected to produce the best inference accuracy of all three'
                             ' kinds of calibration modes if the calibration dataset is representative enough of the'
                             ' inference dataset.')
    parser.add_argument('--quantized-dtype', type=str, default='auto',
                        choices=['auto', 'int8', 'uint8'],
                        help='quantization destination data type for input data')
    parser.add_argument('--quiet', action='store_true', default=False,
                        help='suppress most of log')
    args = parser.parse_args()
    ctx = mx.cpu(0)
    logger = None

    if not args.quiet:
        logging.basicConfig()
        logger = logging.getLogger('logger')
        logger.setLevel(logging.INFO)

    if logger:
        logger.info(args)
        logger.info(f'shuffle_dataset={args.shuffle_dataset}')
        logger.info(f'calibration mode set to {args.calib_mode}')

    calib_mode = args.calib_mode

    # download calibration dataset
    if calib_mode != 'none':
        idx_file_name = os.path.splitext(args.calib_dataset)[0] + '.idx'
        if not os.path.isfile(idx_file_name):
            download_calib_dataset('http://data.mxnet.io/data/val_256_q90.rec', args.calib_dataset)
            creator = IndexCreator(args.calib_dataset, idx_file_name)
            creator.create_index()
            creator.close()

    # get image shape
    image_shape = args.image_shape
    data_shape = [(1,) + tuple(int(i) for i in image_shape.split(','))]

    # check if directory for output model exists
    dir_path = os.path.dirname(os.path.realpath(__file__))
    dir_path = os.path.join(dir_path, 'model')
    if not os.path.exists(dir_path):
        os.mkdir(dir_path)  # without try catch block as we expect to finish script if it fail

    # download model
    if not args.no_pretrained:
        if logger:
            logger.info('Get pre-trained model from Gluon-CV modelzoo.')
            logger.info('If you want to use custom model, please set --no-pretrained.')
        net, prefix = get_from_gluon(model_name=args.model, classes=1000, logger=logger)
        rgb_mean = '0.485,0.456,0.406'
        rgb_std = '0.229,0.224,0.225'
        epoch = 0
        net.hybridize()
        net(mx.np.zeros(data_shape[0])) # dummy forward pass to build graph
        net.export(prefix) # save model
        net.hybridize(active=False) # disable hybridization - it will be handled in quantization API
    else:
        prefix = os.path.join(dir_path, args.model)
        epoch = args.epoch
        net = gluon.SymbolBlock.imports("{}-symbol.json".format(prefix), ['data'], "{}-0000.params".format(prefix))

    # get batch size
    batch_size = args.batch_size
    if logger:
        logger.info(f'batch size = {batch_size} for calibration')

    # get number of batches for calibration
    num_calib_batches = args.num_calib_batches
    if logger:
        if calib_mode == 'none':
            logger.info('skip calibration step as calib_mode is none')
        else:
            logger.info(f'number of batches = {num_calib_batches} for calibration')

    # get number of threads for decoding the dataset
    data_nthreads = args.data_nthreads

    exclude_first_conv = args.exclude_first_conv
    if args.quantized_dtype == "uint8":
        if logger:
            logger.info('quantized dtype is set to uint8, will exclude first conv.')
        exclude_first_conv = True
    excluded_sym_names = []
    if not args.no_pretrained:
        excluded_sym_names += get_exclude_symbols(args.model, args.exclude_first_conv)
    else:
        if logger:
            frameinfo = getframeinfo(currentframe())
            logger.info(F'Please set proper RGB configs inside this script below {frameinfo.filename}:{frameinfo.lineno} for model {args.model}!')
        # add rgb mean/std of your model.
        rgb_mean = '0,0,0'
        rgb_std = '0,0,0'
        # add layer names you donnot want to quantize.
        if logger:
            frameinfo = getframeinfo(currentframe())
            logger.info(F'Please set proper excluded_sym_names inside this script below {frameinfo.filename}:{frameinfo.lineno} for model {args.model} if required!')
        excluded_sym_names += []
        if exclude_first_conv:
            excluded_sym_names += []

    if logger:
        logger.info(f'These layers have been excluded {excluded_sym_names}')
        logger.info(f'Input data shape = {str(data_shape)}')
        logger.info(f'rgb_mean = {rgb_mean}')
        logger.info(f'rgb_std = {rgb_std}')

    rgb_mean = [float(i) for i in rgb_mean.split(',')]
    mean_args = {'mean_r': rgb_mean[0], 'mean_g': rgb_mean[1], 'mean_b': rgb_mean[2]}
    rgb_std = [float(i) for i in rgb_std.split(',')]
    std_args = {'std_r': rgb_std[0], 'std_g': rgb_std[1], 'std_b': rgb_std[2]}
    if calib_mode == 'none':
        if logger:
            logger.info(f'Quantizing FP32 model {args.model}')
        qsym = quantize_net(net, ctx=ctx, exclude_layers_match=excluded_sym_names, data_shapes=data_shape,
                            calib_mode=calib_mode, quantized_dtype=args.quantized_dtype,
                            logger=logger)
        suffix = '-quantized'
    else:
        if logger:
            logger.info('Creating DataLoader for reading calibration dataset')
        dataset = mx.gluon.data.vision.ImageRecordDataset(args.calib_dataset)
        transformer = transforms.Compose([transforms.Resize(256),
                                          transforms.CenterCrop(224),
                                          transforms.ToTensor(),
                                          transforms.Normalize(mean=rgb_mean, std=rgb_std)])
        data_loader = DataLoader(dataset.transform_first(transformer), batch_size, shuffle=args.shuffle_dataset, num_workers=data_nthreads)
        qsym = quantize_net(net, ctx=ctx, exclude_layers_match=excluded_sym_names,
                            calib_mode=calib_mode, calib_data=data_loader, num_calib_batches=num_calib_batches,
                            quantized_dtype=args.quantized_dtype, logger=logger)
        if calib_mode == 'entropy':
            suffix = f'-quantized-{num_calib_batches}batches-entropy'
        elif calib_mode == 'naive':
            suffix = f'-quantized-{num_calib_batches}batches-naive'
        else:
            raise ValueError(f'unknown calibration mode {calib_mode} received, only supports `none`, `naive`, and `entropy`')
    save_path = prefix + suffix
    model_path, params_path = qsym.export(save_path, epoch)
    if logger is not None:
        logger.info(F'Saved quantized model into:\n{model_path}\n{params_path}')