src/utils/plot.py

import numpy as np
import matplotlib.pyplot as plt
import os 
from matplotlib.colors import LinearSegmentedColormap
from collections import Counter
import seaborn as sns

def plot_each_feature_alpha(alpha, size, feature_list):
    cumsum = 0
    fig, ax = plt.subplots(2, 4, figsize=(18, 10))
    fig.tight_layout(pad=3.0)
    color = ['purple', 'orange', 'green', 'pink']
    for id, column in enumerate(feature_list):
        bars = np.arange(size[id])
        ax[id//4][id%4].bar(bars, alpha[cumsum:cumsum+size[id]], color=color[:size[id]])
        ax[id//4][id%4].set_xticks(bars, bars)
        ax[id//4][id%4].set_title(column)
        cumsum += size[id]
    plt.legend()
    plt.show()

def plot_feature_alpha(alpha, size, feature_list, title, course_id='dsp_001'):
    cumsum = 0
    fig, ax = plt.subplots(figsize=(16, 8))
    color = ['lightblue', 'dodgerblue', 'blue', 'navy']
    group = {}
    # replace _ with space for every element in feature_list
    feature_list = [x.replace('_', ' ') for x in feature_list]
    for i in range(len(size)):
        for j in range(max(size)):
            if j not in group:
                group[j] = []
            if j < size[i]:
                group[j].append(alpha[cumsum+j])
            else: 
                group[j].append(0)
        cumsum += size[i]
    bar = np.arange(len(feature_list)) 
    width = 0.18
    for i in range(max(size)):
        plt.bar(bar + width*i, group[i], color=color[i], width=width, edgecolor='black', label=i)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    plt.axhline(alpha[-1], color='black', linestyle='--', label='bias')
    plt.xlabel("Metric") 
    plt.ylabel("Value of alpha") 
    plt.title(title)
    plt.grid(axis='y', linestyle='-', alpha=0.6)
    plt.xticks(bar, feature_list, rotation=45) 
    plt.legend()
    title = title.split(' ')
    title = '-'.join(title)
    if not os.path.exists(f'results/features/{course_id}'):
        os.makedirs(f'results/features/{course_id}')
    plt.savefig(f'results/features/{course_id}/{title}alpha.png') 
    plt.show()

# -----------------------------------------------------------
# Plotting functions for raw clickstream environment

def plot_alpha(alpha, dict_event, dict_action, title=""):
    fig, ax = plt.subplots(1, 2, figsize=(10, 4))
    alpha = alpha
    ax[0].bar(list(dict_event.values()), alpha[:len(dict_event)], align='center',
            color=['red' if v < 0 else 'green' for v in alpha[:len(dict_event)]])
    ax[0].set_title(title + 'Event weight')

    ax[1].bar(list(dict_action.values()), alpha[len(dict_event):len(dict_event)+len(dict_action)], align='center',
            color=['red' if v < 0 else 'green' for v in alpha[len(dict_event):]])
    ax[1].set_xticklabels(labels=list(dict_action.values()), rotation = 90)        
    ax[1].set_title(title + 'Action weight')
    plt.show()

def plot_all_weeks_alpha(history, dict_event, dict_action, 
                         unique_week=np.arange(10), path=None, title="", whatif=False, topics=[], schedule=None):
  
    lines = 2
    cols = len(history)//lines + (len(history)%2!=0)
    fig1, ax = plt.subplots(lines, cols, figsize=(5*cols, 4*lines))
    if lines == 1:
        ax = np.array([ax]).flatten()
    for i in range(len(history)):
        alpha = history[i]['alpha']
        if (len(alpha)==0):
            continue
        subtitle = f'Week {i + 1}'
        axis = ax[i // cols, i % cols] if lines > 1 else ax[i]

        if whatif:
            axis.bar(np.arange(len(alpha)), alpha, align='center',
                     color=['red' if v < 0 else 'green' for v in alpha])
            labels = list(['Chapter_'+str(x) for x in topics.keys()])
            labels.extend(['Chapter Complex', 'Values Prob', 'Value Video', 'Is problem', 'Is video'])
            labels.extend(['Action' + str(x) for x in list(dict_action.keys())])
            axis.set_xticks(np.arange(len(alpha)), labels=labels)
            axis.set_xticklabels(labels=labels, rotation = 90)
        else:
            axis.bar(np.arange(len(dict_event)), alpha[:len(dict_event)], align='center',
                    color=['red' if v < 0 else 'green' for v in alpha[:len(dict_event)]], label='Event')
            n_event_action = len(dict_event) + len(dict_action)
            axis.bar(np.arange(len(dict_event), len(dict_event) + len(dict_action)), alpha[len(dict_event):n_event_action], align='center',
                    color=['violet' if v < 0 else 'cyan' for v in alpha[len(dict_event):]], label='Action')
            if (len(alpha) > n_event_action):
                axis.bar(len(alpha), alpha[-1], align='center', color='black', linewidth=2, label='Intercept')
                axis.axhline(y=alpha[-1], color='black', linestyle='--')
        axis.set_title(subtitle)
    fig1.suptitle(title + 'Event & action weight', weight='bold')
    fig1.tight_layout()
    fig2, ax = plt.subplots(lines, cols, figsize=(20, 4*lines))
    if lines == 1:
        ax = np.array([ax]).flatten()
    for i in range(len(history)):
        train_loss = history[i]['train loss']
        if (len(train_loss)==0):
            continue
        subtitle = f'Week {i + 1}'
        axis = ax[i // cols, i % cols] if lines > 1 else ax[i]
        axis.plot(np.arange(len(train_loss)), train_loss)
        axis.set_title(subtitle)
    fig2.suptitle(title + 'Train loss', weight='bold')
    if path is not None:
        if not os.path.exists(path):
            os.makedirs(path)
        save_title = title.split(' ')
        save_title = '-'.join(save_title)
        fig1.savefig(os.path.join(path, save_title + 'alpha.jpg'))
        fig2.savefig(os.path.join(path, save_title +'train_loss.jpg'))
    fig2.tight_layout()
    plt.legend()
    plt.tight_layout()
    plt.show()
    plt.close('all')
    
def plot_all_weeks_reward(history, dict_event, dict_action, 
                          unique_week=np.arange(10), main_title="", path=None, schedule=None):
    """
    Plot reward for all weeks
    return heatmap of reward
    """
    map_label_event = list(dict_event.keys())
    if schedule is not None:
        map_label_event = [schedule.iloc[x]['chapter'] for x in dict_event.keys()]
    
    # Define layout (rows and columns) dynamically based on number of weeks
    lines = 2
    cols = len(history) // lines + (len(history) % 2 != 0)
    
    fig, ax = plt.subplots(lines, cols, figsize=(8 * cols, 10 * lines))
    if lines == 1:
        ax = np.array([ax]).flatten()
    
    # Custom colormap from red (negative) to green (positive)
    colors = ['#fa8072', 'white', '#00a86b']
    cmap = LinearSegmentedColormap.from_list('custom', list(zip(np.linspace(0, 1, len(colors)), colors)))
    
    # Plot each week's reward
    for i in range(len(history)):
        reward = history[i]['reward']
        if (len(reward)==0):
            continue
        title = f'Week {i + 1}'
        axis = ax[i // cols, i % cols] if lines > 1 else ax[i]
        if schedule is not None:
            sorted_indices = np.argsort(map_label_event)
            sorted_map_event = np.sort(map_label_event)
            sorted_reward = reward[sorted_indices]
            im = axis.imshow(sorted_reward, cmap=cmap, aspect='auto', interpolation='nearest')
            _, unique_indices = np.unique(sorted_map_event, return_index=True)
        
            axis.set_yticks(np.arange(len(sorted_map_event)))
            axis.set_yticklabels(['W' + str(sorted_map_event[x]) if x in unique_indices else None
                                  for x in range(len(sorted_map_event))], fontsize=23)
            axis.set_ylabel('Chapters in Week', fontsize=23)
            # axis.set_ylabel('', fontsize=16)
        else:
            im = axis.imshow(reward, cmap=cmap, aspect='auto', interpolation='nearest', vmin=-1, vmax=1)
       
        
        # axis.set_title(title + 'Reward')
        axis.set_xticks(np.arange(len(list(dict_action.values()))), labels=list(dict_action.values()))
        axis.set_xticklabels(labels=list(dict_action.values()), rotation = 90, fontsize=23)
        axis.set_title(title, fontsize=30)
        
#         # Set the title for each subplot
#         axis.set_title(title, fontsize=18)
        
        # Add colorbar for visualizing reward scale
        fig.colorbar(im, ax=axis, orientation='vertical', fraction=0.04, pad=0.04)
        # fig.subplots_adjust(right=0.8)
        # cbar_ax = fig.add_axes([1, 0.2, 0.02, 0.7]) 
        # cbar = fig.colorbar(im, cax=cbar_ax)
        # Adding gridlines for better readability
        # axis.grid(True, which='both', axis='both', color='black', linestyle='--', linewidth=0.5)
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])  # Adjust spacing for the title
    if path is not None:
        if not os.path.exists(path):
            os.makedirs(path)
        save_title = main_title.split(' ')
        save_title = '-'.join(save_title)
        fig.savefig(os.path.join(path, save_title+'reward.jpg'), dpi=500)
        fig.savefig(os.path.join(path, save_title+'reward.svg'), dpi=500)
    # Optionally, save the plot to a file
    # if path is not None:
    #     plt.savefig(path, dpi=500)
    
    plt.show()

# def plot_all_weeks_reward(history, dict_event, dict_action, 
#                           unique_week=np.arange(10), main_title="", path=None, schedule=None):
#     """
#     Plot reward for all weeks
#     return heatmap of reward
#     """
#     map_label_event = dict_event.keys()
#     if schedule is not None:
#         map_label_event = [schedule.iloc[x]['chapter'] for x in dict_event.keys()]
#     lines = 2
#     cols = len(history)//lines + (len(history)%2!=0)
#     fig, ax = plt.subplots(lines, cols, figsize=(6*cols, 10*lines))
#     if lines == 1:
#         ax = np.array([ax]).flatten()
#     colors = ['red', 'white', 'green']
#     cmap = LinearSegmentedColormap.from_list('custom', list(zip(np.linspace(0, 1, len(colors)), colors)))
#     for i in range(len(history)):
#         reward = history[i]['reward']
#         if (len(reward)==0):
#             continue
#         title = f'Week {i + 1}'
#         axis = ax[i // cols, i % cols] if lines > 1 else ax[i]
#         if schedule is not None:
#             sorted_indices = np.argsort(map_label_event)
#             sorted_map_event = np.sort(map_label_event)
#             sorted_reward = reward[sorted_indices]
#             im = axis.imshow(sorted_reward, cmap=cmap, aspect='auto', interpolation='nearest', vmin=-1, vmax=1)
#             _, unique_indices = np.unique(sorted_map_event, return_index=True)
        
#             axis.set_yticks(np.arange(len(sorted_map_event)))
#             axis.set_yticklabels(['W' + str(sorted_map_event[x]) if x in unique_indices else None
#                                   for x in range(len(sorted_map_event))], fontsize=23)
#         else:
#             im = axis.imshow(reward, cmap=cmap, aspect='auto', interpolation='nearest', vmin=-1, vmax=1)
#         axis.set_title(title + 'Reward')
#         axis.set_xticks(np.arange(len(list(dict_action.values()))), labels=list(dict_action.values()))
#         axis.set_xticklabels(labels=list(dict_action.values()), rotation = 90, fontsize=23)
#         axis.set_title(title, fontsize=30)
#     fig.subplots_adjust(right=0.8)
#     cbar_ax = fig.add_axes([1, 0.2, 0.02, 0.7]) 
#     cbar = fig.colorbar(im, cax=cbar_ax)
#     fig.tight_layout(pad=3.0)
#     plt.xlabel("Student's Action", fontsize=30)
#     plt.ylabel('Content in Week "{x}"', fontsize=30)
#     plt.tight_layout()
#     plt.legend()
    
#     if path is not None:
#         if not os.path.exists(path):
#             os.makedirs(path)
#         save_title = main_title.split(' ')
#         save_title = '-'.join(save_title)
#         fig.savefig(os.path.join(path, save_title+'reward.jpg'), dpi=500)
#     plt.show()
#     plt.close()

def plot_problem_event_alpha(history, dict_event, dict_action, problem_event, unique_week=np.arange(10)):
    cols=5
    lines = len(history)//cols + (1 if len(history)%cols != 0 else 0)
    cols = min(cols, len(history))
    fig, ax = plt.subplots(lines, cols, figsize=(10, 8))
    if lines == 1:
        ax = np.array([ax]).flatten()
    fig.tight_layout(pad=3.0)
    for i in range(len(history)):
        color = ['blue' if dict_event[i] in problem_event else 'orange' for i in range(len(dict_event))]
        alpha = history[i]['alpha']
        if len(alpha) == 0:
            continue
        title = f'Week {unique_week[i] + 1}'
        axis = ax[i // cols, i % cols] if lines > 1 else ax[i]
        axis.bar(list(dict_event.keys()), alpha[:len(dict_event)], align='center', color=color)
        axis.set_title(title)
    plt.show()

def plot_problem_event_reward(history, map_problem_id, map_video_id, dict_action):
    colors = ['red', 'white', 'green']
    _, ax = plt.subplots(1, 2, figsize=(16, 5))
    cmap = LinearSegmentedColormap.from_list('custom', list(zip(np.linspace(0, 1, len(colors)), colors)))
    reward = history['reward']
    print(reward.shape)
    im1 = ax[0].imshow(reward[map_problem_id, :], cmap=cmap, aspect='auto', interpolation='nearest', vmin=-1, vmax=1)
    im2 = ax[1].imshow(reward[map_video_id, :], cmap=cmap, aspect='auto', interpolation='nearest', vmin=-1, vmax=1)   
    ax[0].set_title('Problem Reward')
    ax[0].set_xticks(np.arange(len(list(dict_action.values()))), labels=list(dict_action.values()))
    ax[0].set_xticklabels(labels=list(dict_action.values()), rotation = 90)
    ax[1].set_title('Video Reward')
    ax[1].set_xticks(np.arange(len(list(dict_action.values()))), labels=list(dict_action.values()))
    ax[1].set_xticklabels(labels=list(dict_action.values()), rotation = 90)
    plt.tight_layout()
    plt.show()
    
def distribution_skillset(trajectories_each_week, world, columns=['topic']):
    fig, ax = plt.subplots(2, 5, figsize=(20, 10))
    for i, trajectories in enumerate(trajectories_each_week):
        data = []
        for trajectory in trajectories:
            for state, action, _ in trajectory:
                topic = world.get_features(state, action)['topic']
                data.extend(topic)
        word_counts = Counter(data)
        print(word_counts)
        unique_words = list(word_counts.keys())
        word_counts_values = list(word_counts.values())
        ax[i//5, i%5].bar(unique_words, word_counts_values, color='skyblue')
        ax[i//5, i%5].set_title(f'Week {i + 1}')
    plt.xticks(rotation=90)  
    plt.grid(axis='y') 
    plt.show()

def plot_accuracy_synthesize(accs_training, accs_full_synthesize, accs_personalized_syn, weeks):
    plt.figure(figsize=(12, 6))
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    sns.lineplot(data=accs_training, x='Week', y='Accuracy', label='Prediction Model')
    sns.lineplot(data=accs_full_synthesize, x='Week', y='Accuracy', label='Synthesized Classrooms')
    sns.lineplot(data=accs_personalized_syn, x='Week', y='Accuracy', label='Personalized Trajectories')
    plt.xlabel('Number of Weeks')
    plt.xticks(weeks)
    plt.ylabel('Accuracy')
    plt.legend()
    plt.grid(axis='y')
    plt.show()

def plot_a_week_whatif(week_data, week, course):
    plt.figure(figsize=(12, 8))
    sns.barplot(
        x='Chapter', 
        y='Mean Average Impact', 
        data=week_data, 
        errorbar="sd",
        legend=False
    )
    plt.grid(axis='y', linestyle='-', alpha=0.6)
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    plt.xlabel("Topic")
    plt.ylabel('Mean Average Effectiveness')
    plt.title(f'What-if Classroom in Week {week}')
    plt.savefig(f'results/whatif/{course}/whatif-analysis/whatif_results_week_{week}.png')
    plt.show()
    plt.close()

def plot_heatmap_whatif(whatif_results, course):
    heatmap_data = whatif_results.pivot_table(
        index='Week',
        columns='Chapter',
        values='SD Average Impact',
        aggfunc='mean'  
    )
    plt.figure(figsize=(15, 8))
    sns.heatmap(heatmap_data, annot=False, cmap='Blues', linewidths=.5)
    plt.ylabel('Week')
    plt.xlabel("Week")
    plt.savefig(f'results/whatif/{course}/whatif-analysis/whatif_results_heatmap.png')
    plt.show()