"""
    Calls output of pretrained tensorflow model to make prediction of top k genres for user inputed album art
"""

import numpy as np
import tensorflow as tf
import argparse
import sys

def create_graph(modelFullPath):
    """Creates a graph from saved GraphDef file and returns a saver."""
    # Creates graph from saved graph_def.pb.
    with tf.gfile.FastGFile(modelFullPath, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def, name='')


def run_inference_on_image(imagePath, labelsFullPath, modelFullPath):
    answer = None

    if not tf.gfile.Exists(imagePath):
        tf.logging.fatal('File does not exist %s', imagePath)
        return answer

    image_data = tf.gfile.FastGFile(imagePath, 'rb').read()

    # Creates graph from saved GraphDef.
    create_graph(modelFullPath)

    with tf.Session() as sess:

        softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
        predictions = sess.run(softmax_tensor,
                               {'DecodeJpeg/contents:0': image_data})
        predictions = np.squeeze(predictions)

        top_k = predictions.argsort()[-5:][::-1]  # Getting top 5 predictions
        f = open(labelsFullPath, 'rb')
        lines = f.readlines()
        labels = np.array([str(w).replace("\n", "") for w in lines])
        for node_id in top_k:
            human_string = labels[node_id]
            score = predictions[node_id]
            print('%s (score = %.5f)' % (human_string, score))

        answer = labels[top_k]
        return answer

def get_results(filename):
    '''
        input: prediction.out file generated by predict.py
        output: pandas DataFrame used to plot top 5 results
    '''

    top_5 = dict()
    with open(filename) as f:
        lines = f.readlines()
        top_prediction = lines[0].split()[0]
        for  line in lines:
            genre = line.split()[0]
            score = float(line.split()[-1].split(')')[0])
            top_5[genre] = score


    result = pd.DataFrame(top_5, index= ['score'])
    result = result.T.sort_values('score', axis=0,
                                ascending=True,
                                kind='mergesort')
    return top_prediction, result

def plot_pred(top_prediction, to_plot, save_as= None):
    # take predicted category and ordered prediction strengths
    # plot and save to file

    fig, ax = plt.subplots(figsize=(7,3))
    plt.title('Model Prediction: {}'.format(top_prediction))
    plt.ylabel('Genre')
    plt.xlabel('Prediction Strength')
    ax.set_yticklabels(to_plot.index, rotation=50, ha='right') #FIX ALIGNMENT ISSUE
    plt.barh(range(0,5), to_plot['score'], color='#6982A0', alpha=0.8)

    if save_as:
        plt.savefig('{}.png'.format(save_as), bbox_inches='tight')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Fit a Text Classifier model and save the results.')
    parser.add_argument('--image', help='A jpg file to make predictions on.')
    parser.add_argument('--labels', help='A txt file with the genre labels, also produced as a tensorflow output.')
    parser.add_argument('--model', help='A pb file with the final layer tensorflow outputs.')
    args = parser.parse_args()

    with open('prediction.out', 'w') as f:
        sys.stdout = f
        run_inference_on_image(args.image, args.labels, args.model)
