#
# Script for evaluating precision/recall on a two-class problem given 
# a frozen graph and a COCO-style dataset, which was generated by
# the ../data_management/databases/classification/make_classification_dataset.py script.
# We assume the positive class is at index 0 (with a zero-based indexing)
#

import matplotlib
matplotlib.use('agg')
import argparse
import tensorflow as tf
import os
from PIL import Image
import numpy as np
from pycocotools.coco import COCO
import tqdm
import random
from sklearn.metrics import precision_recall_curve
import matplotlib.pyplot as plt
from inspect import signature


parser = argparse.ArgumentParser('This script evaluates the precision/recall of an exported frozen inference graph on the test data of ' + \
                                'a classification dataset generated with the make_classification_dataset.py script. ')
parser.add_argument('--frozen_graph', type=str,
                    help='Frozen graph of detection network as create by export_inference_graph.py of TFODAPI.')
parser.add_argument('--coco_style_output', type=str,
                    help='Path to directory containing the coco-style output of make_classification_dataset.py')
args = parser.parse_args()

TEST_JSON = os.path.join(args.coco_style_output, 'test_large.json')

# Check that all files exists for easier debugging
assert os.path.exists(args.frozen_graph)
assert os.path.exists(args.coco_style_output)
assert os.path.exists(TEST_JSON), 'Could not find ' + TEST_JSON

# 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

coco = COCO(TEST_JSON)

gt_labels = []
pred_scores = []

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])

        total = 0
        correct = 0
        all_img_keys = list(coco.imgs.keys())
        random.seed(0)
        random.shuffle(all_img_keys)
        for image_id in tqdm.tqdm(all_img_keys):
            # Read image
            image_path = os.path.join(args.coco_style_output, coco.imgs[image_id]['file_name'])
            if not os.path.exists(image_path):
                print('Image {} does not exist'.format(image_path))
                continue
            #with open(image_path, 'rb') as fi:
            #    image =  sess.run(tf.image.decode_jpeg(fi.read(), channels=3))
            #    image = image / 255.
            image = np.array(Image.open(image_path))/255

            # Run inference
            predictions = [0]#sess.run(predictions_tensor, feed_dict={image_tensor: image})
            #predicted_class = np.argmax(predictions)
            pred_scores.append(predictions[0])
            gt_labels.append(1-coco.imgToAnns[image_id][0]['category_id'])

            # Check if correct
            #if coco.imgToAnns[image_id][0]['category_id'] == predicted_class:
            #    correct = correct + 1

            total = total + 1

            if total%100 == 0:
                print('Currently at {:.2f}% top-1 accuracy'.format(correct/total*100))



precision, recall, threshold = precision_recall_curve(gt_labels, pred_scores)

# In matplotlib < 1.5, plt.fill_between does not have a 'step' argument
step_kwargs = ({'step': 'post'}
               if 'step' in signature(plt.fill_between).parameters
                                  else {})

plt.figure()
plt.step(recall, precision, color='b', alpha=0.2,
         where='post')
plt.fill_between(recall, precision, alpha=0.2, color='b', **step_kwargs)

plt.xlabel('Recall')
plt.ylabel('Precision')
plt.ylim([0.0, 1.05])
plt.xlim([0.0, 1.0])
plt.title('2-class Precision-Recall curve')
plt.savefig('pr.png', format='png')
print('Saved PR curve to ./pr.png')

precision = np.array(precision)
recall = np.array(recall)

high_recall_idx = np.where(recall>=0.99)[0]
best_prec_idx = high_recall_idx[np.argmax(precision[high_recall_idx])]
best_prec = precision[best_prec_idx]
best_recall = recall[best_prec_idx]
best_thr = threshold[best_prec_idx]

print('Best precision >= 0.99 recall: recall {}, prec {}, at threshold {}'.format(
    best_recall, best_prec, best_thr))

np.savetxt('prec.txt', precision)
np.savetxt('recall.txt', recall)
np.savetxt('thresholds.txt', threshold)
import ipdb; ipdb.set_trace()
