"""
    Generates the hyperparameter sensitivity plot. Takes as an argument the
    directory that contains the summary file (policy_mean_perf.csv) which is
    generated by create_summary.py

"""

import matplotlib
matplotlib.use('Agg')
import numpy as np
import pandas as pd
import os
import argparse
import seaborn as sns
import matplotlib.pyplot as plt

def main():
    
    sns.set(context='paper', style='darkgrid', font_scale=3, rc={'legend.frameon':False,
    'lines.linewidth':6.0})

    parser = argparse.ArgumentParser()
    parser.add_argument('directory', help='Experiments directory.')
    args = parser.parse_args()

    currd = os.getcwd()

    os.chdir(args.directory)

    exists = os.path.isfile("./policy_mean_perf.csv")
    if not exists:
        raise Exception("policy_mean_perf.csv does not exist in this folder.\
        Check the path. Or you may have forgotten to run create_summary.py")


    df = pd.read_csv("policy_mean_perf.csv")

    # Considering only the neural policies here-

    df = df[df.Policy.str.contains('Network')]

    df['Policy_Type'] = 'NN'
    df.loc[df.Policy.str.contains('RecurrentNetwork'), 'Policy_Type'] = 'RNN'

    # Number of configs scales the line plot to be equal on the x-axis.

    nn_configs = len(df[df.Policy_Type == 'NN'])
    rnn_configs = len(df[df.Policy_Type == 'RNN'])

    df["rank"] = df.groupby("Policy_Type")["Mean_Return"].rank(ascending=False)
    df.loc[df.Policy_Type == 'RNN', 'rank'] = np.linspace(1,nn_configs, num=rnn_configs)
    plt.figure(figsize=(16,9))

    c_list = sns.color_palette()
    c_palette = {'RNN' : c_list[0],
     'NN': c_list[1]}

    ax = sns.lineplot(x='rank', y='Mean_Return', hue='Policy_Type', palette=c_palette,
    data=df, markers=True, marker ='.', markersize=12)

    handles, labels = plt.gca().get_legend_handles_labels()
    order = [1,2]
    plt.legend([handles[idx] for idx in order],[labels[idx] for idx in order])

    # Turn off x-ticks
    plt.tick_params(
        axis='x',
        which='both',
        bottom=False,
        top=False,
        labelbottom=False)

    plt.xlabel('hyperparameter setting (best to worst)')
    plt.ylabel('cumulative reward')
    plt.savefig('hparam_sensitivity.pdf',  bbox_inches='tight',pad_inches = 0)

if __name__ == "__main__":
    main()
