"""
Usage:
python graph_optimizer.py \
--tf_path ../../tensorflow/ \
--model_folder "path_to_the_model_folder" \
--output_names "activation, accuracy" \
--input_names "x"
"""

import os, argparse
from subprocess import call

import freeze_graph
import tensorflow as tf

dir = os.path.dirname(os.path.realpath(__file__))

fr_name = "_frozen.pb"
op_name = "_optimized.pb"


def graph_freez(model_folder, output_names):
    print("Model folder", model_folder)
    checkpoint = tf.train.get_checkpoint_state(model_folder)
    print(checkpoint)
    checkpoint_path = checkpoint.model_checkpoint_path
    output_graph_filename = checkpoint_path + fr_name

    input_saver_def_path = ""
    input_binary = True
    output_node_names = output_names
    restore_op_name = "save/restore_all"
    filename_tensor_name = "save/Const:0"
    clear_devices = False
    input_meta_graph = checkpoint_path + ".meta"

    freeze_graph.freeze_graph(
        "", input_saver_def_path, input_binary, checkpoint_path,
        output_node_names, restore_op_name, filename_tensor_name,
        output_graph_filename, clear_devices, "", "", input_meta_graph)

    return output_graph_filename


def graph_optimization(tf_path, graph_file, input_names, output_names):
    output_file = graph_file[:-len(fr_name)] + op_name
    tf_path += "bazel-bin/tensorflow/tools/graph_transforms/transform_graph"

    call([tf_path,
          "--in_graph=" + graph_file,
          "--out_graph=" + output_file,
          "--inputs=" + input_names,
          "--outputs=" + output_names,
          """--transforms=
          strip_unused_nodes(type=float, shape="1,299,299,3")
          fold_constants(ignore_errors=true)
          fold_batch_norms
          fold_old_batch_norms"""])


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
            "Script freezes graph and optimize it for mobile usage")
    parser.add_argument(
        "--model",
        type=str,
        help="Path of folder + model name (folder_path/model_name)")
    parser.add_argument(
        "--input_names",
        type=str,
        default="",
        help="Input node names, comma separated.")
    parser.add_argument(
        "--output_names",
        type=str,
        default="",
        help="Output node names, comma separated.")
    parser.add_argument(
        "--tf_path",
        type=str,
        default="../../tensorflow/",
        help="Path to the folder with tensorflow (requires bazel build of graph_transforms)")

    args = parser.parse_args()

    graph = graph_freez(args.model, args.output_names)
    graph_optimization(args.tf_path, graph, args.input_names, args.output_names)