#!/home/annemaelle/tensorflow1.4/bin/python
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
r"""Saves out a GraphDef containing the architecture of the model.

To use it, run something like this, with a model name defined by slim:

bazel build tensorflow_models/research/slim:export_inference_graph
bazel-bin/tensorflow_models/research/slim/export_inference_graph \
--model_name=inception_v3 --output_file=/tmp/inception_v3_inf_graph.pb

If you then want to use the resulting model with your own or pretrained
checkpoints as part of a mobile model, you can run freeze_graph to get a graph
def with the variables inlined as constants using:

bazel build tensorflow/python/tools:freeze_graph
bazel-bin/tensorflow/python/tools/freeze_graph \
--input_graph=/tmp/inception_v3_inf_graph.pb \
--input_checkpoint=/tmp/checkpoints/inception_v3.ckpt \
--input_binary=true --output_graph=/tmp/frozen_inception_v3.pb \
--output_node_names=InceptionV3/Predictions/Reshape_1

The output node names will vary depending on the model, but you can inspect and
estimate them using the summarize_graph tool:

bazel build tensorflow/tools/graph_transforms:summarize_graph
bazel-bin/tensorflow/tools/graph_transforms/summarize_graph \
--in_graph=/tmp/inception_v3_inf_graph.pb

To run the resulting graph in C++, you can look at the label_image sample code:

bazel build tensorflow/examples/label_image:label_image
bazel-bin/tensorflow/examples/label_image/label_image \
--image=${HOME}/Pictures/flowers.jpg \
--input_layer=input \
--output_layer=InceptionV3/Predictions/Reshape_1 \
--graph=/tmp/frozen_inception_v3.pb \
--labels=/tmp/imagenet_slim_labels.txt \
--input_mean=0 \
--input_std=255

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
tf.GraphKeys.VARIABLES = tf.GraphKeys.GLOBAL_VARIABLES

import sys
sys.path.append('/mnt/data_b/anne-maelle/hacone/tensorflow/nn_search')
sys.path.append('/mnt/data_b/anne-maelle/hacone/tensorflow/slim')


from tensorflow.python.platform import gfile
from datasets import dataset_factory
from nets import nets_factory
import os

slim = tf.contrib.slim
PATH_TO_HACONE = '/home/anne-maelle/test_dir/hacone'

tf.app.flags.DEFINE_string(
    'model_name', 'cifarnet', 'The name of the architecture to save.')

tf.app.flags.DEFINE_boolean(
    'is_training', False,
    'Whether to save out a training-focused version of the model.')

tf.app.flags.DEFINE_integer(
    'image_size', None,
    'The image size to use, otherwise use the model default_image_size.')

tf.app.flags.DEFINE_integer(
    'batch_size', None,
    'Batch size for the exported model. Defaulted to "None" so batch size can '
    'be specified at model runtime.')

tf.app.flags.DEFINE_string('dataset_name', 'cifar10',
                           'The name of the dataset to use with the model.')

tf.app.flags.DEFINE_integer(
    'labels_offset', 0,
    'An offset for the labels in the dataset. This flag is primarily used to '
    'evaluate the VGG and ResNet architectures which do not use a background '
    'class for the ImageNet dataset.')

tf.app.flags.DEFINE_string(
    'job_name', 'cifar10_nns', 'Name for the running job')

output_dir = os.path.join(PATH_TO_HACONE,'outputs', tf.app.flags.FLAGS.job_name)
if not os.path.exists(output_dir):
   print('Output directory {} non-existent, creating it'.format(output_dir))
   os.makedirs(output_dir)
output_file = os.path.join(output_dir, 'inference_graph.pb')
tf.app.flags.DEFINE_string(
    'output_file', output_file, 'Where to save the resulting file to.')

tf.app.flags.DEFINE_string(
    'dataset_dir', PATH_TO_HACONE+'/dataset/cifar10', 'Directory to save intermediate dataset files to')

FLAGS = tf.app.flags.FLAGS


def main(_):
  if not FLAGS.output_file:
    raise ValueError('You must supply the path to save to with --output_file')
  tf.logging.set_verbosity(tf.logging.INFO)
  with tf.Graph().as_default() as graph:
    dataset = dataset_factory.get_dataset(FLAGS.dataset_name, 'train',
                                          FLAGS.dataset_dir)
    candidate = [-1, -1, 0, 3, -1, -1, 0, 3, -1, -1, 0, 3, -1, -1, 0, 3, -1, -1, 0, 3]
    #candidate=[-2, -1, 4, 1, 0, 0, 3, 4, -2, 0, 1, 7, -2, -1, 5, 2, 3, 1, 3, 3]
    #candidate = [-1, -1, 0, 2, 0, -2, 6, 5, 0, -2, 5, 4, -1, 1, 5, 2, 1, 3, 4, 5]
    #candidate = [-2, -2, 1, 2, 0, -1, 3, 6, -1, -1, 1, 4, -2, -2, 2, 6, -2, 3, 6, 4]
    #candidate = [-1, -1, 4, 5, -1, -2, 5, 6, 1, -1, 3, 2, 2, 0, 5, 0, 2, -1, 3, 0]
    #candidate = [-1, -2, 3, 0, -2, 0, 6, 7, 0, 0, 5, 5, 1, 2, 0, 5, 3, 2, 4, 7]
    #candidate = [-2, -1, 3, 1, -1, 0, 1, 4, 0, 0, 0, 0, -1, -1, 6, 1, 1, 2, 4, 4]
    #candidate = [-1, -2, 7, 0, -2, -2, 0, 5, 0, -1, 5, 3, -2, 2, 1, 1, -1, -2, 4, 0]
    #candidate = [-1, -2, 7, 0, 0, -2, 1, 1, 0, -2, 0, 2, 1, -2, 4, 1, 2, -2, 0, 1]
    #candidate = [-2, -2, 1, 7, -1, -1, 1, 2, -1, 0, 1, 0, -1, -1, 0, 4, 3, 2, 3, 2]
    N = 2
    F = 24
    network_fn = nets_factory.get_network_fn(
        FLAGS.model_name, candidate, N, F,
        num_classes=(dataset.num_classes - FLAGS.labels_offset),
        is_training=FLAGS.is_training)

    image_size = FLAGS.image_size or network_fn.default_image_size
    placeholder = tf.placeholder(name='input', dtype=tf.float32,
                                 shape=[FLAGS.batch_size, image_size,
                                        image_size, 3])
    network_fn(placeholder)
    graph_def = graph.as_graph_def()
    with gfile.GFile(FLAGS.output_file, 'wb') as f:
      f.write(graph_def.SerializeToString())


if __name__ == '__main__':
  tf.app.run()
