# 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.
# ==============================================================================
r"""Generate captions for images using default beam search parameters."""

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

import math
import os


import tensorflow as tf

import configuration
import inference_wrapper
from inference_utils import caption_generator
from inference_utils import vocabulary

import glob
import json
import datetime

FLAGS = tf.flags.FLAGS

base_dir_local = "/home/jbs/ImageCaptioning/im2txt/im2txt"
base_dir_ext = "/media/jbs/0b83a3c6-83d6-481b-9400-ab5e4e011efa/error404/image_captioning/data/mscoco/raw-data"

tf.flags.DEFINE_string("checkpoint_path", base_dir_local+"/model/3m/train",
                       "Model checkpoint file or directory containing a "
                       "model checkpoint file.")
tf.flags.DEFINE_string("vocab_file", base_dir_local+"/model/3m/train/word_counts.txt", "Text file containing the vocabulary.")

set = []

#/media/deckyal/78DCCAA06404648F/mscoco/raw-data

test = False # test=True -> test dataset ||| test=False -> val dataset

if test:
    set = glob.glob(base_dir_ext+"/test2014/*.jpg")
else:
    set = glob.glob(base_dir_ext+"/val2014/*.jpg")
    
#set = ["/media/jbs/0b83a3c6-83d6-481b-9400-ab5e4e011efa/error404/image_captioning/data/mscoco/raw-data/test2014/COCO_test2014_000000000229.jpg"]
set = ','.join(map(str, set))


tf.flags.DEFINE_string("input_files", set, 
                       "File pattern or comma-separated list of file patterns "
                       "of image files.")

tf.logging.set_verbosity(tf.logging.INFO)

def writeToJSONFile(path, fileName, data):
  filePathNameWExt = './' + path + '/' + fileName + '.json'
  with open(filePathNameWExt, 'w') as fp:
    json.dump(data, fp)

def main(_):
  # Build the inference graph.
  g = tf.Graph()
  with g.as_default():
    model = inference_wrapper.InferenceWrapper()
    restore_fn = model.build_graph_from_config(configuration.ModelConfig(),
                                               FLAGS.checkpoint_path)
  g.finalize()

  # Create the vocabulary.
  vocab = vocabulary.Vocabulary(FLAGS.vocab_file)
  filenames = []
  fileptt = 0
  for file_pattern in FLAGS.input_files.split(","):
    fileptt += 1
    if((fileptt%100) == 0):
        print(fileptt)
    filenames.extend(tf.gfile.Glob(file_pattern))
  tf.logging.info("Running caption generation on %d files matching %s",
                  len(filenames), FLAGS.input_files)
  with tf.Session(graph=g) as sess:
    # Load the model from checkpoint.
    restore_fn(sess)

    # Prepare the caption generator. Here we are implicitly using the default
    # beam search parameters. See caption_generator.py for a description of the
    # available beam search parameters.
    generator = caption_generator.CaptionGenerator(model, vocab)

    diverseGroupsCount = [0, 0, 0] #array to track occurrence of each group
    jsonData = []
    filenum = 0
    
    print(datetime.datetime.now().time())
    
    for filename in filenames:
      filenum += 1
      if ((filenum%100) == 0):
          print(" %d images processed" % filenum )
      with tf.gfile.GFile(filename, "rb") as f:
        image = f.read()
      captions = generator.beam_search(sess, image)
      #print("Captions for image %s:" % os.path.basename(filename))
      for i, caption in enumerate(captions):
        # Ignore begin and end words.
        sentence = [vocab.id_to_word(w) for w in caption.sentence[1:-1]]
        sentence = " ".join(sentence)
        #print("  %d) %s (p=%f)" % (i, sentence, math.exp(caption.logprob)))
        
        diverseGroupsCount[caption.group] += 1
        
        data = {}
          
        if test:
          data['image_id'] = int(os.path.basename(filename)[14:-4])
        else:
          data['image_id'] = int(os.path.basename(filename)[13:-4])
            
        if "." in sentence:
          data['caption'] = sentence[0:-2]
        else:
          data['caption'] = sentence 
               
        jsonData.append(data)
        
        
    
    writeToJSONFile('./','testInferences',jsonData)  
      
    print(datetime.datetime.now().time())
      
    print("  Group 0 = %d" % (diverseGroupsCount[0]))
    print("  Group 1 = %d" % (diverseGroupsCount[1]))
    print("  Group 2 = %d" % (diverseGroupsCount[2]))
           
      

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