# Copyright (C) 2020-2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import json
import os
from addict import Dict

from openvino.tools.pot import load_model, save_model, create_pipeline
from openvino.tools.pot.utils.logger import init_logger
from openvino.tools.pot.api.samples.utils.argument_parser import get_common_argparser
from openvino.tools.pot.api.samples.speech.data_loader import ArkDataLoader
from openvino.tools.pot.api.samples.speech.utils import ArkEngine


def parse_args():
    parser = get_common_argparser()
    parser.add_argument(
        '-i',
        '--input_names',
        help='List of input names of network',
        required=True
    )
    parser.add_argument(
        '-f',
        '--files_for_input',
        help='List of filenames mapped to input names (without .ark extension)',
        required=True
    )
    parser.add_argument(
        '-is',
        '--input_shapes',
        help='List of input shapes mapped to input names of network. Input data will be reshaped with provided shape. '
             'Use this argument only for 3D shape inputs. '
             'List example: [[1,2,3],[4,5,6]]',
        type=json.loads,
        default=[]
    )
    parser.add_argument(
        '-p',
        '--preset',
        help='Preset for quantization. '
             '-performance for INT8 weights and INT16 inputs; '
             '-accuracy for INT16 weights and inputs',
        default='accuracy',
        choices=['performance', 'accuracy'])
    parser.add_argument(
        '-o',
        '--output',
        help='Path to save the quantized model',
        default='./model/optimized')
    parser.add_argument(
        '--log-level',
        type=str,
        default='INFO',
        choices=['CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG'],
        help='Log level to print')
    parser.add_argument(
        '-s',
        '--subset_size',
        help='Subset size for calibration',
        default=2000,
        type=int)
    return parser.parse_args()


def get_configs(args):
    if not args.weights:
        args.weights = '{}.bin'.format(os.path.splitext(args.model)[0])

    model_config = Dict({
        'model_name': 'gna_model',
        'model': os.path.expanduser(args.model),
        'weights': os.path.expanduser(args.weights),
        'exec_log_dir': args.output
    })
    engine_config = Dict({
        'device': 'CPU'
    })
    dataset_config = {
        'data_source': os.path.expanduser(args.dataset),
        'type': 'simplified',
        'input_names': args.input_names.split(','),
        'input_shapes': args.input_shapes
    }

    if args.files_for_input is not None:
        dataset_config['input_files'] = args.files_for_input.split(',')

    algorithms = [
        {
            'name': 'DefaultQuantization',
            'params': {
                'target_device': 'GNA',
                'preset': args.preset,
                # The custom configuration is for speech recognition models
                'stat_subset_size': args.subset_size,
                'activations': {
                    'range_estimator': {
                        'max': {
                            'type': 'abs_max',
                            'aggregator': 'max'
                        }
                    }
                }
            }
        }
    ]

    return model_config, engine_config, dataset_config, algorithms


def optimize_model(args):
    model_config, engine_config, dataset_config, algorithms = get_configs(args)
    data_loader = ArkDataLoader(dataset_config)
    engine = ArkEngine(config=engine_config, data_loader=data_loader)
    pipeline = create_pipeline(algorithms, engine)

    model = load_model(model_config, target_device='GNA')
    return pipeline.run(model)


def main():
    args = parse_args()
    out_dir = os.path.expanduser(args.output)
    if not os.path.isdir(out_dir):
        os.makedirs(out_dir)
    init_logger(level=args.log_level, file_name=os.path.join(out_dir, 'log.txt'))
    compressed_model = optimize_model(args)
    save_model(compressed_model, out_dir)


if __name__ == '__main__':
    main()
