utils.py

import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.colors import LinearSegmentedColormap
from matplotlib import cm
from matplotlib.widgets import Slider
import csv
from datetime import datetime
from config import *
import os

def draw_q_values_v3(ax, q_values, grid_size, vmax=None, fig=None, cbar=None, center_zero=False, plot_type="Q+"):
    ax.clear()
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    ax.set_aspect('equal')
    ax.invert_yaxis()
    ax.axis('off')
    
    # Normalize values from 0 to 1
    vmin = 0
    vmax = 1  # Normalize to a fixed range from 0 to 1

    # Apply different colormaps based on plot_type (Q+ or Q-)
    if plot_type == "Q+":
        # Q+ should transition from white to blue
        cmap = LinearSegmentedColormap.from_list("white_blue", ["#FFFFFF", "#0000FF"], N=256)
    elif plot_type == "Q-":
        # Q- should transition from white to red
        cmap = LinearSegmentedColormap.from_list("white_red", ["#FFFFFF", "#FF0000"], N=256)

    norm = plt.Normalize(vmin=vmin, vmax=vmax)

    directions = {
        0: (0.3, 0.05),   # Up
        1: (0.3, 0.55),   # Down
        2: (0.05, 0.3),   # Left
        3: (0.55, 0.3),   # Right
    }
    size = 0.4

    for s in range(grid_size * grid_size):
        row, col = divmod(s, grid_size)
        center_x, center_y = col + 0.5, row + 0.5
        circle = patches.Circle((center_x, center_y), 0.05, facecolor='black', edgecolor='none')
        ax.add_patch(circle)

        for a in range(4):
            if (a == 0 and row == 0) or (a == 1 and row == grid_size - 1) or \
               (a == 2 and col == 0) or (a == 3 and col == grid_size - 1):
                continue

            dx, dy = directions[a]
            color = cmap(norm(q_values[s, a]))  # Normalize the values for color mapping
            square = patches.Rectangle((col + dx, row + dy), size, size, facecolor=color, edgecolor='black', linewidth=0.2)
            ax.add_patch(square)

    # Create/update colorbar only if needed
    if cbar is None and fig is not None:
        colorbar_width = 0.75
        cbar_ax = fig.add_axes([(1 - colorbar_width) / 2, 0.07, colorbar_width, 0.02])

        # Clear the old colorbar if it exists
        if cbar is not None:
            cbar.remove()

        # Create a new colorbar
        cbar = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=cbar_ax, orientation='horizontal')
        
        # Set tick values and labels
        tick_values = np.linspace(vmin, vmax, num=5)
        cbar.set_ticks(tick_values)
        cbar.ax.set_xticklabels([f"{v:.2f}" for v in tick_values], fontsize=10)

    # Update the plot title based on Q+ or Q-
    ax.set_title(f"{plot_type} Values (Up, Down, Left, Right)")

    return cbar

def update_all_live_plots_v3(agent, rewards, plots, grid_size):
    # Update cumulative reward line
    line = plots["reward_line"]
    line.set_data(range(len(rewards)), rewards)

    ax = plots["reward_ax"]
    ax.set_xlim(0, len(rewards))
    ax.set_ylim(min(rewards) - 5, max(rewards) + 5)

    # For Q+ values
    draw_q_values_v3(plots["qplus_ax"], agent.q_plus, grid_size, fig=plots["qplus_fig"], plot_type="Q+")
    plots["qplus_ax"].set_title("Q+ Values")

    # For Q- values
    draw_q_values_v3(plots["qminus_ax"], agent.q_minus, grid_size, fig=plots["qminus_fig"], plot_type="Q-")
    plots["qminus_ax"].set_title("Q- Values")

    # Update Q-Net (Q+ - Q−) plot
    plots["qnet_cbar"] = draw_q_net_difference(plots["qnet_ax"], agent.q_plus, agent.q_minus, grid_size, fig=plots["qnet_fig"], cbar=plots["qnet_cbar"])
    plots["qnet_ax"].set_title("Q-Net Difference (Q⁺ - Q⁻)")

    # Update policy arrows
    draw_policy_arrows(plots["policy_ax"], agent.q_plus, agent.q_minus, grid_size)
    plots["policy_ax"].set_title("Policy Arrows (Derived from Q⁺ - Q⁻)")

    # Refresh all plots
    for fig in [plots["reward_fig"], plots["qnet_fig"], plots["policy_fig"], plots["qplus_fig"], plots["qminus_fig"]]:
        fig.canvas.draw()
        fig.canvas.flush_events()

    plt.pause(0.00001)

