# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.  See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# This script benchmarks gpt2 model with past state.
# For gpt2 model without past state, use benchmark.py to measure performance.

import argparse
import csv
import logging
import os
import sys
from datetime import datetime

import psutil
import torch
from gpt2_beamsearch_helper import MODEL_CLASSES, Gpt2HelperFactory
from gpt2_helper import DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS, Gpt2Helper
from packaging import version
from transformers import AutoConfig

sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))

from benchmark_helper import Precision, create_onnxruntime_session, prepare_environment, setup_logger
from quantize_helper import QuantizeHelper

logger = logging.getLogger("")


def parse_arguments(argv=None):
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-m",
        "--model_name_or_path",
        required=True,
        type=str,
        help="Model path, or pretrained model name selected in the list: " + ", ".join(PRETRAINED_GPT2_MODELS),
    )

    parser.add_argument(
        "--model_class",
        required=False,
        type=str,
        default="GPT2LMHeadModel",
        choices=list(MODEL_CLASSES.keys()),
        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
    )

    parser.add_argument(
        "--cache_dir",
        required=False,
        type=str,
        default=os.path.join(".", "cache_models"),
        help="Directory to cache pre-trained models",
    )

    parser.add_argument(
        "--onnx_dir",
        required=False,
        type=str,
        default=os.path.join(".", "onnx_models"),
        help="Directory to store onnx models",
    )

    parser.add_argument(
        "--test_times",
        required=False,
        default=100,
        type=int,
        help="Number of repeat times to get average inference latency.",
    )

    parser.add_argument(
        "-v",
        "--validate_onnx",
        required=False,
        action="store_true",
        help="Validate ONNX model",
    )

    parser.add_argument(
        "-o",
        "--optimize_onnx",
        required=False,
        action="store_true",
        help="Use optimizer.py to optimize onnx model",
    )
    parser.set_defaults(optimize_onnx=False)

    parser.add_argument("--use_gpu", required=False, action="store_true", help="use GPU for inference")
    parser.set_defaults(use_gpu=False)

    parser.add_argument(
        "-p",
        "--precision",
        type=Precision,
        default=Precision.FLOAT32,
        choices=list(Precision),
        help="Precision of model to run. fp32 for full precision, fp16 for half precision, and int8 for quantization",
    )

    parser.add_argument("--torchscript", required=False, action="store_true", help="use Torchscript")
    parser.set_defaults(torchscript=False)

    parser.add_argument("-b", "--batch_sizes", nargs="+", type=int, default=[1], help="batch size")
    parser.add_argument(
        "--beam_size",
        type=int,
        default=4,
        help="Beam size if greedy/top-p/top-k sampling is needed",
    )

    parser.add_argument(
        "--sequence_lengths",
        nargs="+",
        type=int,
        default=[1],
        help="sequence lengths (excluding past)",
    )

    parser.add_argument(
        "-s",
        "--past_sequence_lengths",
        nargs="+",
        type=int,
        default=[8, 16, 32, 64, 128, 256],
        help="past sequence lengths",
    )

    parser.add_argument(
        "-r",
        "--result_csv",
        required=False,
        default=None,
        help="CSV file for saving summary results.",
    )

    parser.add_argument("--thread_num", required=False, type=int, default=-1, help="Threads to use")

    parser.add_argument("--include_copy_output_latency", required=False, action="store_true")
    parser.set_defaults(include_copy_output_latency=False)

    parser.add_argument("--verbose", required=False, action="store_true")
    parser.set_defaults(verbose=False)

    search_option_group = parser.add_argument_group("configurable one step search options")

    search_option_group.add_argument(
        "--ignore_eos",
        type=bool,
        default=False,
        help="If ignore end of sentence token in model inference.",
    )
    search_option_group.add_argument(
        "--repetition_penalty",
        type=float,
        default=1,
        help="Positive. >1 to penalize and <1 to encorage.",
    )
    search_option_group.add_argument(
        "--temperature",
        type=float,
        default=1,
        help="Softmax temperature for output logits.",
    )
    search_option_group.add_argument(
        "--excluded_token_ids",
        required=False,
        nargs="+",
        type=float,
        help="A list of token ids to be excluded in inference.",
    )
    search_option_group.add_argument(
        "--length_penalty",
        type=float,
        default=1,
        help="Positive. >1 to penalize and <1 to encorage short sentence.",
    )

    sampling_option_group = parser.add_argument_group("one step sampling options")
    sampling_option_group.add_argument(
        "--do_sample",
        action="store_true",
        help="If to do sampling instead of beam search or greedy.",
    )
    sampling_option_group.add_argument(
        "--do_sample_top_p",
        type=float,
        default=0.95,
        help="Nuclear/top-p sampling accumulation probability.",
    )
    sampling_option_group.add_argument("--do_sample_top_k", type=int, default=0, help="Use top-k if non-zero.")

    args = parser.parse_args(argv)

    return args


