import os
import matplotlib.pyplot as plt
import numpy as np
import json
import seaborn as sns; sns.set()
import glob2
import argparse


def smooth_reward_curve(x, y, padding, a_range=20000):
    if a_range != 20000:
        len_m = a_range
    else:
        len_m = len(x)
    halfwidth = int(np.ceil(len_m / 100))
    k = halfwidth
    xsmoo = x
    ysmoo = np.convolve(y, np.ones(padding * k + 1), mode='same') / np.convolve(np.ones_like(y), np.ones(padding * k + 1),
        mode='same')
    return xsmoo, ysmoo


def load_results(file):
    if not os.path.exists(file):
        return None
    with open(file, 'r') as f:
        lines = [line for line in f]
    if len(lines) < 2:
        return None
    keys = [name.strip() for name in lines[0].split(',')]
    data = np.genfromtxt(file, delimiter=',', skip_header=1, filling_values=0.)
    if data.ndim == 1:
        data = data.reshape(1, -1)
    assert data.ndim == 2
    assert data.shape[-1] == len(keys)
    result = {}
    for idx, key in enumerate(keys):
        result[key] = data[:, idx]
    return result


def pad(xs, value=np.nan):
    maxlen = np.max([len(x) for x in xs])
    max_index = 0
    for index, x in enumerate(xs):
        if len(x) == maxlen:
            max_index = index
    
    padded_xs = []
    index=-1
    for x in xs:
        if x.shape[0] >= maxlen:
            padded_xs.append(x)
            index += 1
        else:
            padding = np.ones((maxlen - x.shape[0],) + x.shape[1:]) * value
            x_padded = np.concatenate([x, padding], axis=0)
            assert x_padded.shape[1:] == x.shape[1:]
            assert x_padded.shape[0] == maxlen
            padded_xs.append(x_padded)
    return np.array(padded_xs), index, max_index


parser = argparse.ArgumentParser()
parser.add_argument('--dir', type=str, default='')
parser.add_argument('--title', type=str, default='')
parser.add_argument('--type', type=str, default='test')
parser.add_argument('--dir1', type=str, default=0)
parser.add_argument('--dir2', type=str, default=0)
parser.add_argument('--dir3', type=str, default=0)
parser.add_argument('--dir4', type=str, default=0)
parser.add_argument('--dir5', type=str, default=0)
parser.add_argument('--dir6', type=str, default=0)
parser.add_argument('--dir7', type=str, default=0)
parser.add_argument('--padding', type=int, default=10)
parser.add_argument('--plot_name', type=str, default='')
parser.add_argument('--smooth', type=int, default=1)
parser.add_argument('--range', type=int, default=20000)

args = parser.parse_args()
path_list = []

path_list_dict = dict()
dir_list_dict = dict()

if args.dir1:
    dir_list_dict['dir1'] = args.dir1
if args.dir2:
    dir_list_dict['dir2'] = args.dir2
if args.dir2:
    dir_list_dict['dir3'] = args.dir2
if args.dir2:
    dir_list_dict['dir4'] = args.dir2
if args.dir2:
    dir_list_dict['dir5'] = args.dir2
if args.dir2:
    dir_list_dict['dir6'] = args.dir2
if args.dir2:
    dir_list_dict['dir7'] = args.dir2

# Load all data.
data = {}
for i in range(len(dir_list_dict)):
    i += 1
    dir_list_dict['dir'+str(i)] = './models' + args.dir + '/' + dir_list_dict['dir'+str(i)]

print(dir_list_dict)
        
for i in range(len(dir_list_dict)):
    curr_path = dir_list_dict['dir'+str(i+1)].split(':')[0]
    print(curr_path)
    if not os.path.isdir(curr_path):
        print("Not os path is current path")
        continue
    results = load_results(os.path.join(curr_path, 'progress.csv'))
    if not results:
        print('skipping {}'.format(curr_path))
        continue
    print('loading {} ({})'.format(curr_path, len(results['1.  Epoch'])))
    with open(os.path.join(curr_path, 'params.json'), 'r') as f:
        params = json.load(f) #load the parameters from the param file
    
    if args.type == 'train':
        success_rate = np.array(results['5.  Train success rate'])
    else:
        success_rate = np.array(results['9.  Test success rate'])
    epoch = np.array(results['2.  Timesteps']) + 1
    env_id = params['env_name']
    replay_strategy = params['replay_strategy']

    iden = dir_list_dict['dir'+str(i+1)].split(':')[-1]
    env_id += iden
    config = iden

    # Process and smooth data.
    assert success_rate.shape == epoch.shape
    x = epoch
    y = success_rate
    if args.smooth:
        x, y = smooth_reward_curve(epoch, success_rate, args.padding, args.range) #smoothen the data
    assert x.shape == y.shape

    z=y
    if env_id not in data:
        data[env_id] = {}
    if config not in data[env_id]:
        data[env_id][config] = []
    data[env_id][config].append((x, y, z))

# Plot data.
plt.clf()
plt.figure(1)
plt.xlabel('Epoch')
plt.ylabel('Success Rate')
plt.title(args.title, fontsize=12)
plt.legend()
i = -1
colours = ['black', 'r', 'deepskyblue', 'limegreen', 'g', 'hotpink', 'saddlebrown', 'blue', 'gray']
for env_id in sorted(data.keys()): # for all the env ids
    i += 1

    for config in sorted(data[env_id].keys()): #for all the configs in env ids
        xs, ys, zs = zip(*data[env_id][config])
        xs, _, index = pad(xs)
        ys,_,_ = pad(ys)
        zs,_,_ = pad(zs)
        assert xs.shape == ys.shape
        assert xs.shape == zs.shape
        # if args.range:
        xs = xs[:,:args.range]
        ys = ys[:,:args.range]
        zs = zs[:,:args.range]
        
        plt.plot(xs[index], np.nanmedian(ys, axis=0), color=colours[i], label=config)
        plt.fill_between(xs[index], np.nanpercentile(ys, 25, axis=0), np.nanpercentile(ys, 75, axis=0), alpha=0.25)
        plt.xlabel('Timesteps', fontsize=12)
        plt.ylabel('Success Rate', fontsize=12)
        current_values = plt.gca().get_xticks()
        plt.gca().set_xticklabels(['{:.1f}'.format(x/1000000.)+'M' if x >= 1000000 else '{:.0f}'.format(x/1000)+'k' if x >= 1000 else '{:.0f}'.format(x) for x in current_values])

        plt.legend(fontsize=11)

plt.grid(b=None)
plt.subplots_adjust(top=0.92, bottom=0.1, left=0.15, right=0.95, hspace=0.3, wspace=0.35)

temp_name = '_success_rate'
plt.savefig(os.path.join('./figs/', '{}.png'.format(args.plot_name+ temp_name)))
plt.show()
        


    
    


    
