# SPDX-License-Identifier: Apache-2.0


"""Tool to convert and test pre-trained tensorflow models."""

# pylint: disable=broad-except,logging-not-lazy,unused-argument,unnecessary-lambda,import-outside-toplevel
# pylint: disable=wrong-import-position,too-many-nested-blocks

import argparse
import os
import random
import re
import shutil
import sys
import tarfile
import tempfile
import time
import zipfile
from collections import namedtuple

import numpy as np
import PIL.Image
import yaml
from packaging.version import Version

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf

# contrib ops are registered only when the module is imported, the following import statement is needed,
# otherwise tf runtime error will show up when the tf model is restored from pb file because of un-registered ops.
try:
    import tensorflow.contrib.rnn  # pylint: disable=unused-import
except:  # pylint: disable=bare-except
    # not needed for tf-2.0
    pass

try:
    pass  # pylint: disable=unused-import
except Exception as err:
    pass

from tfjs_runner import run_tfjs

from tf2onnx import constants, logging, optimizer, tf_loader, tf_utils, utils
from tf2onnx.graph import ExternalTensorStorage
from tf2onnx.tf_loader import tf_reset_default_graph, tf_session
from tf2onnx.tfonnx import process_tf_graph

logger = logging.getLogger("run_pretrained")

TEMP_DIR = os.path.join(utils.get_temp_directory(), "run_pretrained")
PERF_STEP = 10
PERF_TIME = 10

def get_img(shape, path, dtype, should_scale=True):
    """Get image as input."""
    resize_to = shape[1:3]
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), path)
    img = PIL.Image.open(path)
    img = img.resize(resize_to, PIL.Image.LANCZOS)
    img_np = np.array(img).astype(dtype)
    img_np = np.stack([img_np] * shape[0], axis=0).reshape(shape)
    if should_scale:
        img_np = img_np / 255
    return img_np


def get_beach(shape):
    """Get beach image as input."""
    return get_img(shape, "beach.jpg", np.float32, should_scale=True)


def get_beach_uint8(shape):
    """Get beach image as uint8."""
    return get_img(shape, "ade20k.jpg", np.uint8, should_scale=False)


def get_car(shape):
    """Get car image as input."""
    return get_img(shape, "car.JPEG", np.float32, should_scale=True)


def get_ade20k(shape):
    """Get truck image from ade20k segmentation dataset."""
    return get_img(shape, "ade20k.jpg", np.float32, should_scale=True)


def get_ade20k_uint8(shape):
    """Get truck image from ade20k segmentation dataset."""
    return get_img(shape, "ade20k.jpg", np.uint8, should_scale=False)


def get_random(shape):
    """Get random input."""
    np.random.seed(42)
    return np.random.sample(shape).astype(np.float32)


def get_random256(shape):
    """Get random imput between 0 and 255."""
    np.random.seed(42)
    return np.round(np.random.sample(shape) * 256).astype(np.float32)


def get_ramp(shape):
    """Get ramp input."""
    size = np.prod(shape)
    return np.linspace(1, size, size).reshape(shape).astype(np.float32)


def get_ones(shape):
    """Get ones."""
    return np.ones(shape).astype(np.float32)

def get_zeros(shape):
    """Get zeros."""
    return np.zeros(shape).astype(np.float32)

def get_zeros_int32(shape):
    """Get zeros."""
    return np.zeros(shape).astype(np.int32)

def get_zeros_int64(shape):
    """Get zeros."""
    return np.zeros(shape).astype(np.int64)

def get_ones_int32(shape):
    """Get ones."""
    return np.ones(shape).astype(np.int32)

def get_small_rand_int32(shape):
    """Get random ints in range [1, 99]"""
    np.random.seed(42)
    return np.random.randint(low=1, high=100, size=shape, dtype=np.int32)