def draw_q_values(ax, q_values, grid_size, vmax=None, fig=None, cbar=None, center_zero=False):
    ax.clear()
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    ax.set_aspect('equal')
    ax.invert_yaxis()
    ax.axis('off')
    # Normalize values
    if vmax is None:
        vmax = np.max(np.abs(q_values)) + 1e-8
    if center_zero:
        cmap = LinearSegmentedColormap.from_list("red_white_blue", ["#FF0000", "#FFFFFF", "#0000FF"], N=21)
        vmin = -vmax
    else:
        cmap = plt.cm.viridis
        vmin = 0
    norm = plt.Normalize(vmin=vmin, vmax=vmax)
    directions = {
        0: (0.3, 0.05),   # Up
        1: (0.3, 0.55),   # Down
        2: (0.05, 0.3),   # Left
        3: (0.55, 0.3),   # Right
    }
    size = 0.4

    for s in range(grid_size * grid_size):
        row, col = divmod(s, grid_size)
        center_x, center_y = col + 0.5, row + 0.5
        circle = patches.Circle((center_x, center_y), 0.05, facecolor='black', edgecolor='none')
        ax.add_patch(circle)

        for a in range(4):
            if (a == 0 and row == 0) or (a == 1 and row == grid_size - 1) or \
               (a == 2 and col == 0) or (a == 3 and col == grid_size - 1):
                continue

            dx, dy = directions[a]
            color = cmap(norm(q_values[s, a]))
            square = patches.Rectangle((col + dx, row + dy), size, size, facecolor=color, edgecolor='black', linewidth=0.2)
            ax.add_patch(square)

    if cbar is None and fig is not None:
        colorbar_width = 0.75
        cbar_ax = fig.add_axes([(1 - colorbar_width) / 2, 0.07, colorbar_width, 0.02])
        cbar = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=cbar_ax, orientation='horizontal')
        tick_values = np.linspace(vmin, vmax, num=5)
        cbar.set_ticks(tick_values)
        cbar.ax.set_xticklabels([f"{v:.2f}" for v in tick_values], fontsize=10)

    return cbar

def init_all_live_plots(grid_size):
    plt.ion()

    # Create a single figure with a grid layout for subplots
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    fig.tight_layout(pad=4.0)

    # Unpack axes
    (ax1, ax2, ax3), (ax4, ax5, ax6) = axes

    # Reward plot
    line_reward, = ax1.plot([], [], label="Cumulative Reward")
    ax1.set_title("Cumulative Reward")
    ax1.set_xlabel("Step")
    ax1.set_ylabel("Reward")
    ax1.grid(True)
    ax1.legend()

    # Q+ Values
    step_text = ax2.text(0.5, -0.5, "", fontsize=12, color="black")
    ax2.set_xlim(0, grid_size)
    ax2.set_ylim(0, grid_size)
    ax2.set_aspect('equal')
    ax2.invert_yaxis()
    ax2.axis('off')

    # Q− Values
    ax3.set_xlim(0, grid_size)
    ax3.set_ylim(0, grid_size)
    ax3.set_aspect('equal')
    ax3.invert_yaxis()
    ax3.axis('off')

    # Q-Net Difference
    ax4.set_xlim(0, grid_size)
    ax4.set_ylim(0, grid_size)
    ax4.set_aspect('equal')
    ax4.invert_yaxis()
    ax4.axis('off')

    # Policy Arrows
    ax5.set_xlim(0, grid_size)
    ax5.set_ylim(0, grid_size)
    ax5.set_aspect('equal')
    ax5.invert_yaxis()
    ax5.axis('off')

    # Empty slot for now (ax6)
    ax6.axis('off')  # Or use it for reward rate later

    return {
        "fig": fig,
        "reward_ax": ax1, "reward_line": line_reward,
        "qplus_ax": ax2, "qplus_text": step_text,
        "qminus_ax": ax3,
        "qnet_ax": ax4,
        "policy_ax": ax5,
        'qnet_cbar': None
    }

