import tensorflow as tf
import tf2onnx
import argparse
import logging
import sys
import onnx
import onnx.utils
from tensorflow.python.platform import gfile
from tools import combo, eliminating, replacing, other


def tf2onnx_flow(pb_path: str, test_mode=False) -> onnx.ModelProto:
    """Convert frozen graph pb file into onnx

    Args:
        pb_path (str): input pb file path
        test_mode (bool, optional): test mode. Defaults to False.

    Raises:
        Exception: invalid input file

    Returns:
        onnx.ModelProto: converted onnx
    """
    TF2ONNX_VERSION = int(tf2onnx.version.version.replace(".", ""))

    if 160 <= TF2ONNX_VERSION:
        from tf2onnx import tf_loader
    else:
        from tf2onnx import loader as tf_loader

    if pb_path[-3:] == ".pb":
        model_name = pb_path.split("/")[-1][:-3]

        # always reset tensorflow session at begin
        tf.reset_default_graph()

        with tf.Session() as sess:
            with gfile.FastGFile(pb_path, "rb") as f:
                graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name="")

            if 160 <= int(tf2onnx.version.version.replace(".", "")):
                (
                    onnx_nodes,
                    op_cnt,
                    attr_cnt,
                    output_shapes,
                    dtypes,
                    functions,
                ) = tf2onnx.tf_utils.tflist_to_onnx(sess.graph, {})
            else:
                (
                    onnx_nodes,
                    op_cnt,
                    attr_cnt,
                    output_shapes,
                    dtypes,
                ) = tf2onnx.tfonnx.tflist_to_onnx(
                    sess.graph.get_operations(), {}
                )

            for n in onnx_nodes:
                if len(n.output) == 0:
                    onnx_nodes.remove(n)

            # find inputs and outputs of graph
            nodes_inputs = set()
            nodes_outputs = set()

            for n in onnx_nodes:
                if n.op_type == "Placeholder":
                    continue
                for input in n.input:
                    nodes_inputs.add(input)
                for output in n.output:
                    nodes_outputs.add(output)

            graph_input_names = set()
            for input_name in nodes_inputs:
                if input_name not in nodes_outputs:
                    graph_input_names.add(input_name)

            graph_output_names = set()
            for n in onnx_nodes:
                if n.input and n.input[0] not in nodes_outputs:
                    continue
                if len(n.output) == 0:
                    n.output.append(n.name + ":0")
                    graph_output_names.add(n.output[0])
                else:
                    output_name = n.output[0]
                    if (output_name not in nodes_inputs) and (
                        0 < len(n.input)
                    ):
                        graph_output_names.add(output_name)

        logging.info("Model Inputs: %s", str(list(graph_input_names)))
        logging.info("Model Outputs: %s", str(list(graph_output_names)))

        graph_def, inputs, outputs = tf_loader.from_graphdef(
            model_path=pb_path,
            input_names=list(graph_input_names),
            output_names=list(graph_output_names),
        )

        with tf.Graph().as_default() as tf_graph:
            tf.import_graph_def(graph_def, name="")

        if 160 <= TF2ONNX_VERSION:
            with tf_loader.tf_session(graph=tf_graph):
                onnx_graph = tf2onnx.tfonnx.process_tf_graph(
                    tf_graph=tf_graph,
                    input_names=inputs,
                    output_names=outputs,
                    opset=11,
                )
        else:
            with tf.Session(graph=tf_graph):
                onnx_graph = tf2onnx.tfonnx.process_tf_graph(
                    tf_graph=tf_graph,
                    input_names=inputs,
                    output_names=outputs,
                    opset=11,
                )

        # Optimize with tf2onnx.optimizer
        onnx_graph = tf2onnx.optimizer.optimize_graph(onnx_graph)
        model_proto = onnx_graph.make_model(model_name)

        # Make tf2onnx output compatible with the spec. of other.polish_model
        replacing.replace_initializer_with_Constant(model_proto.graph)
        model_proto = other.polish_model(model_proto)

    else:
        raise Exception(
            'expect .pb file as input, but got "' + str(pb_path) + '"'
        )

    # rename
    m = model_proto

    m = combo.preprocess(m)
    m = combo.common_optimization(m)
    m = combo.tensorflow_optimization(m)
    m = combo.postprocess(m)

    if not test_mode:
        g = m.graph
        eliminating.eliminate_shape_changing_after_input(g)

    m = other.polish_model(m)
    return m


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert tensorflow pb file to onnx file and optimized "
                    "onnx file. Or just optimize tensorflow onnx file."
    )
    parser.add_argument("in_file", help="input file")
    parser.add_argument("out_file", help="output optimized model file")
    parser.add_argument(
        "-t",
        "--test_mode",
        default=False,
        help="test mode will not eliminate shape changes after input",
    )

    args = parser.parse_args()
    logging.basicConfig(
        stream=sys.stdout,
        format="[%(asctime)s] %(levelname)s: %(message)s",
        level=logging.INFO,
    )
    m = tf2onnx_flow(args.in_file, args.test_mode)
    onnx.save(m, args.out_file)
    logging.info("Save Optimized ONNX: %s", args.out_file)