def main(args):
    from transformers import __version__ as transformers_version

    if version.parse(transformers_version) < version.parse(
        "3.1.0"
    ):  # past_key_values name does not exist in 3.0.2 or older
        raise RuntimeError("This tool requires transformers 3.1.0 or later.")

    logger.info(f"Arguments:{args}")
    if args.precision == Precision.FLOAT16:
        assert args.optimize_onnx and args.use_gpu, "fp16 requires --optimize_onnx --use_gpu"

    if args.precision == Precision.INT8:
        assert not args.use_gpu, "quantization only supports CPU"

    torch.set_num_threads(psutil.cpu_count(logical=True) if args.thread_num <= 0 else args.thread_num)
    print(torch.__config__.parallel_info())

    cache_dir = args.cache_dir
    output_dir = args.onnx_dir
    prepare_environment(cache_dir, output_dir, args.use_gpu)

    model_class = MODEL_CLASSES[args.model_class][0]
    if args.model_class == "GPT2LMHeadModel_BeamSearchStep":
        model_type = "beam_search_step"
    elif args.model_class == "GPT2LMHeadModel_ConfigurableOneStepSearch":
        model_type = "configurable_one_step_search"
    else:
        model_type = "default"

    gpt2helper = Gpt2HelperFactory.create_helper(model_type)
    config = AutoConfig.from_pretrained(args.model_name_or_path, torchscript=args.torchscript, cache_dir=cache_dir)
    if model_type == "beam_search_step":
        model = model_class.from_pretrained(
            args.model_name_or_path,
            config=config,
            batch_size=1,
            beam_size=args.beam_size,
            cache_dir=cache_dir,
        )
    elif model_type == "configurable_one_step_search":
        model = model_class.from_pretrained(
            args.model_name_or_path,
            config=config,
            batch_size=1,
            beam_size=args.beam_size,
            ignore_eos=args.ignore_eos,
            temperature=args.temperature,
            repetition_penalty=args.repetition_penalty,
            excluded_token_ids=args.excluded_token_ids,
            length_penalty=args.length_penalty,
            do_sample=args.do_sample,
            do_sample_top_p=args.do_sample_top_p,
            do_sample_top_k=args.do_sample_top_k,
            cache_dir=cache_dir,
        )
    else:
        model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir)

    # This scirpt does not support float16 for PyTorch.
    # if args.float16:
    #    model.half()

    device = torch.device("cuda:0" if args.use_gpu else "cpu")
    model.to(device)
    use_external_data_format = config.n_layer > 24  # TODO: find a way to check model size > 2GB
    onnx_model_paths = gpt2helper.get_onnx_paths(
        output_dir,
        args.model_name_or_path,
        args.model_class,
        has_past=True,
        new_folder=use_external_data_format,
    )

    onnx_model_path = onnx_model_paths["raw"]
    use_padding = MODEL_CLASSES[args.model_class][2]
    gpt2helper.export_onnx(
        model,
        device,
        onnx_model_path,
        args.verbose,
        use_external_data_format,
        has_position_ids=use_padding,
        has_attention_mask=use_padding,
    )

    if args.optimize_onnx or args.precision != Precision.FLOAT32:
        onnx_model_path = onnx_model_paths[str(args.precision) if args.precision != Precision.INT8 else "fp32"]
        gpt2helper.optimize_onnx(
            onnx_model_paths["raw"],
            onnx_model_path,
            args.precision == Precision.FLOAT16,
            model.config.num_attention_heads,
            model.config.hidden_size,
            use_external_data_format,
            auto_mixed_precision=True,
        )

        if args.precision == Precision.INT8:
            logger.info("quantizing model...")
            QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_paths["int8"], use_external_data_format)
            model = QuantizeHelper.quantize_torch_model(model)
            logger.info("finished quantizing model")
            onnx_model_path = onnx_model_paths["int8"]

    if args.torchscript:
        model = gpt2helper.torchscript(
            model,
            config,
            device,
            has_position_ids=use_padding,
            has_attention_mask=use_padding,
        )

    session = create_onnxruntime_session(
        onnx_model_path,
        args.use_gpu,
        enable_all_optimization=False,
        num_threads=args.thread_num,
        verbose=args.verbose,
    )
    if session is None:
        return

    # Allocate output buffers for IO Binding
    if model_type == "beam_search_step" or model_type == "configurable_one_step_search":
        max_output_shapes = gpt2helper.get_output_shapes(
            max(args.batch_sizes),
            context_len=max(args.past_sequence_lengths),
            past_sequence_length=max(args.past_sequence_lengths),
            sequence_length=max(args.sequence_lengths),
            beam_size=args.beam_size,
            step=0,
            config=config,
            model_class=args.model_class,
        )

        output_buffers = gpt2helper.get_output_buffers(max_output_shapes, device, args.precision == Precision.FLOAT16)

    else:
        max_output_shapes = gpt2helper.get_output_shapes(
            max(args.batch_sizes),
            max(args.past_sequence_lengths),
            max(args.sequence_lengths),
            config,
            args.model_class,
        )
        output_buffers = gpt2helper.get_output_buffers(max_output_shapes, device, args.precision == Precision.FLOAT16)

    csv_filename = args.result_csv or "benchmark_result_{}.csv".format(datetime.now().strftime("%Y%m%d-%H%M%S"))
    with open(csv_filename, mode="a", newline="") as csv_file:
        column_names = [
            "model_name",
            "model_class",
            "gpu",
            "precision",
            "optimizer",
            "torchscript",
            "batch_size",
            "sequence_length",
            "past_sequence_length",
            "torch_latency",
            "onnxruntime_latency",
            "onnxruntime_io_binding_latency",
        ]
        csv_writer = csv.DictWriter(csv_file, fieldnames=column_names)
        csv_writer.writeheader()

        for batch_size in args.batch_sizes:
            for sequence_length in args.sequence_lengths:
                for past_sequence_length in args.past_sequence_lengths:
                    assert batch_size > 0 and sequence_length > 0 and past_sequence_length >= 0
                    logger.debug(
                        f"Running test for batch_size={batch_size} sequence_length={sequence_length} past_sequence_length={past_sequence_length}..."
                    )
                    if model_type == "beam_search_step" or model_type == "configurable_one_step_search":
                        dummy_inputs = gpt2helper.get_dummy_inputs(
                            batch_size,
                            past_sequence_length,
                            sequence_length,
                            config.num_attention_heads,
                            config.hidden_size,
                            config.n_layer,
                            config.vocab_size,
                            device,
                            float16=(args.precision == Precision.FLOAT16),
                            has_position_ids=use_padding,
                            has_attention_mask=use_padding,
                        )
                        output_shapes = gpt2helper.get_output_shapes(
                            batch_size,
                            past_sequence_length,
                            past_sequence_length,
                            sequence_length,
                            args.beam_size,
                            0,
                            config,
                            args.model_class,
                        )
                    else:
                        dummy_inputs = gpt2helper.get_dummy_inputs(
                            batch_size,
                            past_sequence_length,
                            sequence_length,
                            config.num_attention_heads,
                            config.hidden_size,
                            config.n_layer,
                            config.vocab_size,
                            device,
                            float16=(args.precision == Precision.FLOAT16),
                            has_position_ids=use_padding,
                            has_attention_mask=use_padding,
                        )
                        output_shapes = gpt2helper.get_output_shapes(
                            batch_size,
                            past_sequence_length,
                            sequence_length,
                            config,
                            args.model_class,
                        )

                    try:
                        outputs, torch_latency = gpt2helper.pytorch_inference(model, dummy_inputs, args.test_times)

                        # Dump Torch output shape
                        for i, value in enumerate(outputs):
                            if isinstance(value, tuple):
                                logger.debug(f"torch output {i} is tuple of size {len(value)}, shape {value[0].shape}")
                            else:
                                logger.debug(f"torch output {i} shape {value.shape}")

                        ort_outputs, ort_latency = gpt2helper.onnxruntime_inference(
                            session, dummy_inputs, args.test_times
                        )

                        (ort_io_outputs, ort_io_latency,) = gpt2helper.onnxruntime_inference_with_binded_io(
                            session,
                            dummy_inputs,
                            output_buffers,
                            output_shapes,
                            args.test_times,
                            return_numpy=False,
                            include_copy_output_latency=args.include_copy_output_latency,
                        )

                        if args.validate_onnx:
                            if gpt2helper.compare_outputs(
                                outputs,
                                ort_outputs,
                                model_class=args.model_class,
                                rtol=DEFAULT_TOLERANCE[args.precision],
                                atol=DEFAULT_TOLERANCE[args.precision],
                            ):
                                logger.info(
                                    f"Pytorch and ONNX Runtime outputs are all close (tolerance={DEFAULT_TOLERANCE[args.precision]})."
                                )

                            # Results of IO binding might be in GPU. Copy outputs to CPU for comparison.
                            copy_outputs = []
                            for output in ort_io_outputs:
                                copy_outputs.append(output.cpu().numpy())

                            if gpt2helper.compare_outputs(
                                outputs,
                                copy_outputs,
                                model_class=args.model_class,
                                rtol=DEFAULT_TOLERANCE[args.precision],
                                atol=DEFAULT_TOLERANCE[args.precision],
                            ):
                                logger.info(
                                    f"Pytorch and ONNX Runtime IO Binding outputs are all close (tolerance={DEFAULT_TOLERANCE[args.precision]})."
                                )

                        logger.info(
                            f"batch_size={batch_size}, sequence_length={sequence_length}, past_sequence_length={past_sequence_length}, torch_latency={torch_latency:.2f}, onnxruntime_latency={ort_latency:.2f}, onnxruntime_io_binding_latency={ort_io_latency:.2f}"
                        )

                        row = {
                            "model_name": args.model_name_or_path,
                            "model_class": args.model_class,
                            "gpu": args.use_gpu,
                            "precision": args.precision,
                            "optimizer": args.optimize_onnx,
                            "torchscript": args.torchscript,
                            "batch_size": batch_size,
                            "sequence_length": sequence_length,
                            "past_sequence_length": past_sequence_length,
                            "torch_latency": f"{torch_latency:.2f}",
                            "onnxruntime_latency": f"{ort_latency:.2f}",
                            "onnxruntime_io_binding_latency": f"{ort_io_latency:.2f}",
                        }
                        csv_writer.writerow(row)
                    except:
                        logger.error(f"Exception", exc_info=True)
                        return None

    logger.info(f"Results are saved to file {csv_filename}")
    return csv_filename


if __name__ == "__main__":
    args = parse_arguments()
    setup_logger(args.verbose)
    main(args)