def get_zeros_then_ones(shape):
    """Fill half the tensor with zeros and the rest with ones"""
    cnt = np.prod(shape)
    zeros_cnt = cnt // 2
    ones_cnt = cnt - zeros_cnt
    return np.concatenate((np.zeros(zeros_cnt, dtype=np.int32), np.ones(ones_cnt, dtype=np.int32))).reshape(shape)

def get_wav(shape):
    """Get sound data."""
    return np.sin(np.linspace(-np.pi, np.pi, shape[0]), dtype=np.float32)

def get_sentences(shape):
    """Get sentences of shape"""
    words = "the quick brown fox jumps over a lazy dog".split(' ')
    random.seed(42)
    def get_sentence():
        length = random.randint(2, 7)
        return ' '.join(random.choice(words) for _ in range(length))
    return np.array([get_sentence() for _ in range(np.product(shape))]).reshape(shape)


_INPUT_FUNC_MAPPING = {
    "get_beach": get_beach,
    "get_beach_uint8": get_beach_uint8,
    "get_car": get_car,
    "get_ade20k": get_ade20k,
    "get_ade20k_uint8": get_ade20k_uint8,
    "get_random": get_random,
    "get_random256": get_random256,
    "get_ramp": get_ramp,
    "get_ones": get_ones,
    "get_zeros": get_zeros,
    "get_wav": get_wav,
    "get_zeros_int32": get_zeros_int32,
    "get_zeros_int64": get_zeros_int64,
    "get_ones_int32": get_ones_int32,
    "get_small_rand_int32": get_small_rand_int32,
    "get_zeros_then_ones": get_zeros_then_ones,
    "get_sentences": get_sentences,
}


OpsetConstraint = namedtuple("OpsetConstraint", "domain, min_version, max_version, excluded_version")