def update_all_live_plots(agent, rewards, plots, grid_size):
    # Update cumulative reward line
    line = plots["reward_line"]
    line.set_data(range(len(rewards)), rewards)

    reward_ax = plots["reward_ax"]
    reward_ax.set_xlim(0, len(rewards))
    reward_ax.set_ylim(min(rewards) - 5, max(rewards) + 5)
    reward_ax.set_title("Live Cumulative Reward")

    # Update Q+ detailed grid
    draw_q_values(plots["qplus_ax"], agent.q_plus, grid_size, fig=plots["fig"])
    plots["qplus_ax"].set_title("Q+ Values (Up, Down, Left, Right)")

    # Update Q− detailed grid
    draw_q_values(plots["qminus_ax"], agent.q_minus, grid_size, fig=plots["fig"])
    plots["qminus_ax"].set_title("Q- Values (Up, Down, Left, Right)")

    # Update Q-Net (Q+ - Q−) plot
    plots["qnet_cbar"] = draw_q_net_difference(
        plots["qnet_ax"], agent.q_plus, agent.q_minus, grid_size, fig=plots["fig"], cbar=plots["qnet_cbar"]
    )
    plots["qnet_ax"].set_title("Q-Net Difference (Q⁺ - Q⁻)")

    # Update policy arrows
    draw_policy_arrows(plots["policy_ax"], agent.q_plus, agent.q_minus, grid_size)
    plots["policy_ax"].set_title("Policy Arrows (Derived from Q⁺ - Q⁻)")

    # Refresh the unified figure
    plots["fig"].canvas.draw()
    plots["fig"].canvas.flush_events()
    plt.pause(0.00001)

def init_all_live_plots_v1(grid_size):
    plt.ion()

    # Reward plot
    fig1, ax1 = plt.subplots(figsize=(6, 4))
    line_reward, = ax1.plot([], [], label="Cumulative Reward")
    ax1.set_title("Live Cumulative Reward")
    ax1.set_xlabel("Step")
    ax1.set_ylabel("Reward")
    ax1.grid(True)
    ax1.legend()

    # Q+ detailed plot
    fig2, ax2 = plt.subplots(figsize=(8, 8))
    ax2.set_xlim(0, grid_size)
    ax2.set_ylim(0, grid_size)
    ax2.set_aspect('equal')
    ax2.invert_yaxis()
    ax2.axis('off')
    step_text = ax2.text(0.5, -0.5, "", fontsize=12, color="black")

    # Q− detailed plot
    fig3, ax3 = plt.subplots(figsize=(8, 8))
    ax3.set_xlim(0, grid_size)
    ax3.set_ylim(0, grid_size)
    ax3.set_aspect('equal')
    ax3.invert_yaxis()
    ax3.axis('off')

    # Q-Net difference plot (Q+ - Q−)
    fig4, ax4 = plt.subplots(figsize=(8, 8))
    ax4.set_xlim(0, grid_size)
    ax4.set_ylim(0, grid_size)
    ax4.set_aspect('equal')
    ax4.invert_yaxis()
    ax4.axis('off')

    # Reward rate plot
    # fig_rate, ax_rate = plt.subplots(figsize=(6, 4))
    # line_rate, = ax_rate.plot([], [], label=f"Reward Rate (window={WINDOW_SIZE})", color='orange')
    # ax_rate.set_title("Live Reward Rate")
    # ax_rate.set_xlabel("Step")
    # ax_rate.set_ylabel("Reward Rate")
    # ax_rate.set_ylim(0, 1)  # reward rate ranges from 0 to 1
    # ax_rate.grid(True)
    # ax_rate.legend()

    # Policy arrows plot
    fig_policy, ax_policy = plt.subplots(figsize=(8, 8))
    ax_policy.set_xlim(0, grid_size)
    ax_policy.set_ylim(0, grid_size)
    ax_policy.set_aspect('equal')
    ax_policy.invert_yaxis()
    ax_policy.axis('off')
    
    return {
        "reward_fig": fig1, "reward_ax": ax1, "reward_line": line_reward,
        "qplus_fig": fig2, "qplus_ax": ax2, "qplus_text": step_text,
        "qminus_fig": fig3, "qminus_ax": ax3,
        # "reward_rate_fig": fig_rate, "reward_rate_ax": ax_rate, "reward_rate_line": line_rate,
        "qnet_fig": fig4, "qnet_ax": ax4,
        "policy_fig": fig_policy, "policy_ax": ax_policy,

        'qnet_cbar': None
    }

