# Copyright (C) 2021-2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from __future__ import print_function

import argparse
import logging
import os

import torch
import yaml

from wenet.utils.init_model import init_model
import intel_extension_for_pytorch as ipex
from intel_extension_for_pytorch.quantization import prepare, convert


def get_args():
    parser = argparse.ArgumentParser(description='export your script model')
    parser.add_argument('--config', required=True, help='config file')
    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
    parser.add_argument('--output_file', default=None, help='output file')
    parser.add_argument('--dtype',
                        default="fp32",
                        help='choose the dtype to run:[fp32,bf16]')
    parser.add_argument('--output_quant_file',
                        default=None,
                        help='output quantized model file')
    args = parser.parse_args()
    return args


def scripting(model):
    with torch.inference_mode():
        script_model = torch.jit.script(model)
        script_model = torch.jit.freeze(
            script_model,
            preserved_attrs=[
                "forward_encoder_chunk", "ctc_activation",
                "forward_attention_decoder", "subsampling_rate",
                "right_context", "sos_symbol", "eos_symbol",
                "is_bidirectional_decoder"
            ])
    return script_model


def main():
    args = get_args()
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(levelname)s %(message)s')
    # No need gpu for model export
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

    with open(args.config, 'r') as fin:
        configs = yaml.load(fin, Loader=yaml.FullLoader)
    model, configs = init_model(args, configs)
    print(model)

    # Apply IPEX optimization
    model.eval()
    torch._C._jit_set_texpr_fuser_enabled(False)
    model.to(memory_format=torch.channels_last)
    if args.dtype == "fp32":
        ipex_model = ipex.optimize(model)
    elif args.dtype == "bf16":  # For Intel 4th generation Xeon (SPR)
        ipex_model = ipex.optimize(model,
                                   dtype=torch.bfloat16,
                                   weights_prepack=False)

    # Export jit torch script model
    if args.output_file:
        if args.dtype == "fp32":
            script_model = scripting(ipex_model)
        elif args.dtype == "bf16":
            torch._C._jit_set_autocast_mode(True)
            with torch.cpu.amp.autocast():
                script_model = scripting(ipex_model)
        script_model.save(args.output_file)
        print('Export model successfully, see {}'.format(args.output_file))

    # Export quantized jit torch script model
    if args.output_quant_file:
        dynamic_qconfig = ipex.quantization.default_dynamic_qconfig
        dummy_data = (torch.zeros(1, 67, 80), 16, -16,
                      torch.zeros(12, 4, 32, 128), torch.zeros(12, 1, 256, 7))
        model = prepare(model, dynamic_qconfig, dummy_data)
        model = convert(model)
        script_quant_model = scripting(model)
        script_quant_model.save(args.output_quant_file)
        print('Export quantized model successfully, '
              'see {}'.format(args.output_quant_file))


if __name__ == '__main__':
    main()