import numpy as np
import matplotlib.pyplot as plt
import os


def plot_csv(paths, x_axis, y_axis):
    for path_ in paths:
        data = np.genfromtxt(path_, delimiter=',', skip_header=0, dtype=float)

        mean = np.mean(data, axis=0)
        std = np.std(data, axis=0)

        x = np.linspace(0, mean.shape[0], mean.shape[0])

        # Extract the first part of the filename and use it as a label
        label = os.path.basename(path_).split('-')[0:5]
        label = f"{label[1]}," \
                f" {label[2]}," \
                f"  {label[3].replace('_', '.') if label[3] != '' else 'base'}," \
                f" nrbfs = {int(label[4])}"

        plt.plot(x, mean, label=label)
        plt.fill_between(
            x,
            mean - 1.96 * std,
            mean + 1.96 * std,
            alpha=0.5
        )
    plt.xlabel(x_axis)
    plt.xlim([0, mean.shape[0]])
    plt.ylabel(y_axis)
    plt.grid(True)
    plt.legend(loc="upper left")
    plt.show()


if __name__ == '__main__':
    filenames = ['franka-pei-random-1_0-6-1694787936_385925.csv',
                 ]
    home_dir = os.path.expanduser('~')
    file_path = os.path.join(home_dir, 'Documents/IntRLResults/Franka-Results')
    paths = [os.path.join(file_path, filename) for filename in filenames]
    plot_csv(paths, 'Episodes', 'Reward')
    #
    # filenames2 = ['mc-ei-random-1_0-10-1685708208_6402516.csv',
    #               'mc-pi-random-1_0-10-1685709168_7113624.csv',
    #               'mc-cb-random-1_0-10-1685714487_4603446.csv']
    #
    # paths2 = [os.path.join(file_path, filename) for filename in filenames2]
    # plot_csv(paths2, 'Episodes', 'Reward')