#! /usr/bin/env python

import os
import sys
from pprint import pprint
from argparse import ArgumentParser
import numpy as np
import onnx
import onnx_graphsurgeon as gs
from onnx_graphsurgeon.ir.tensor import Constant
from typing import Tuple, Optional

class Color:
    BLACK          = '\033[30m'
    RED            = '\033[31m'
    GREEN          = '\033[32m'
    YELLOW         = '\033[33m'
    BLUE           = '\033[34m'
    MAGENTA        = '\033[35m'
    CYAN           = '\033[36m'
    WHITE          = '\033[37m'
    COLOR_DEFAULT  = '\033[39m'
    BOLD           = '\033[1m'
    UNDERLINE      = '\033[4m'
    INVISIBLE      = '\033[08m'
    REVERCE        = '\033[07m'
    BG_BLACK       = '\033[40m'
    BG_RED         = '\033[41m'
    BG_GREEN       = '\033[42m'
    BG_YELLOW      = '\033[43m'
    BG_BLUE        = '\033[44m'
    BG_MAGENTA     = '\033[45m'
    BG_CYAN        = '\033[46m'
    BG_WHITE       = '\033[47m'
    BG_DEFAULT     = '\033[49m'
    RESET          = '\033[0m'


def shrinking(
    input_onnx_file_path: str,
    output_onnx_file_path: str,
    mode: Optional[str] = 'shrink',
    non_verbose: Optional[bool] = False,
) -> Tuple[onnx.ModelProto, str]:

    """
    Parameters
    ----------
    input_onnx_file_path: str
        Input onnx file path.

    output_onnx_file_path: str
        Outpu onnx file path.

    mode: Optional[str]
        Constant Value Compression Mode.\n\
        'shrink': Share constant values inside the model as much as possible.\n\
            The model size is slightly larger because some shared constant values remain\n\
            inside the model, but performance is maximized.\n\
        'npy': Outputs constant values used repeatedly in the model to an external file .npy.\n\
            Instead of the smallest model body size, the file loading overhead is greater.\n\
        Default: shrink

    non_verbose: Optional[bool]
        Do not show all information logs. Only error logs are displayed.

    Returns
    -------
    shrunken_graph: onnx.ModelProto
        Shrunken onnx ModelProto

    npy_file_paths:
        List of paths to externally output .npy files.
        An empty list is always returned when in 'shrink' mode.
    """
    # Loading Graphs
    graph = gs.import_onnx(onnx.load(input_onnx_file_path))

    # Constant Value Extraction
    constants = {}
    for graph_node in graph.nodes:
        for graph_node_input in graph_node.inputs:
            if not isinstance(graph_node_input, Constant):
                continue
            if len(graph_node_input.shape) == 0:
                continue
            if np.isscalar(graph_node_input.values):
                continue
            constants[graph_node_input.name] = graph_node_input
    if not non_verbose:
        print(
            f'{Color.GREEN}INFO:{Color.RESET} '+
            f'Number of constant values to be studied: {len(constants)}'
        )

    constants_list = list(constants.items())

    # Processing is performed only when the number of constant tensors to be processed is 2 or more
    if len(constants_list) >= 2:
        # Aggregate list generation
        aggregate_constants_name_list = {}
        aggregate_constants_value_list = {}
        aggregate_constants_matched_list = [False] * len(constants_list)

        for comparator_idx in range(0, len(constants_list)-1):
            if not aggregate_constants_matched_list[comparator_idx]:
                for comparison_dest_idx in range(comparator_idx+1, len(constants_list)):
                    if not aggregate_constants_matched_list[comparison_dest_idx]:
                        # constants_list[0] = Key, constants_list[1] = Value
                        if (constants_list[comparator_idx][1].values.shape == constants_list[comparison_dest_idx][1].values.shape) and \
                            (constants_list[comparator_idx][1].values.dtype == constants_list[comparison_dest_idx][1].values.dtype) and \
                            (constants_list[comparator_idx][1].values == constants_list[comparison_dest_idx][1].values).all():

                            aggregate_constants_matched_list[comparator_idx] = True
                            aggregate_constants_matched_list[comparison_dest_idx] = True
                            aggregate_constants_name_list.setdefault(constants_list[comparator_idx][0], [])
                            aggregate_constants_name_list[constants_list[comparator_idx][0]].append(constants_list[comparison_dest_idx][0])
                            aggregate_constants_value_list.setdefault(constants_list[comparator_idx][0], constants_list[comparator_idx][1])

    """
    aggregate_constants
    {'425': ['434', '5506'], '1646': ['1910', '2608', '2872', '3570', '3834'], '5550': ['4107']}
    """
    npy_file_paths = []
    for layer_name_quoted_src, layer_name_quoted_dists in aggregate_constants_name_list.items():
        i = None
        if mode == 'npy':
            # Export constant values to a numpy file
            external_file_name = f"{os.path.splitext(os.path.basename(output_onnx_file_path))[0]}_exported_{layer_name_quoted_src.replace(':','_').replace(';','_').replace('/','_').replace(',','_')}.npy"
            np.save(
                external_file_name,
                aggregate_constants_value_list[layer_name_quoted_src].values
            )
            npy_file_paths.append(external_file_name)
            # Generate Inputs
            i = gs.Variable(
                name=external_file_name,
                dtype=aggregate_constants_value_list[layer_name_quoted_src].values.dtype,
                shape=aggregate_constants_value_list[layer_name_quoted_src].values.shape,
            )
            aggregate_constants_value_list[layer_name_quoted_src] = i
            graph.inputs.append(i)

        for graph_node_idx, graph_node in enumerate(graph.nodes):
            for input_idx, graph_node_input in enumerate(graph_node.inputs):
                for layer_name_quoted_dist in layer_name_quoted_dists:
                    if not graph_node_input.name == layer_name_quoted_src and graph_node_input.name == layer_name_quoted_dist:
                        graph.nodes[graph_node_idx].inputs[input_idx] = aggregate_constants_value_list[layer_name_quoted_src]
                    if mode == 'npy' and graph_node_input.name == layer_name_quoted_src:
                        graph.nodes[graph_node_idx].inputs[input_idx] = i

    graph.cleanup().toposort()
    if not non_verbose:
        print(f'{Color.GREEN}INFO:{Color.RESET} Results:')
        pprint(aggregate_constants_name_list)
        if len(npy_file_paths) > 0:
            print(f'{Color.GREEN}INFO:{Color.RESET} .npy files:')
            pprint(npy_file_paths)

    shrunken_graph = gs.export_onnx(graph)

    new_model = None
    try:
        new_model = onnx.shape_inference.infer_shapes(shrunken_graph)
    except:
        new_model = shrunken_graph
        if not non_verbose:
            print(
                f'{Color.YELLOW}WARNING:{Color.RESET} '+
                'The input shape of the next OP does not match the output shape. '+
                'Be sure to open the .onnx file to verify the certainty of the geometry.'
            )

    # Save
    onnx.save(new_model, f'{output_onnx_file_path}')

    return new_model, npy_file_paths


