import json
import os
import matplotlib.pyplot as plt

keys = ['action_history', 'noop_count', 'failed']
alphas = ['1.0', '1.5', '2.0', '2.5', '3.0']
levels = ['level_2', 'level_3', 'level_4']
levels = ['level_2']
directories = [ '3agents_07_08', '2agents_07_08']

print('alpha', '\t\t\t', 'success_rate', '\t\t\t',
            'noop_rate', '\t\t\t',
            'action_success_rate' )

for level in levels:
    plt.figure()
    for directory in directories:
        files =  os.listdir(directory)

        datas = []
        for file in files:
            if level in file:
                with open(f'{directory}/{file}', 'r') as f:
                    data = json.load(f)
            
                break
                # datas.append(data)
        title = level
        if directory.startswith('2agents'):
            line_types = [None, None, None]
            color = 'red'
            name = '2agents'
        else:
            # line_types = ['--', '--', '--']
            line_types = [None, None, None]
            color = 'blue'
            name = '3agents'
        print(name)
        
        success_rates = []
        noop_rates = []
        action_success_rates = []
        for alpha in alphas:
            if not alpha in data:
                continue
            success = 0
            total = 0
            for eposide_index in range(3):
                for  action, action_success in zip( data[alpha]['action_history'][eposide_index] ,
                                                   data[alpha]['action_success_history'][eposide_index] ):
                    for ac, aS in zip(action,action_success):
                        if not ac.startswith('noop'):
                            success += int(aS)
                        total += 1

            success_rate = "{:.3f}".format(data[alpha]['success']/(data[alpha]['success'] + data[alpha]['failed']))
            noop_rate = "{:.3f}".format(data[alpha]['noop_count']/(60*int(directory[0])))
            action_success_rate = "{:.3f}".format(success/(total))
            
            success_rates.append(float(success_rate))
            noop_rates.append(float(noop_rate))
            action_success_rates.append(float(action_success_rate))

            print(alpha, '\t\t\t', success_rate, '\t\t\t',
            noop_rate, '\t\t\t',
            action_success_rate )

        

        plt.plot(alphas, success_rates, label=f"success_rates-{name}", color=color, linestyle=line_types[0])
        # plt.plot(alphas, noop_rates, label=f"noop_rates-{directory}", linestyle=line_types[1], color='red')
        # plt.plot(alphas, action_success_rates, label=f"action_success_rates-{directory}", 
                #  color = 'blue', linestyle=line_types[2])
        plt.legend()
        plt.ylim([0, 1.1])
        plt.grid(True)
        plt.title(title)
        plt.xlabel('task intervals')
        plt.xlabel('task intervals')
        # plt.show()

    plt.savefig(title+'.png')
