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

import argparse
import glob
import hashlib
import json
import os
import re
import struct
import sys
import tarfile
import time

import numpy as np
from six.moves import urllib
import tensorflow as tf
from tensorflow.contrib.learn import ModeKeys

from tensorflow.python.platform import gfile
from tensorflow.python.util import compat

import transfer_learning

# If you've already downloaded the inception model, and it's elsewhere, 
# edit this path to reflect that so you don't need to re-download.
INCEPTION_MODEL_DIR = '/tmp/imagenet'
DATA_URL = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'
BOTTLENECK_TENSOR_NAME = 'pool_3/_reshape:0'
BOTTLENECK_TENSOR_SIZE = 2048
MODEL_INPUT_WIDTH = 299
MODEL_INPUT_HEIGHT = 299
MODEL_INPUT_DEPTH = 3
JPEG_DATA_TENSOR_NAME = 'DecodeJpeg/contents:0'
RESIZED_INPUT_TENSOR_NAME = 'ResizeBilinear:0'

LABELS_FILENAME = "output_labels.json"


"""
Explanation: Copyright 2016 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.
This notebook lets you run predictions against an image classification model trained with transfer_learning.py (and bootstrapping from the saved Inception v3 image classification model), in this same directory. See the README in this directory for more information on running the training on a set of photos first.
First, some imports and definitions:
End of explanation
"""


# Replace MODEL_DIR with the path to the directory in which your learned model resides.
MODEL_DIR = '/tmp/tfmodels/img_classify/your-model-dir'

"""
Explanation: Edit the following to point to the model directory in which the trained model that you want to use resides.  If you just did a training run, the directory name will have been printed to STDOUT.
End of explanation
"""


def create_inception_graph():
  """"Creates a graph from saved GraphDef file and returns a Graph object.
  """
  with tf.Session() as sess:
    model_filename = os.path.join(
        INCEPTION_MODEL_DIR, 'classify_image_graph_def.pb')
    with gfile.FastGFile(model_filename, 'rb') as f:
      graph_def = tf.GraphDef()
      graph_def.ParseFromString(f.read())
      bottleneck_tensor, jpeg_data_tensor, resized_input_tensor = (
          tf.import_graph_def(graph_def, name='', return_elements=[
              BOTTLENECK_TENSOR_NAME, JPEG_DATA_TENSOR_NAME,
              RESIZED_INPUT_TENSOR_NAME]))
  return sess.graph, bottleneck_tensor, jpeg_data_tensor, resized_input_tensor


def run_bottleneck_on_image(sess, image_data, image_data_tensor,
                            bottleneck_tensor):
  """Runs inference on an image to extract the 'bottleneck' summary layer.
  """
  bottleneck_values = sess.run(
      bottleneck_tensor,
      {image_data_tensor: image_data})
  bottleneck_values = np.squeeze(bottleneck_values)
  return bottleneck_values

"""
Explanation: Now define some helper functions. You'll find these functions in transfer_learning.py also.
(In run_bottleneck_on_image, note that we're calling sess.run() to get the value of the 'bottleneck' layer of the Inception graph, with image data fed to the JPEG_DATA_TENSOR_NAME node.)
End of explanation
"""


# load the labels list, needed to create the model; if it's 
# not there, we can't proceed
output_labels_file = os.path.join(MODEL_DIR, "output_labels.json")
if gfile.Exists(output_labels_file):
  with open(output_labels_file, 'r') as lfile:
    labels_string = lfile.read()
    labels_list = json.loads(labels_string)
    print("labels list: %s" % labels_list)
    class_count = len(labels_list)
else:
  print("Labels list %s not found; we can't proceed." % output_labels_file)


"""
Explanation: Next, we need to load the file that gives us the class name ordering used for the result vectors during training. (Since this info was generated from reading the photos directories structure, the ordering can potentially change.  We need to make sure that doesn't happen, so that we interpret the prediction results consistently).
End of explanation
"""


def make_image_predictions(
    classifier, jpeg_data_tensor, bottleneck_tensor, path_list, labels_list):
  """Use the learned model to make predictions."""

  if not labels_list:
    output_labels_file = os.path.join(MODEL_DIR, LABELS_FILENAME)
    if gfile.Exists(output_labels_file):
      with open(output_labels_file, 'r') as lfile:
        labels_string = lfile.read()
        labels_list = json.loads(labels_string)
        print("labels list: %s" % labels_list)
    else:
      print("Labels list %s not found" % output_labels_file)
      return None

  sess = tf.Session()
  bottlenecks = []
  print("Predicting for images: %s" % path_list)
  for img_path in path_list:
    # get bottleneck for an image path.
    if not gfile.Exists(img_path):
      tf.logging.fatal('File does not exist %s', img_path)
    image_data = gfile.FastGFile(img_path, 'rb').read()
    bottleneck_values = run_bottleneck_on_image(sess, image_data,
                                                jpeg_data_tensor,
                                                bottleneck_tensor)
    bottlenecks.append(bottleneck_values)
  prediction_input = np.array(bottlenecks)
  predictions = classifier.predict(x=prediction_input, as_iterable=True)
  print("Predictions:")
  for _, p in enumerate(predictions):
    print("---------")
    for k in p.keys():
      print("%s is: %s " % (k, p[k]))
      if k == "index":
        print("index label is: %s" % labels_list[p[k]])

"""
Explanation: Define a function to run the image predictions. First, we need to get the 'bottleneck' values, using the graph loaded from the Inception model. Then, we feed that data to our own trained model.
classifier is a custom Estimator, and we will use its predict method. (We'll define the Estimator in a few more cells).
End of explanation
"""


# Set up the pre-trained graph 
transfer_learning.maybe_download_and_extract(INCEPTION_MODEL_DIR)
graph, bottleneck_tensor, jpeg_data_tensor, resized_image_tensor = (
    create_inception_graph())
print("Finished.")

"""
Explanation: Define the Inception-based graph we'll use to generate the 'bottleneck' values. Wait for this to print "Finished" before continuing.
End of explanation
"""


# Define the custom estimator
model_fn = transfer_learning.make_model_fn(class_count, 'final_result')
model_params = {}
classifier = tf.contrib.learn.Estimator(
    model_fn=model_fn, params=model_params, model_dir=MODEL_DIR)

img_list = transfer_learning.get_prediction_images('prediction_images')

"""
Explanation: Define our custom Estimator. (As the lab exercise, you will write some of the code that does this).
End of explanation
"""


# If PIL/Pillow is not installed, this step is not important
import PIL.Image
from IPython.display import display
for imgfile in img_list:
    img = PIL.Image.open(imgfile)
    display(img)
    

"""
Explanation: For fun, display the images that we're going to predict the classification for.
End of explanation
"""


make_image_predictions(
    classifier, jpeg_data_tensor, bottleneck_tensor, img_list, labels_list)

"""
Explanation: Run the predict() method of our Estimator to predict the classifications of our list of images.
End of explanation
"""