def add_episode_background(ax, total_steps, episode_length=UPDATE_SCHEDULE):
    [child.remove() for child in ax.get_children() if isinstance(child, patches.Polygon)]
    num_episodes = (total_steps // episode_length) + 1
    for i in range(num_episodes):
        start = i * episode_length
        end = (i + 1) * episode_length
        color = 'lightgray' if i % 2 == 0 else 'white'
        ax.axvspan(start, end, facecolor=color, alpha=0.3, zorder=0)

def update_all_live_plots_v1(agent, rewards, plots, grid_size):
    # Update cumulative reward line
    line = plots["reward_line"]
    line.set_data(range(len(rewards)), rewards)

    ax = plots["reward_ax"]
    ax.set_xlim(0, len(rewards))
    ax.set_ylim(min(rewards) - 5, max(rewards) + 5)
    # add_episode_background(ax, len(rewards), episode_length=100)

    # reward rate plot
    # reward_events = [rewards[0]] + \
    #             [rewards[i] - rewards[i-1] for i in range(1, len(rewards))]

    # reward_rate = []
    # for i in range(len(reward_events)):
    #     if i < WINDOW_SIZE:
    #         reward_rate.append(np.nan)  # not enough data yet
    #     else:
    #         window_sum = sum(reward_events[i-WINDOW_SIZE:i])
    #         reward_rate.append(window_sum / WINDOW_SIZE)

    # # Update reward rate plot
    # rate_line = plots["reward_rate_line"]
    # rate_line.set_data(range(len(reward_rate)), reward_rate)
    # ax_rate = plots["reward_rate_ax"]
    # ax_rate.set_xlim(0, len(reward_rate))
    # ax_rate.set_ylim(0, 1)  # reward rate from 0 to 1
    # add_episode_background(ax_rate, len(reward_rate), episode_length=100)
    # Update Q+ detailed grid
    draw_q_values(plots["qplus_ax"], agent.q_plus, grid_size)
    plots["qplus_ax"].set_title("Q+ Values (Up, Down, Left, Right)")
    # Update Q− detailed grid
    draw_q_values(plots["qminus_ax"], agent.q_minus, grid_size)
    plots["qminus_ax"].set_title("Q- Values (Up, Down, Left, Right)")
    # max_val = max(np.max(np.abs(agent.q_plus)), np.max(np.abs(agent.q_minus)))
    # plots["qplus_cbar"] = draw_q_values(
    #     plots["qplus_ax"], agent.q_plus, grid_size, vmax=max_val, fig=plots["qplus_fig"], center_zero=True
    # )
    # plots["qminus_cbar"] = draw_q_values(
    #     plots["qminus_ax"], agent.q_minus, grid_size, vmax=max_val, fig=plots["qminus_fig"], center_zero=True
    # )


    # Update Q-Net (Q+ - Q−) plot
    plots["qnet_cbar"] = draw_q_net_difference(plots["qnet_ax"], agent.q_plus, agent.q_minus, grid_size, fig=plots["qnet_fig"], cbar=plots["qnet_cbar"])
    plots["qnet_ax"].set_title("Q-Net Difference (Q⁺ - Q⁻)")
    
    draw_policy_arrows(plots["policy_ax"], agent.q_plus, agent.q_minus, grid_size)
    plots["policy_ax"].set_title("Policy Arrows (Derived from Q⁺ - Q⁻)")

    # Refresh all plots
    # for fig in [plots["qplus_fig"], plots["qminus_fig"],  plots["qnet_fig"]]: 
    #     fig.canvas.draw()
    #     fig.canvas.flush_events()
    for fig in [plots["reward_fig"], plots["qnet_fig"], plots["policy_fig"],plots["qplus_fig"], plots["qminus_fig"]]:
        fig.canvas.draw()
        fig.canvas.flush_events()
    # plots["fig"].canvas.draw()
    # plots["fig"].canvas.flush_events()
    plt.pause(0.00001)

def plot_reward_rate_with_slider(rewards):
    initial_window = 10
    max_window = max(1, len(rewards)//2)
    fig, ax = plt.subplots(figsize=(8, 5))
    plt.subplots_adjust(bottom=0.25)
    ax.set_title("Reward Rate with Adjustable Window Size")
    ax.set_xlabel("Step")
    ax.set_ylabel("Reward Rate")
    ax.set_xlim(0, len(rewards))
    ax.set_ylim(0, 1)
    ax.grid(True)
    reward_events = [rewards[0]] + [rewards[i] - rewards[i-1] for i in range(1, len(rewards))]
    def compute_reward_rate(window_size):
        rate = []
        for i in range(len(reward_events)):
            if i < window_size:
                rate.append(np.nan)
            else:
                window_sum = sum(reward_events[i-window_size:i])
                rate.append(window_sum / window_size)
        return rate

    reward_rate = compute_reward_rate(initial_window)
    line, = ax.plot(range(len(reward_rate)), reward_rate, color='orange', label=f'Window size = {initial_window}')
    ax.legend()
    ax_slider = plt.axes([0.15, 0.1, 0.7, 0.03], facecolor='lightgoldenrodyellow')
    slider = Slider(ax_slider, 'Window Size', 1, max_window, valinit=initial_window, valstep=1)
    def update(val):
        w = int(slider.val)
        rate = compute_reward_rate(w)
        line.set_ydata(rate)
        line.set_label(f'Window size = {w}')
        ax.legend()
        fig.canvas.draw_idle()

    slider.on_changed(update)

    plt.show()

def draw_q_net_difference(ax, q_plus, q_minus, grid_size, fig=None, cbar=None):
    q_net_raw = q_plus - q_minus
    max_abs = np.max(np.abs(q_net_raw))
    if max_abs == 0:
        q_net = np.zeros_like(q_net_raw)
    else:
        q_net = q_net_raw / max_abs

    ax.clear()
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    ax.set_aspect('equal')
    ax.invert_yaxis()
    ax.axis('off')

    cmap = LinearSegmentedColormap.from_list("red_white_blue", ["#FF0000", "#FFFFFF", "#0000FF"], N=21)
    vmin, vmax = -1, 1
    norm = plt.Normalize(vmin=vmin, vmax=vmax)

    directions = {
        0: (0.3, 0.05),   # Up
        1: (0.3, 0.55),   # Down
        2: (0.05, 0.3),   # Left
        3: (0.55, 0.3),   # Right
    }
    size = 0.4

    for s in range(grid_size * grid_size):
        row, col = divmod(s, grid_size)
        center_x, center_y = col + 0.5, row + 0.5
        circle = patches.Circle((center_x, center_y), 0.05, facecolor='black', edgecolor='none')
        ax.add_patch(circle)

        for a in range(4):
            if (a == 0 and row == 0) or (a == 1 and row == grid_size - 1) or \
               (a == 2 and col == 0) or (a == 3 and col == grid_size - 1):
                continue

            dx, dy = directions[a]
            color = cmap(norm(q_net[s, a]))
            square = patches.Rectangle((col + dx, row + dy), size, size, facecolor=color, edgecolor='black', linewidth=0.2)
            ax.add_patch(square)

    if cbar is None and fig is not None:
        colorbar_width = 0.75
        cbar_ax = fig.add_axes([(1 - colorbar_width) / 2, 0.07, colorbar_width, 0.02])
        cbar = fig.colorbar(cm.ScalarMappable(cmap=cmap, norm=norm), cax=cbar_ax, orientation='horizontal')
        tick_values = np.linspace(-1, 1, num=5)
        cbar.set_ticks(tick_values)
        cbar.ax.set_xticklabels([f"{v:.2f}" for v in tick_values], fontsize=10)

    return cbar

def save_best_params_to_file(best_params, best_score, file_name="./results/best_params"):
    results_dir = os.path.dirname(file_name)
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)
    timestamp = datetime.now().strftime("%m%d_%H%M")
    filename = f"{file_name}_{timestamp}.txt"
    with open(filename, "w") as f:
        f.write(f"Best Hyperparameter Configuration:\n")
        f.write(f"LEARNING_RATE= {best_params[0]}\n")
        f.write(f"FORGET_REWARD= {best_params[1]}\n")
        f.write(f"FORGET_PUNISH= {best_params[2]}\n")
        f.write(f"Average Reward: {best_score}\n")
    print(f"Best parameters saved to {file_name}")

def plot_grid_search_results(results):
    df = pd.DataFrame([
        {"learning_rate": p[0], "forget_reward": p[1], "forget_punish": p[2], "reward": score}
        for (p, score) in results
    ])
    plot_heatmap_with_slider(df)

def save_numbers_to_csv(results, file_path):
    results_dir = os.path.dirname(file_path)
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)
    with open(file_path, mode='w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['learning_rate', 'forget_reward', 'forget_punish', 'avg_reward'])
        for params, score in results:
            row = list(params) + [score]
            writer.writerow(row)

def load_numbers_from_csv(file_path):
    results = []
    with open(file_path, newline='') as f:
        reader = csv.reader(f)
        next(reader)
        for row in reader:
            params = tuple(float(x) for x in row[:-1])
            score = float(row[-1])
            results.append((params, score))
    return results

def load_and_show_results(file_path):
    results = load_numbers_from_csv(file_path)
    plot_grid_search_results(results)
    
def plot_cumulative_and_reward_rate_with_slider(rewards):
    initial_window = 10
    max_window = max(1, len(rewards) // 2)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
    plt.subplots_adjust(bottom=0.25)
    # Plot cumulative reward
    ax1.set_title("Cumulative Reward")
    ax1.set_xlabel("Step")
    ax1.set_ylabel("Cumulative Reward")
    ax1.plot(range(len(rewards)), rewards, label="Cumulative Reward", color="blue")
    ax1.grid(True)

    # Plot reward rate
    reward_events = [rewards[0]] + [rewards[i] - rewards[i - 1] for i in range(1, len(rewards))]

    def compute_reward_rate(window_size):
        return [
            np.nan if i < window_size else sum(reward_events[i - window_size:i]) / window_size
            for i in range(len(reward_events))
        ]

    reward_rate = compute_reward_rate(initial_window)
    line2, = ax2.plot(range(len(reward_rate)), reward_rate, color='orange', label=f'Window = {initial_window}')
    ax2.set_title("Reward Rate with Adjustable Window")
    ax2.set_xlabel("Step")
    ax2.set_ylabel("Reward Rate")
    ax2.set_ylim(0, 1)
    ax2.grid(True)
    ax2.legend()

    # Slider
    ax_slider = plt.axes([0.2, 0.08, 0.6, 0.03], facecolor='lightgoldenrodyellow')
    slider = Slider(ax_slider, 'Window Size', 1, max_window, valinit=initial_window, valstep=1)

    def update(val):
        w = int(slider.val)
        new_rate = compute_reward_rate(w)
        line2.set_ydata(new_rate)
        line2.set_label(f'Window = {w}')
        ax2.legend()
        fig.canvas.draw_idle()

    slider.on_changed(update)

    plt.show()

def draw_policy_arrows(ax, q_plus, q_minus, grid_size):
    ax.clear()
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    ax.set_aspect('equal')
    ax.invert_yaxis()
    ax.axis('off')

    state_values = np.mean((q_plus + q_minus) / 2, axis=1)
    min_val, max_val = np.min(state_values), np.max(state_values)

    if max_val - min_val > 1e-6:
        normalized_values = (state_values - min_val) / (max_val - min_val)
    else:
        normalized_values = np.zeros_like(state_values)

    for s in range(grid_size * grid_size):
        row, col = divmod(s, grid_size)
        x, y = col, row
        gray = 1.0 - normalized_values[s]
        rect = patches.Rectangle((x, y), 1, 1, facecolor=(gray, gray, gray), edgecolor='none')
        ax.add_patch(rect)

    X, Y = [], []
    U, V = [], []

    for s in range(grid_size * grid_size):
        row, col = divmod(s, grid_size)
        center_x, center_y = col + 0.5, row + 0.5

        circle = patches.Circle((center_x, center_y), 0.04, facecolor='red', edgecolor='none')
        ax.add_patch(circle)

        available_actions = []
        if row > 0: available_actions.append(0)  # up
        if row < grid_size - 1: available_actions.append(1)  # down
        if col > 0: available_actions.append(2)  # left
        if col < grid_size - 1: available_actions.append(3)  # right

        if not available_actions:
            continue
        Qp = q_plus[s]
        Qm = q_minus[s]
        v = 0.0
        h = 0.0

        if 0 in available_actions: v += 0.25 * (Qp[0] - Qm[0])
        if 1 in available_actions: v -= 0.25 * (Qp[1] - Qm[1])
        if 2 in available_actions: h -= 0.25 * (Qp[2] - Qm[2])
        if 3 in available_actions: h += 0.25 * (Qp[3] - Qm[3])
        scale_factor = len(available_actions) / 4.0
        v /= scale_factor
        h /= scale_factor

        X.append(center_x)
        Y.append(center_y)
        U.append(h)
        V.append(v)

    magnitudes = np.sqrt(np.array(U)**2 + np.array(V)**2)
    if np.any(magnitudes > 1e-6):
        scale = np.max(magnitudes)
        U = np.array(U) / scale
        V = np.array(V) / scale
    # U = np.array(U)
    # V = np.array(V)

    # Normalize each arrow to unit length
    # magnitudes = np.sqrt(U**2 + V**2)
    # nonzero = magnitudes > 1e-8
    # U[nonzero] = U[nonzero] / magnitudes[nonzero]
    # V[nonzero] = V[nonzero] / magnitudes[nonzero]
    # ax.quiver(X, Y, U, -np.array(V), angles='xy', scale_units='xy', scale=2.5, color='red', width=0.01)
    ax.quiver(X, Y, U, -np.array(V),angles='xy',scale_units='xy',scale=2.5,color='red',width=0.005,headwidth=4,headlength=3,headaxislength=3.5)

def plot_heatmap_with_slider(df, fixed_param="learning_rate", ax=None, fig=None):
    df = df[(df["learning_rate"] > 0) &
            (df["forget_reward"] > 0) &
            (df["forget_punish"] > 0)]
    
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 8))
        plt.subplots_adjust(bottom=0.2)
    row_var, col_var = "forget_reward", "forget_punish"
    
    heatmap = None
    cbar = None  
    annotations = [] 

    def log_format(val, _):
        return r"${:.2f}$".format(np.log10(val))

    def get_sparse_labels(values, num_labels=5):
        total = len(values)
        if total <= num_labels:
            return [log_format(v, None) for v in values] 

        idxs = np.linspace(0, total - 1, num_labels, dtype=int)
        labels = []
        for i in range(total):
            if i in idxs:
                labels.append(log_format(values[i], None))
            else:
                labels.append("")
        return labels

    def update_heatmap(fixed_value):
        nonlocal heatmap, cbar, annotations
        
        df_fixed = df[df[fixed_param] == fixed_value]
        pivot = df_fixed.pivot(index=row_var, columns=col_var, values="reward")
        
        pivot = pivot.dropna(how="any", axis=0).dropna(how="any", axis=1)
        pivot = pivot[(pivot.index > 0)]
        all_pivot = pivot.copy()
        pivot = pivot.loc[:, (pivot.columns > 0)]
        
        if pivot.empty:
            print("Cannot plot heatmap.")
            return
        
        # Clear previous annotations if they exist
        for annotation in annotations:
            annotation.remove()
        annotations.clear()

        # Set min/max for color normalization across all data to keep color consistency
        vmin = np.min(df["reward"])
        vmax = np.max(df["reward"])

        # Initial plot
        if heatmap is None:
            heatmap = ax.imshow(pivot.values, cmap="viridis", aspect="auto", origin='lower', vmin=vmin, vmax=vmax)
            cbar = fig.colorbar(heatmap, ax=ax) 
            cbar.set_label('Average Reward')

            xticks = np.arange(len(pivot.columns))
            yticks = np.arange(len(pivot.index))

            x_tick_labels = get_sparse_labels(pivot.columns.to_numpy(), num_labels=5)
            y_tick_labels = get_sparse_labels(pivot.index.to_numpy(), num_labels=5)

            ax.set_xticks(xticks)
            ax.set_yticks(yticks)

            ax.set_xticklabels(x_tick_labels, fontsize=8)
            ax.set_yticklabels(y_tick_labels, fontsize=8)
            
            ax.set_xlabel(r"$\log_{10}(\phi_{-})$")
            ax.set_ylabel(r"$\log_{10}(\phi_{+})$")
            ax.set_title(f"Heatmap with fixed {fixed_param}={fixed_value:.1e}")
        else:
            # Update the heatmap data for subsequent calls
            heatmap.set_data(pivot.values)
            heatmap.set_clim(vmin, vmax)  # Update the color limits as well

            xticks = np.arange(len(pivot.columns))
            yticks = np.arange(len(pivot.index))

            x_tick_labels = get_sparse_labels(pivot.columns.to_numpy(), num_labels=5)
            y_tick_labels = get_sparse_labels(pivot.index.to_numpy(), num_labels=5)

            ax.set_xticks(xticks)
            ax.set_yticks(yticks)

            ax.set_xticklabels(x_tick_labels, fontsize=8)
            ax.set_yticklabels(y_tick_labels, fontsize=8)

            ax.set_title(f"Heatmap with fixed {fixed_param}={fixed_value:.1e}")
        
        # Add annotations
        for i in range(len(pivot.index)):
            for j in range(len(pivot.columns)):
                annotation = ax.text(j, i, f"{pivot.values[i, j]:.0f}", ha="center", va="center", color="white", fontsize=5)
                annotations.append(annotation)

        plt.draw()

    # Get unique values of the fixed parameter (learning rate)
    learning_rate_values = np.sort(df["learning_rate"].unique())
    initial_value = learning_rate_values[0]

    # Slider setup
    ax_slider = plt.axes([0.1, 0.01, 0.8, 0.05], facecolor='lightgoldenrodyellow')
    slider = Slider(ax_slider, 'Learning Rate', 0, len(learning_rate_values) - 1,
                valinit=0, valstep=1)

    def slider_update(val):
        idx = int(val)
        fixed_value = learning_rate_values[idx]
        update_heatmap(fixed_value)
        ax_slider.set_title(f"Learning Rate: {fixed_value:.5e}")
        plt.draw()

    ax_slider.set_title(f"Learning Rate: {initial_value:.5e}")
    slider.on_changed(slider_update)
    update_heatmap(initial_value)
    plt.show()

    return slider