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


def plot_csv_with_titles(paths_dict, x_axis, y_axis, subplot_titles, y_limits=(0, 200)):
    # Adjustments for dynamic subplot creation based on the dictionary size
    num_subplots = len(paths_dict)
    fig, axs = plt.subplots(1, num_subplots, figsize=(4*num_subplots, 4))

    # If only one subplot, axs is not an array, so we need to convert it
    if num_subplots == 1:
        axs = [axs]

    for idx, (_, file_list) in enumerate(paths_dict.items()):
        for path_, color in file_list:
            data = np.genfromtxt(path_, delimiter=',', skip_header=0, dtype=float)

            if data.shape[0]> data.shape[1]:
                mean = np.mean(data, axis=1)
                std = np.std(data, axis=1)
            else:
                mean = np.mean(data, axis=0)
                std = np.std(data, axis=0)

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

            axs[idx].plot(x, mean, color=color)
            axs[idx].fill_between(
                x,
                mean - 1.96 * std,
                mean + 1.96 * std,
                color=color,
                alpha=0.5
            )
            axs[idx].set_title(subplot_titles[idx])
            axs[idx].set_xlabel(x_axis)
            axs[idx].set_xlim([0, mean.shape[0]])
            axs[idx].grid(True)
            axs[idx].set_ylim(y_limits)

            # Only label the y-axis for the leftmost plot
            if idx == 0:
                axs[idx].set_ylabel(y_axis)
            else:
                axs[idx].set_ylabel('')

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':

    home_dir = os.path.expanduser('~')
    file_path = os.path.join(home_dir, 'Documents/IntRLResults/CP-Results')

    filenames = [
        'cp-e150r10-bf15-base/cp-ei-random-1_0-15-1690282051_2959082.csv',
        'cp-e150r10-bf15-noshaping/cp-pei-random-0_95-15-1690276946_1944933.csv',
        'cp-e150r10-bf15-noshaping/cp-pei-regular-25_0-15-1690290021_6843266.csv',
        'cp-e150r10-bf15-noshaping/cp-pei-improvement-0_1-15-1690292664_0382216.csv',
        'cp-e150r10-bf15-shaping/cp-ei-random-0_95-15-1690451164_0115042.csv',
        'cp-e150r10-bf15-shaping/cp-ei-regular-25_0-15-1690456185_1115792.csv',
        'cp-e150r10-bf15-shaping/cp-ei-improvement-0_1-15-1690465143_0114875.csv',
        'cp-e150r10-bf15-shaping/cp-pei-random-0_95-15-1690467921_7118568.csv',
        'cp-e150r10-bf15-shaping/cp-pei-regular-25_0-15-1690470012_0117908.csv',
        'cp-e150r10-bf15-shaping/cp-pei-improvement-0_1-15-1690472449_4115295.csv',
    ]
    filepaths = [os.path.join(file_path, filename) for filename in filenames]

    # Demonstrating the adjusted function with subplot titles
    titles = ["Preference", "Shaping", "Combination", "Regular"]

    data_dict_colored = {
        'subplot1': [(filepaths[0], 'C0'), (filepaths[1], 'C1'), (filepaths[2], 'C2'), (filepaths[3], 'C3')],
        'subplot2': [(filepaths[0], 'C0'), (filepaths[4], 'C1'), (filepaths[5], 'C2'), (filepaths[6], 'C3')],
        'subplot3': [(filepaths[0], 'C0'), (filepaths[7], 'C1'), (filepaths[8], 'C2'), (filepaths[9], 'C3')],
        'subplot4': [(filepaths[0], 'C0'), (filepaths[2], 'C6'), (filepaths[5], 'C8'), (filepaths[8], 'C5')],
    }

    plot_csv_with_titles(data_dict_colored, 'Episodes', 'Reward', titles)

    file_path_reacher = os.path.join(home_dir, 'Documents/IntRLResults/RE-Results')
    filenames_reacher = ['base_line/re-ei-random-1_0-5-1694370994_0363934.csv',
                         'shaping/re-ei-random-1_0-10-1694359559_616903.csv',
                         'shaping/re-ei-regular-10_0-5-1694371946_5364418.csv'
                         ]

    filepaths_reacher = [os.path.join(file_path_reacher, filename) for filename in filenames_reacher]

    titles_reacher = ["Shaping"]

    data_dict_reacher = {
        'subplot1': [(filepaths_reacher[0], 'C0'), (filepaths_reacher[1], 'C1'), (filepaths_reacher[2], 'C2')]
    }

    plot_csv_with_titles(data_dict_reacher, 'Episodes', 'Reward', titles_reacher, y_limits=(-150, 50))