class Test(object):
    """Main Test class."""

    cache_dir = None
    target = []

    def __init__(self, url, local, input_func, input_names, output_names,
                 disabled=False, rtol=0.01, atol=1e-6, ptol=0, dequantize=False,
                 check_only_shape=False, model_type="frozen", force_input_shape=False,
                 skip_tensorflow=False, opset_constraints=None, tf_min_version=None, tag=None,
                 skip_conversion=False, converted_model=None, signature_def=None, concrete_function=None,
                 large_model=False, structured_outputs=None, run_tf_frozen=None, use_custom_ops=False,
                 ort_profile=None, tf_profile=None):
        self.url = url
        self.input_func = input_func
        self.local = local
        self.input_names = input_names
        self.output_names = output_names
        self.disabled = disabled
        self.large_model = large_model
        self.ort_profile = ort_profile
        self.tf_profile = tf_profile
        self.use_custom_ops = use_custom_ops
        if run_tf_frozen is None:
            run_tf_frozen = not self.large_model
        self.run_tf_frozen = run_tf_frozen
        self.structured_outputs = structured_outputs  # Needed to determine output order for tf_function
        self.rtol = rtol
        self.atol = atol
        self.ptol = ptol
        self.dequantize = dequantize
        self.check_only_shape = check_only_shape
        self.perf = None
        self.tf_runtime = 0
        self.onnx_runtime = 0
        self.model_type = model_type
        self.tag = tag
        self.force_input_shape = force_input_shape
        self.skip_tensorflow = skip_tensorflow
        self.skip_conversion = skip_conversion
        self.converted_model = converted_model
        self.opset_constraints = opset_constraints
        self.tf_min_version = tf_min_version
        self.signatures = [signature_def] if signature_def else None
        self.concrete_function = concrete_function

    def make_input(self, v):
        """Allows each input to specify its own function while defaulting to the input_get function"""
        if isinstance(v, dict):
            if "input_get" in v:
                return _INPUT_FUNC_MAPPING[v["input_get"]](v["shape"])
            if "value" in v:
                return np.array(v["value"])
        return self.input_func(v)

    def download_model(self):
        """Download model from url."""
        cache_dir = Test.cache_dir
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
        url = self.url
        if url.startswith(r'module://'):
            return self.download_from_module()
        k = url.rfind('/')
        fname = self.url[k + 1:]
        dir_name = fname + "_dir"
        ftype = None
        if url.endswith(".tar.gz") or url.endswith(".tgz"):
            ftype = 'tgz'
            dir_name = fname.replace(".tar.gz", "").replace(".tgz", "")
        elif url.endswith('.zip'):
            ftype = 'zip'
            dir_name = fname.replace(".zip", "")
        elif url.endswith('.tflite'):
            ftype = 'tflite'
            dir_name = fname.replace(".tflite", "")
        elif self.model_type == 'tflite':
            fname = self.local
            dir_name = fname.replace(".tflite", "") + "_dir"
        elif self.model_type == 'tfjs':
            ftype = 'tgz'
            fname = 'model.tar.gz'
            dir_name = "_".join(url.split("/")[5:-3]) + "_dir"
        dir_name = os.path.join(cache_dir, dir_name)
        os.makedirs(dir_name, exist_ok=True)
        fpath = os.path.join(dir_name, fname)
        if not os.path.exists(fpath):
            utils.get_url(url, fpath)
        model_path = os.path.join(dir_name, self.local)
        if not os.path.exists(model_path) or self.local == ".":
            if ftype == 'tgz':
                tar = tarfile.open(fpath)
                tar.extractall(dir_name)
                tar.close()
            elif ftype == 'zip':
                zip_ref = zipfile.ZipFile(fpath, 'r')
                zip_ref.extractall(dir_name)
                zip_ref.close()
        return fpath, dir_name

    def download_from_module(self):
        """Download a model from a python module"""
        cache_dir = Test.cache_dir
        from importlib import import_module
        i = self.url.rfind('//')
        module, model_name = self.url[i + 2:].split('/')
        mod_object = import_module(module)
        model_class = getattr(mod_object, model_name)
        model = model_class()
        fpath = os.path.join(cache_dir, self.local)
        model.save(fpath)
        return fpath, cache_dir

    def run_tensorflow(self, sess, inputs):
        """Run model on tensorflow so we have a reference output."""
        feed_dict = {}
        for k, v in inputs.items():
            k = sess.graph.get_tensor_by_name(k)
            feed_dict[k] = v
        logger.info("Running TF")
        result = sess.run(self.output_names, feed_dict=feed_dict)
        if self.perf:
            logger.info("Running TF perf")
            n = 0
            start = time.time()
            stop = start + PERF_TIME
            while time.time() < stop:
                for _ in range(PERF_STEP):
                    _ = sess.run(self.output_names, feed_dict=feed_dict)
                n += PERF_STEP
            self.tf_runtime = 1000 * (time.time() - start) / n
            logger.info("TF perf {:.2f}ms/inference, n={}".format(self.tf_runtime, n))
        return result

    def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, input_names=None,
                const_node_values=None, initialized_tables=None, tflite_path=None, tensors_to_rename=None,
                tfjs_path=None):
        """Convert graph to tensorflow."""
        if extra_opset is None:
            extra_opset = []
        if self.use_custom_ops:
            extra_opset.append(utils.make_opsetid(constants.CONTRIB_OPS_DOMAIN, 1))
        return process_tf_graph(tf_graph, continue_on_error=False, opset=opset,
                                extra_opset=extra_opset, target=Test.target, shape_override=shape_override,
                                input_names=input_names, output_names=self.output_names,
                                const_node_values=const_node_values, initialized_tables=initialized_tables,
                                tflite_path=tflite_path, dequantize=self.dequantize,
                                tensors_to_rename=tensors_to_rename, tfjs_path=tfjs_path)

    def run_onnxruntime(self, name, model_proto, inputs, outputs, external_tensor_storage=None):
        """Run test against onnxruntime backend."""
        import onnxruntime as rt
        model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=True,
                                           as_text=utils.is_debug_mode(),
                                           external_tensor_storage=external_tensor_storage)
        logger.info("Model saved to %s", model_path)
        providers = ['CPUExecutionProvider']
        if rt.get_device() == "GPU":
            gpus = os.environ.get("CUDA_VISIBLE_DEVICES")
            if gpus is None or len(gpus) > 1:
                providers = ['CUDAExecutionProvider']

        opt = rt.SessionOptions()
        if self.use_custom_ops:
            from onnxruntime_extensions import get_library_path
            opt.register_custom_ops_library(get_library_path())
        if self.ort_profile is not None:
            opt.enable_profiling = True
        m = rt.InferenceSession(model_path, sess_options=opt, providers=providers)
        results = m.run(outputs, inputs)
        if self.perf:
            n = 0
            start = time.time()
            stop = start + PERF_TIME
            while time.time() < stop:
                for _ in range(PERF_STEP):
                    _ = m.run(outputs, inputs)
                n += PERF_STEP
            self.onnx_runtime = 1000 * (time.time() - start) / n
            logger.info("ORT perf {:.2f}ms/inference, n={}".format(self.onnx_runtime, n))
        if self.ort_profile is not None:
            tmp_path = m.end_profiling()
            shutil.move(tmp_path, self.ort_profile)
        return results

    @staticmethod
    def create_onnx_file(name, model_proto, inputs, outdir, external_tensor_storage=None):
        os.makedirs(outdir, exist_ok=True)
        if external_tensor_storage is None:
            model_path = os.path.join(outdir, name + ".onnx")
            utils.save_protobuf(model_path, model_proto)
        else:
            model_path = os.path.join(outdir, name + ".zip")
            utils.save_onnx_zip(model_path, model_proto, external_tensor_storage)
        logger.info("Created %s", model_path)

    def run_test(self, name, backend="onnxruntime", onnx_file=None, opset=None, extra_opset=None, perf=None):
        """Run complete test against backend."""
        self.perf = perf

        # get the model
        if self.url:
            _, dir_name = self.download_model()
            logger.info("Downloaded to %s", dir_name)
            model_path = os.path.join(dir_name, self.local) if self.local != "." else dir_name
        else:
            model_path = self.local

        logger.info("Load model from %s", model_path)
        tf_reset_default_graph()
        tf.keras.backend.clear_session()
        input_names = list(self.input_names.keys())
        initialized_tables = {}
        outputs = self.output_names
        tflite_path = None
        tfjs_path = None
        to_rename = {}
        if self.model_type in ["checkpoint"]:
            graph_def, input_names, outputs = tf_loader.from_checkpoint(model_path, input_names, outputs)
        elif self.model_type in ["saved_model"]:
            loaded = tf_loader.from_saved_model(model_path, None, None, self.tag, self.signatures,
                                                self.concrete_function, self.large_model,
                                                return_concrete_func=not self.run_tf_frozen,
                                                return_initialized_tables=True, return_tensors_to_rename=True)
            if not self.run_tf_frozen:
                # Must maintain ref to imported since concrete_func uses weak refs
                # pylint: disable=unused-variable
                graph_def, input_names, outputs, concrete_func, imported, initialized_tables, to_rename = loaded
            else:
                graph_def, input_names, outputs, initialized_tables, to_rename = loaded
        elif self.model_type in ["keras"]:
            graph_def, input_names, outputs = tf_loader.from_keras(model_path, input_names, outputs)
        elif self.model_type in ["tflite"]:
            tflite_path = model_path
            graph_def = None
        elif self.model_type in ["tfjs"]:
            tfjs_path = model_path
            graph_def = None
        else:
            graph_def, input_names, outputs = tf_loader.from_graphdef(model_path, input_names, outputs)

        if utils.is_debug_mode():
            utils.save_protobuf(os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"), graph_def)

        if tflite_path is not None:
            inputs = {}
            for k in input_names:
                v = self.input_names[k]
                inputs[k] = self.make_input(v)

            interpreter = tf.lite.Interpreter(tflite_path)
            input_details = interpreter.get_input_details()
            output_details = interpreter.get_output_details()
            input_name_to_index = {n['name'].split(':')[0]: n['index'] for n in input_details}
            for k, v in inputs.items():
                interpreter.resize_tensor_input(input_name_to_index[k], v.shape)
            interpreter.allocate_tensors()
            def run_tflite():
                for k, v in inputs.items():
                    interpreter.set_tensor(input_name_to_index[k], v)
                interpreter.invoke()
                result = [interpreter.get_tensor(output['index']) for output in output_details]
                return result
            if not self.skip_tensorflow:
                tf_results = run_tflite()
                if self.perf:
                    logger.info("Running TFLite perf")
                    n = 0
                    start = time.time()
                    stop = start + PERF_TIME
                    while time.time() < stop:
                        for _ in range(PERF_STEP):
                            _ = run_tflite()
                        n += PERF_STEP
                    self.tf_runtime = 1000 * (time.time() - start) / n
                    logger.info("TFLite perf {:.2f}ms/inference, n={}".format(self.tf_runtime, n))
                logger.info("TFLite OK")

        if tfjs_path is not None:
            inputs = {}
            for k in input_names:
                v = self.input_names[k]
                inputs[k] = self.make_input(v)
            if not self.skip_tensorflow:
                logger.info("Running TFJS")
                tf_results = run_tfjs(tfjs_path, inputs, outputs)
                logger.info("TFJS OK")

        if not self.run_tf_frozen:
            inputs = {}
            for k in input_names:
                v = self.input_names[k]
                inputs[to_rename.get(k, k)] = tf.constant(self.make_input(v))
            tf_func = tf.function(concrete_func)
            logger.info("Running TF")
            tf_results_d = tf_func(**inputs)
            # If there is only a single output a dict might not be returned
            if isinstance(tf_results_d, tf.Tensor):
                tf_results = [tf_results_d]
            else:
                tf_results = [tf_results_d[k] for k in sorted(tf_results_d.keys())]
            tf_results = [tf_res.numpy() for tf_res in tf_results]
            if self.perf:
                logger.info("Running TF perf")
                n = 0
                start = time.time()
                stop = start + PERF_TIME
                if self.tf_profile is not None:
                    tf.profiler.experimental.start(self.tf_profile)
                while time.time() < stop:
                    for _ in range(PERF_STEP):
                        _ = concrete_func(**inputs)
                    n += PERF_STEP
                if self.tf_profile is not None:
                    tf.profiler.experimental.stop()
                self.tf_runtime = 1000 * (time.time() - start) / n
                logger.info("TF perf {:.2f}ms/inference, n={}".format(self.tf_runtime, n))
            logger.info("TensorFlow OK")

        const_node_values = None
        tf_graph = None

        if graph_def is not None:
            inputs = {}
            tf_reset_default_graph()

            with tf.Graph().as_default() as tf_graph:
                from tf2onnx.tf_utils import compress_graph_def
                if self.large_model:
                    const_node_values = compress_graph_def(graph_def)
                tf.import_graph_def(graph_def, name='')

            with tf_session(graph=tf_graph) as sess:
                # create the input data
                for k in input_names:
                    v = self.input_names[k]
                    t = sess.graph.get_tensor_by_name(k)
                    expected_dtype = tf.as_dtype(t.dtype).name
                    if isinstance(v, str) and v.startswith("np."):
                        np_value = eval(v)  # pylint: disable=eval-used
                        if expected_dtype != np_value.dtype:
                            logger.warning("dtype mismatch for input %s: expected=%s, actual=%s", k, expected_dtype,
                                           np_value.dtype)
                        inputs[k] = np_value.astype(expected_dtype)
                    else:
                        if expected_dtype == "string":
                            inputs[k] = self.make_input(v).astype(str).astype(object)
                        else:
                            inputs[k] = self.make_input(v).astype(expected_dtype)

                # run the model with tensorflow
                if self.skip_tensorflow:
                    logger.info("TensorFlow SKIPPED")
                elif self.run_tf_frozen:
                    if self.tf_profile is not None:
                        tf.profiler.experimental.start(self.tf_profile)
                    tf_results = self.run_tensorflow(sess, inputs)
                    if self.tf_profile is not None:
                        tf.profiler.experimental.stop()
                    logger.info("TensorFlow OK")
                tf_graph = sess.graph


        model_proto = None
        if self.skip_conversion:
            if self.large_model:
                external_tensor_storage = ExternalTensorStorage()
                model_proto = utils.model_proto_from_zip(self.converted_model, external_tensor_storage)
            else:
                external_tensor_storage = None
                model_proto = utils.model_proto_from_file(self.converted_model)
            logger.info("ONNX loaded from file")
        else:
            try:
                # convert model to onnx
                if self.force_input_shape:
                    shape_override = {k: list(v.shape) for k, v in inputs.items()}
                else:
                    shape_override = None
                onnx_graph = self.to_onnx(tf_graph, opset=opset, extra_opset=extra_opset,
                                          shape_override=shape_override, input_names=inputs.keys(),
                                          const_node_values=const_node_values,
                                          initialized_tables=initialized_tables, tflite_path=tflite_path,
                                          tensors_to_rename=to_rename, tfjs_path=tfjs_path)
                onnx_graph = optimizer.optimize_graph(onnx_graph)
                print("ONNX", onnx_graph.dump_node_statistics())
                external_tensor_storage = ExternalTensorStorage() if self.large_model else None
                model_proto = onnx_graph.make_model("converted from tf2onnx",
                                                    external_tensor_storage=external_tensor_storage)
                logger.info("To_ONNX, OK")
                if onnx_file:
                    self.create_onnx_file(name, model_proto, inputs, onnx_file, external_tensor_storage)
                if self.converted_model:
                    if self.large_model:
                        utils.save_onnx_zip(self.converted_model, model_proto, external_tensor_storage)
                    else:
                        utils.save_protobuf(self.converted_model, model_proto)
                    logger.info("Created %s", self.converted_model)

            except Exception:
                logger.error("To_ONNX FAIL", exc_info=1)
                return False

        try:
            onnx_results = None
            if backend == "onnxruntime":
                struc_outputs = [to_rename.get(k, k) for k in self.output_names]
                struc_inputs = {to_rename.get(k, k): v for k, v in inputs.items()}
                onnx_results = self.run_onnxruntime(
                    name, model_proto, struc_inputs, struc_outputs, external_tensor_storage)
            else:
                raise ValueError("unknown backend")
            logger.info("Run_ONNX OK")

            try:
                if self.skip_tensorflow:
                    logger.info("Results: skipped tensorflow")
                else:
                    if self.check_only_shape:
                        for tf_res, onnx_res in zip(tf_results, onnx_results):
                            np.testing.assert_array_equal(tf_res.shape, onnx_res.shape)
                    else:
                        for tf_res, onnx_res in zip(tf_results, onnx_results):
                            good_cnt = np.count_nonzero(np.isclose(tf_res, onnx_res, rtol=self.rtol, atol=self.atol))
                            bad_cnt = tf_res.size - good_cnt
                            if bad_cnt > self.ptol * tf_res.size:
                                # Prints a nice error message with stats
                                np.testing.assert_allclose(tf_res, onnx_res, rtol=self.rtol, atol=self.atol)
                    logger.info("Results: OK")
                return True
            except Exception:
                logger.error("Results", exc_info=1)

        except Exception:
            logger.error("Run_ONNX FAIL", exc_info=1)

        return False

    def check_opset_constraints(self, opset, extra_opset=None):
        """ Return (condition, reason) tuple, condition is True if constraints are met. """
        if not self.opset_constraints:
            return True, None

        opsets = {"onnx": opset}
        if extra_opset:
            for e in extra_opset:
                opsets[e.domain] = e.version

        for constraint in self.opset_constraints:
            domain = constraint.domain
            opset_version = opsets.get(domain)
            if not opset_version:
                return False, "conversion requires opset {}".format(domain)

            if constraint.min_version and opset_version < constraint.min_version:
                reason = "conversion requires opset {} >= {}".format(domain, constraint.min_version)
                return False, reason

            if constraint.max_version and opset_version > constraint.max_version:
                reason = "conversion requires opset {} <= {}".format(domain, constraint.max_version)
                return False, reason

            if constraint.excluded_version:
                if utils.is_list_or_tuple(constraint.excluded_version):
                    skip = opset_version in constraint.excluded_version
                else:
                    skip = opset_version == constraint.excluded_version
                if skip:
                    reason = "conversion requires opset {} != {}".format(domain, constraint.excluded_version)
                    return False, reason

        return True, None


def get_args():
    """Parse commandline."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--cache", default=os.path.join(tempfile.gettempdir(), 'pre-trained'),
                        help="pre-trained models cache dir")
    parser.add_argument("--config", default="tests/run_pretrained_models.yaml", help="yaml config to use")
    parser.add_argument("--tests", help="tests to run")
    parser.add_argument("--target", default="", help="target platform")
    parser.add_argument("--backend", default="onnxruntime",
                        choices=["onnxruntime"], help="backend to use")
    parser.add_argument("--opset", type=int, default=None, help="opset to use")
    parser.add_argument("--extra_opset", default=None,
                        help="extra opset with format like domain:version, e.g. com.microsoft:1")
    parser.add_argument("--skip_tf_tests", help="skip non-tflite tests", default="False")
    parser.add_argument("--skip_tflite_tests", help="skip tflite tests", default="False")
    parser.add_argument("--skip_tfjs_tests", help="skip tfjs tests", default="False")
    parser.add_argument("--verbose", "-v", help="verbose output, option is additive", action="count")
    parser.add_argument("--debug", help="debug mode", action="store_true")
    parser.add_argument("--list", help="list tests", action="store_true")
    parser.add_argument("--onnx-file", help="create onnx file in directory")
    parser.add_argument("--perf", help="capture performance numbers")
    parser.add_argument("--include-disabled", help="include disabled tests", action="store_true")
    args = parser.parse_args()

    args.target = args.target.split(",")
    args.skip_tf_tests = args.skip_tf_tests.upper() == "TRUE"
    args.skip_tflite_tests = args.skip_tflite_tests.upper() == "TRUE"
    args.skip_tfjs_tests = args.skip_tfjs_tests.upper() == "TRUE"
    if args.extra_opset:
        tokens = args.extra_opset.split(':')
        if len(tokens) != 2:
            raise ValueError("invalid extra_opset argument")
        args.extra_opset = [utils.make_opsetid(tokens[0], int(tokens[1]))]
    return args


def load_tests_from_yaml(path):
    """Create test class from yaml file."""
    path = os.path.abspath(path)
    base_dir = os.path.dirname(path)

    tests = {}
    config = yaml.safe_load(open(path, 'r').read())
    for name, settings in config.items():
        if name in tests:
            raise ValueError("Found duplicated test: {}".format(name))

        # parse model and url, non-absolute local path is relative to yaml directory
        model = settings.get("model")
        url = settings.get("url")
        if not url and not os.path.isabs(model):
            model = os.path.join(base_dir, model)

        # parse input_get
        input_func = settings.get("input_get")
        input_func = _INPUT_FUNC_MAPPING[input_func]

        # parse inputs, non-absolute npy file path for np.load is relative to yaml directory
        inputs = settings.get("inputs")
        for k, v in list(inputs.items()):
            if isinstance(v, str):
                # assume at most 1 match
                matches = re.findall(r"np\.load\((r?['\"].*?['\"])", v)
                if matches:
                    npy_path = matches[0].lstrip('r').strip("'").strip('"')
                    if not os.path.isabs(npy_path):
                        abs_npy_path = os.path.join(base_dir, npy_path)
                        inputs[k] = v.replace(matches[0], "r'{}'".format(abs_npy_path))

        # parse opset_constraints
        opset_constraints = []
        section = settings.get("opset_constraints")
        if section:
            for k, v in section.items():
                c = OpsetConstraint(k, min_version=v.get("min"), max_version=v.get("max"),
                                    excluded_version=v.get("excluded"))
                opset_constraints.append(c)

        kwargs = {}
        for kw in ["rtol", "atol", "ptol", "disabled", "check_only_shape", "model_type", "concrete_function",
                   "skip_tensorflow", "force_input_shape", "tf_min_version", "tag", "skip_conversion",
                   "converted_model", "signature_def", "large_model", "structured_outputs", "run_tf_frozen",
                   "use_custom_ops", "dequantize", "ort_profile", "tf_profile"]:
            if settings.get(kw) is not None:
                kwargs[kw] = settings[kw]

        test = Test(url, model, input_func, inputs, settings.get("outputs"),
                    opset_constraints=opset_constraints, **kwargs)
        tests[name] = test
    return tests


def main():
    args = get_args()
    logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
    if args.debug:
        utils.set_debug_mode(True)

    Test.cache_dir = args.cache
    Test.target = args.target
    tests = load_tests_from_yaml(args.config)
    if args.list:
        logger.info(sorted(tests.keys()))
        return 0
    if args.tests:
        test_keys = args.tests.split(",")
    else:
        test_keys = list(tests.keys())

    failed = 0
    count = 0
    for test in test_keys:
        logger.info("===================================")

        t = tests[test]
        if args.tests is None:
            if t.disabled and not args.include_disabled:
                logger.info("Skip %s: disabled", test)
                continue

            if args.skip_tfjs_tests and t.model_type == "tfjs":
                logger.info("Skip %s: tfjs test", test)
                continue
            if args.skip_tflite_tests and t.model_type == "tflite":
                logger.info("Skip %s: tflite test", test)
                continue
            if args.skip_tf_tests and t.model_type not in ["tflite", "tfjs"]:
                logger.info("Skip %s: tf test", test)
                continue

            condition, reason = t.check_opset_constraints(args.opset, args.extra_opset)
            if not condition:
                logger.info("Skip %s: %s", test, reason)
                continue

            if t.tf_min_version:
                if tf_utils.get_tf_version() < Version(str(t.tf_min_version)):
                    logger.info("Skip %s: %s %s", test, "Min TF version needed:", t.tf_min_version)
                    continue

        count += 1
        try:
            logger.info("Running %s", test)
            ret = t.run_test(test, backend=args.backend, onnx_file=args.onnx_file,
                             opset=args.opset, extra_opset=args.extra_opset, perf=args.perf)
        except Exception:
            logger.error("Failed to run %s", test, exc_info=1)
            ret = None
        finally:
            if not utils.is_debug_mode():
                utils.delete_directory(TEMP_DIR)
        if not ret:
            failed += 1

    logger.info("===================================")
    logger.info("RESULT: %s failed of %s, backend=%s", failed, count, args.backend)

    if args.perf:
        with open(args.perf, "w") as f:
            f.write("test,tensorflow,onnx\n")
            for test in test_keys:
                t = tests[test]
                if t.perf:
                    # Report perf in ms per inference
                    f.write("{},{},{}\n".format(test, t.tf_runtime, t.onnx_runtime))
    return failed


if __name__ == "__main__":
    sys.exit(main())