def main():
    parser = ArgumentParser()
    parser.add_argument(
        'input_onnx_file_path',
        type=str,
        help='Input onnx file path.'
    )
    parser.add_argument(
        'output_onnx_file_path',
        type=str,
        help='Output onnx file path.'
    )
    parser.add_argument(
        '--mode',
        type=str,
        choices=[
            'shrink',
            'npy',
        ],
        default='shrink',
        help="\
            Constant Value Compression Mode. \
            shrink: Share constant values inside the model as much as possible. \
            The model size is slightly larger because some shared constant values remain \
            inside the model, but performance is maximized. \
            npy: Outputs constant values used repeatedly in the model to an external file .npy. \
            Instead of the smallest model body size, the file loading overhead is greater. \
            Default: shrink"
    )
    parser.add_argument(
        '--non_verbose',
        action='store_true',
        help='Do not show all information logs. Only error logs are displayed.'
    )
    args = parser.parse_args()

    # file existence check
    if not os.path.exists(args.input_onnx_file_path) or \
        not os.path.isfile(args.input_onnx_file_path) or \
        not os.path.splitext(args.input_onnx_file_path)[-1] == '.onnx':

        print(
            f'{Color.RED}ERROR:{Color.RESET} '+
            f'The specified file (.onnx) does not exist. or not an onnx file. File: {args.input_onnx_file_path}'
        )
        sys.exit(1)

    # Model shrink
    shrunken_graph, npy_file_paths = shrinking(args.input_onnx_file_path, args.output_onnx_file_path, args.mode, args.non_verbose)

    if not args.non_verbose:
        print(f'{Color.GREEN}INFO:{Color.RESET} Finish!')
