import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import web_predict as p
import sys

#def retrieve_image()

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][:-1])
            top_5[genre] = score

    # f = open(filename)
    # lines = f.readlines()
    # top_prediction = lines[0].split('(')[0]
    # for  line in lines:
    #     genre = line.split('(')[0]
    #     score = float(line.split()[-1][:-1])
    #     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.xlabel('Predictive Probability')
    # print 'YOUR GRAPH LABELS\n\n\n'
    # print np.append([0], to_plot.index)
    labels = [item.strip() for item in to_plot.index]
    ax.set_yticklabels(np.append([0], labels), 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__':

    img_name = 'example.jpg'
    img_name_sin_ext = img_name.split('.')[0]
    img_path = 'static/tmp/{}'.format(img_name)
    labels_path = 'model/output_labels.txt'
    model_path = 'model/output_graph.pb'

    output_name = 'static/tmp/prediction{}.out'.format(img_name_sin_ext)
    with open(output_name, 'w') as f:
        sys.stdout = f
        p.run_inference_on_image(img_path, labels_path, model_path)

    top_prediction, results = get_results(output_name)
    plot_name = 'static/tmp/{}_plt'.format(img_name_sin_ext)
    #print plot_name
    plot_pred(top_prediction, results, save_as= plot_name)
