#
# Script for selecting testing images of a COCO-style dataset generated by the script
# ../data_management/databases/classification/make_classification_dataset.py in a consistent 
# manner and predicting the class for it. 
#

import json
import numpy as np
import argparse
import tensorflow as tf
import os
import random
import collections
import shutil


parser = argparse.ArgumentParser(description='Tools for sampling images from the testing split of a dataset and ' + \
                                 'predicting results with it.')
# Mandatory parameters
parser.add_argument('--frozen_graph', type=str,
                    help='Frozen graph of the classification model, which includes preprocessing. You can generate ' + \
                    'it with the scripts in this directory', metavar='PATH_TO_CLASSIFIER_W_PREPROCESSING')
parser.add_argument('--test_json', type=str, help='Path to test.json generated by the make_classification_dataset.py script')
# Optional parameters
parser.add_argument('--output_dir', type=str, help='Path to output directory. We will copy sample images to this location ' + \
                                  ' and also save the output as text file here. Default: "./sample_output"', default='./sample_output')
parser.add_argument('--num_samples', type=int, help='Number of samples to selected. Default: 5', default=10)
args = parser.parse_args()

# Validate parameters
assert os.path.exists(args.test_json)
assert os.path.exists(args.frozen_graph)
os.makedirs(args.output_dir, exist_ok=True)
assert args.num_samples > 0

# Derived parameters

# We assume that the dataset was generated with the make_classification_dataset.py script, 
# hence the images should be located in the same folder as the json
IMAGE_DIR = os.path.dirname(args.test_json)

# Make seletion deterministic
random.seed(0)

# Load frozen graph
model_graph = tf.Graph()
with model_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(args.frozen_graph, 'rb') as fid:
      od_graph_def.ParseFromString(fid.read())
      tf.import_graph_def(od_graph_def, name='')
graph = model_graph

# Get dataset information
with open(args.test_json, 'rt') as fi:
    js = json.load(fi)
    
# Get classes
class_id_to_name = {cat['id']:cat['name'] for cat in js['categories']}

# ...and the class list corresponding to the model outputs by assuming
# that they are in order of their ids
classlist = [class_id_to_name[idx] for idx in sorted(list(class_id_to_name.keys()))]

# Get images of each class
image_id_to_file_name = {im['id']:im['file_name'] for im in js['images']}
class_images = collections.defaultdict(lambda: list())
for ann in js['annotations']:
    class_images[ann['category_id']].append(image_id_to_file_name[ann['image_id']])
# Shuffle the image list
for class_id in class_images.keys():
    random.shuffle(class_images[class_id])

# Start the image sampling
    
# Set of avaiable class IDs, will be filled below
available_classes = list()
selected_images = list()

while len(selected_images) < args.num_samples:
    
    if len(available_classes) < 1:
        available_classes = list(class_images.keys())
        random.shuffle(available_classes)
    sampled_class = available_classes.pop()
    
    # If there are still images left for that class
    if len(class_images[sampled_class]) > 0:
        # Get image for the sampled class, we already shuffled the class images before so 
        # we can simply pop()
        selected_images.append(class_images[sampled_class].pop())

# Start prediction
with model_graph.as_default():
    
    with tf.Session() as sess:
        
        # Collect tensors for input and output
        image_tensor = tf.get_default_graph().get_tensor_by_name('input:0')
        predictions_tensor = tf.get_default_graph().get_tensor_by_name('output:0')
        predictions_tensor = tf.squeeze(predictions_tensor, [0])

        for image_path in selected_images:
            # Read image
            full_image_path = os.path.join(IMAGE_DIR, image_path)
            with open(full_image_path, 'rb') as fi:
                image = sess.run(tf.image.decode_jpeg(fi.read(), channels=3))
                image = image / 255.

            # Run inference
            predictions = sess.run(predictions_tensor, feed_dict={image_tensor: image})

            out_image = os.path.join(args.output_dir, '___'.join(os.path.normpath(image_path).split(os.path.sep)))
            out_txt = os.path.splitext(out_image)[0] + '.txt'
            print('Predicting', image_path)
            with open(out_txt, 'wt') as log:
                # Print output to log file
                print('Predicting', image_path, file=log)
                print('Most likely classes:', file=log)
                for class_id in np.argsort(-predictions)[:5]:
                    print('\t"{}" with confidence {:.2f}%'.format(classlist[class_id], predictions[class_id]*100), file=log)
                shutil.copy2(full_image_path, out_image)
            print('Wrote output to', out_txt)
