### tf 1.x scripts ###
import tensorflow as tf
from typing import List

def load_graph(frozen_graph_filename):
    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        return graph_def

tf.import_graph_def(load_graph("frozen_inference_graph.pb"))

# the saved model is a model generated by tf.saved_model.builder and is has to be imported into a
#  session, this file contains the full graph with all training weights (just like the frozen 
# graph) but here can be trained upon, and this one is not serialized and needs to be loaded by
#  this snippet. The [] are tagconstants which can be read by the saved_model_cli. 
# This model is also often served to predict on, like google ml engine par example:

with tf.Session() as sess:
    tf.saved_model.loader.load(sess, [], "foldername to saved_model.pb, only folder")


# ANOTHER ADDITION:

# The code to load models is already provided above. To actually predict you need a session, 
# for a saved model this session is already created, for a frozen model, it's not.

# saved model:

with tf.Session() as sess:
    tf.saved_model.loader.load(sess, [], "foldername to saved_model.pb, only folder")
    output_tensor = None # set your tensor here
    prediction = sess.run(output_tensor, feed_dict={input_tensor: test_images})

# Frozen model:

tf.import_graph_def(load_graph("frozen_inference_graph.pb"))
with tf.Session() as sess:
    prediction = sess.run(output_tensor, feed_dict={input_tensor: test_images})

# To further understand what your input and output layers are, you need to check them out with
#  tensorboard, simply add the following line of code into your session:

tf.summary.FileWriter("path/to/folder/to/save/logs", sess.graph)







#######################
#######################
#######################

# I wrote a simple script to analyze the dependency relations in a computational graph 
# (usually a DAG, directly acyclic graph). 
# It's so obvious that the inputs are the nodes that lack a input. 
# However, outputs can be defined as any nodes in a graph 
# because, in the weirdest but still valid case, outputs can be inputs while the other nodes are 
# all dummy. I still define the output operations as nodes without output in the code. 
# You could neglect it at your willing.

import tensorflow as tf

def load_graph(frozen_graph_filename):
    with tf.io.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.compat.v1.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def)
    return graph

def analyze_inputs_outputs(graph):
    ops = graph.get_operations()
    outputs_set = set(ops)
    inputs = []
    for op in ops:
        if len(op.inputs) == 0 and op.type != 'Const':
            inputs.append(op)
        else:
            for input_tensor in op.inputs:
                if input_tensor.op in outputs_set:
                    outputs_set.remove(input_tensor.op)
    outputs = list(outputs_set)
    return (inputs, outputs)

#####################################
def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    """
    Freezes the state of a session into a pruned computation graph.

    Creates a new computation graph where variable nodes are replaced by
    constants taking their current value in the session. The new graph will be
    pruned so subgraphs that are not necessary to compute the requested
    outputs are removed.
    @param session The TensorFlow session to be frozen.
    @param keep_var_names A list of variable names that should not be frozen,
                          or None to freeze all the variables in the graph.
    @param output_names Names of the relevant graph outputs.
    @param clear_devices Remove the device directives from the graph for better portability.
    @return The frozen graph definition.
    """
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = tf.graph_util.convert_variables_to_constants(
            session, input_graph_def, output_names, freeze_var_names)
        return frozen_graph



def freeze_pb(metafile:str,output_nodes_list:List[str],name='frozen_pb'): 
    # name is without .pb extension
    pbname = name+'.pb'
    folder_to_save = "./freezed_pb"
    ckptfile = metafile.replace("meta",'')
    tf.reset_default_graph()
    saver = tf.train.import_meta_graph(metafile)
    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)) as sess:
        saver.restore(sess.ckptfile)
        print("Model Restored Successfully")
        output_nodes=output_nodes_list
        frozen_graph =  freeze_session (sess, output_names=output_nodes) # provide your output node list: by default freezing all nodes names
        tf.train.write_graph(frozen_graph, folder_to_save, pbname, as_text=False) 
        print("Saved file in frozen_pb format named as '%s'" %(pbname))





import tensorflow as tf
from argparse import ArgumentParser

def main():
    parser = ArgumentParser()
    parser.add_argument('--checkpoint', type=str,
                        dest='checkpoint',
                        help='dir or .ckpt file to load checkpoint from',
                        metavar='CHECKPOINT', required=True)
    parser.add_argument('--model', type=str,
                        dest='model',
                        help='.meta for your model',
                        metavar='MODEL', required=True)
    parser.add_argument('--out-path', type=str,
                        dest='out_path',
                        help='model output directory',
                        metavar='MODEL_OUT', required=True)
    opts = parser.parse_args()
    tf.reset_default_graph()
    saver = tf.train.import_meta_graph(opts.model)
    builder = tf.saved_model.builder.SavedModelBuilder(opts.out_path)
    with tf.Session() as sess:
        # Restore variables from disk.
        saver.restore(sess, opts.checkpoint)
        print("Model restored.")
        builder.add_meta_graph_and_variables(sess,
                                       ['tfckpt2pb'],
                                       strip_default_attrs=False)
        builder.save()

# if __name__ == '__main__':
#     main()


from keras import backend as K

# Create, compile and train model...

frozen_graph = freeze_session(K.get_session(),
                              output_names=[out.op.name for out in model.outputs])