#
# 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 argparse
import os
import random
import collections
import shutil

import tqdm
import numpy as np
import tensorflow as tf
import PIL.Image

# Assumes the MegaDetector repo root is on the path
from api.batch_processing.postprocessing.postprocess_batch_results import \
    render_bounding_boxes, prepare_html_subpages, PostProcessingOptions

# Make seletion deterministic
random.seed(0)


parser = argparse.ArgumentParser(description='Tools for sampling images, applying a classifier, and generating a html ' + \
                                'summary for 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('--class_list', type=str, help='Simple plain text file containing the list of classes corresponding to the model.')

parser.add_argument('--image_list', type=str, help='Simple plain text file containing a list of image paths.')

# Optional parameters
parser.add_argument('--image_dir', type=str, help='Root directory of the images in the image list.', default=None)
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: 3000', default=3000)
parser.add_argument('--binary_threshold', type=float, help='Switches to two-class mode using the threshold in [0,1] provided')
args = parser.parse_args()

# Validate parameters
assert os.path.isfile(args.frozen_graph)
assert os.path.isfile(args.image_list) and os.path.isfile(args.class_list)
if args.image_dir:
    os.path.isdir(args.image_dir)
os.makedirs(args.output_dir, exist_ok=True)
assert args.num_samples > 0

# Tranfer parameters to post-processing format
options = PostProcessingOptions()
options.output_dir = args.output_dir

# Load frozen graph
print('Loading model...')
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

# Reading image list
with open(args.image_list, 'rt') as fi:
    image_list = fi.read().splitlines()

# Reading class list
with open(args.class_list, 'rt') as fi:
    class_list = fi.read().splitlines()

# Image sampling
random.shuffle(image_list)
selected_images = image_list[:args.num_samples]
    
# Start prediction
with model_graph.as_default():
    
    with tf.Session(config=tf.ConfigProto()) 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])

        # Array for collecting infos for rendering the html
        image_infos = collections.defaultdict(lambda: [])

        for image_path in tqdm.tqdm(selected_images):

            # Read image
            if args.image_dir:
                full_image_path = os.path.join(args.image_dir, image_path)
            else:
                full_image_path = image_path

            image = np.array(PIL.Image.open(full_image_path).convert("RGB"))
            # Scale pixel values to [0,1]
            image = image / 255

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

            if args.binary_threshold:
                cur_class_id = 0 if predictions[0] > args.binary_threshold else 1
                cur_conf = predictions[0]
                conf_name = 'Score'
            else:
                cur_class_id = np.argsort(-predictions)[0]
                cur_conf = predictions[cur_class_id]
                conf_name = 'Conf'

            cur_class_name = class_list[cur_class_id]
            out_image_name = image_path
            display_name = '<b>Class</b>: {}, <b>{}</b>: {:.4%}, <b>Image</b>: {}'.format(
                    cur_class_name, conf_name, cur_conf, out_image_name)

            cur_image_out_dir = os.path.join(args.output_dir, cur_class_name)
            if not os.path.isdir(cur_image_out_dir):
                os.makedirs(cur_image_out_dir, exist_ok=True)
            cur_image_html = render_bounding_boxes(args.image_dir, out_image_name, display_name, [], cur_class_name, options=options)
            image_infos[cur_class_name].append(cur_image_html)

class_counts = prepare_html_subpages(image_infos, args.output_dir)

index_page = """<html>
        <head>
        <style type="text/css">
        <!--
        a { text-decoration:none; }
        body { font-family:segoe ui, calibri, "trebuchet ms", verdana, arial, sans-serif; }
        div.contentdiv { margin-left:20px; }
        -->
        </style>
        </head>
        <body>"""
index_page += """
        <h2>Visualization of classifications</h2>
        <p>A sample of {} images.</p>""".format(sum(class_counts.values()))

if args.binary_threshold:
    index_page += """
            <p>Evaluated a binary classifier at threshold {0:g}%.
            All scores above {0:g}% are classified as '{1}', all below as '{2}'.</p>""".format(
            args.binary_threshold*100, class_list[0], class_list[1])

index_page += """
        <h3>Sample images</h3>
        <div class="contentdiv">"""

# Add links to all available classes
for cname in sorted(class_list):
    assert cname != 'index'
    if cname in class_counts:
        index_page += "<a href='{0}.html'>{0}</a> ({1})<br>".format(
                cname,
                class_counts[cname])

index_page += "</div>"

with open(os.path.join(args.output_dir, 'index.html'), 'wt') as fi:
    fi.write(index_page)