#!/usr/bin/env python3
# -*- mode: python -*-
# =============================================================================
#
#  Copyright (c) 2015-2020, 2024 Qualcomm Technologies, Inc.
#  All Rights Reserved.
#  Confidential and Proprietary - Qualcomm Technologies, Inc.
#
# =============================================================================
from __future__ import print_function
import logging
import argparse
import os
import sys
import traceback

try:
    from qti.aisw.dlc_utils import snpe_dlc_utils
except ImportError as ie:
    print("Failed to find necessary package:")
    print(str(ie))
    print("Please ensure that $SNPE_ROOT/lib/python is in your PYTHONPATH")
    sys.exit(1)


def main():
    try:
        parser = argparse.ArgumentParser()
        parser._action_groups.pop()
        required = parser.add_argument_group('required arguments')
        required.add_argument('-i', '--input_dlc', required=True, type=str,
                              help="Path to a DLC file")

        optional = parser.add_argument_group('optional arguments')
        optional.add_argument('-s', '--save', type=str, required=False,
                              help="Save the output to a csv file. Specify a target file path.")
        optional.add_argument('-m', '--memory', action="store_true", required=False,
                              help="Show detailed information about memory usage.")
        optional.add_argument('-d', '--display_all_encodings', action="store_true", required=False,
                              help="Show detailed axis-quantization encoding information.")
        optional.add_argument('-t', '--dump_framework_trace', action="store_true", required=False,
                              help="Save framework trace info into the csv file that was passed to --save option.")

        args = parser.parse_args()

        output_file_name = args.save
        snpe_dlc_utils.setUpLogger(True)

        logger = logging.getLogger()
        if not os.path.exists(args.input_dlc):
            logger.debug("Cannot find archive DLC file " + args.input_dlc)
            sys.exit(-1)

        model_info = snpe_dlc_utils.ModelInfo(args.input_dlc)
        model_info.dump_info(args.memory, output_file_name, args.display_all_encodings, args.dump_framework_trace)

    except Exception as e:
        traceback.print_exc()
        sys.exit(-2)

if __name__ == "__main__":
    main()
