import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

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['groups'], rotation=50, ha='right')
    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__':
    top_prediction, results = get_results('prediction.out')
    plot_pred(top_prediction, results, save_as= 'plt_name')
