<filename>utils/ERLplot_utils.py
#!/usr/bin/env python3

# Import modules
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib import ticker, gridspec
from deprecated import deprecated
import utils.ERLexperiment_utils as exp
import numpy as np
from scipy import ndimage
import seaborn as sns
from collections import Iterable
from utils.tile_coder import TileCoder
import matplotlib as mpl
import pickle


# TODO: Add documentation for plot_updates and plot_updates_
# TODO: Update action plots
# TODO: state_update_heatmap should do 2D heatmap of multi-dimensional
#       states, just give the two dimensions to plot in a heatmap. same for the
#       surface plot!!
# TODO: separate state_update_heatmap into two functions, one of which will
#   simply take a vector of update counts and create a heatmap of it given the
#   dimensions to reduce it to (e.g. reshape(100, 100))

# TODO: make _pendulum_ticks more general to label any number of ticks

# TODO: Plot mean with runs episodic
# TODO: Plot mean with stderr episodic for multiple data files
# TODO: Change action plot to histogram
# TODO: Make the state surface plot use either state updates or states visited
#       Like the new scheme in state heatmap plot
# TODO: State surface plot should take similar params to heatmap
# TODO: Improve documentation
# TODO: plot_mean_with_runs() should plot avg rewards for continuing


# Global variables
CMAP = "rocket"
DEFAULT_COLOURS = list(sns.color_palette(CMAP, 6).as_hex())
plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=sns.color_palette(CMAP))
# sns.set_theme(palette=CMAP)
OFFSET = 0  # The offset to start in DEFAULT_COLOURS


def updates(data, type_, ind, smooth_over, names, fig=None, ax=None,
            figsize=(12, 6), last_ind=-1, colours=None, xlim=None, ylim=None,
            env_type="continuing", label="", alpha=0.2, xscale="log",
            yscale="linear", base=10, xlabel=None, ylabel=None,
            runs=False, keep_shape=False, bootstrap=False, significance=0.1):
    """
    Plots the performance vs number of samples used in updates

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    names : list of str
        The name of the agents, used for the legend
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    colours : list of list of str
        The colours to use for each hyperparameter settings plot for each data
        dictionary
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    xscale : str, optional
        The scale at which to plot the x axis, by default "log"
    yscale : str, optional
        The scale at which to plot the y axis, by default "linear"
    base : int, optional
        The base of the scale of the x axis, by default 10
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    title : str, optional
        The title of the plot, by default ""
    runs : bool, optional
        Whether or not to plot the individual runs or error bars, by default
        False. If True, plots individual runs instead of error bars.
    keep_shape : bool, optional
        Whether the shape of the input data should be kept after smoothing, by
        default False. If True, then the first smooth_over episodes are
        smoothed and kept. If False, then the first smooth_over episodes are
        replaced by their fully smoothed value.
    bootstrap : bool, optional
        Whether or plot standard error or bootstrap confidence intervals, by
        default False. If True, bootstrap confidence intervals are plotted
        instead of standard error.
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.Figure, plt.Axes
        The figure and axes plotted on
    """
    if colours is None:
        colours = _get_default_colours(ind)

    # Setup figure
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel,
                         xscale, yscale, base)

    legend_labels = []
    for i in range(len(data)):
        for j in range(len(ind[i])):
            fig, ax = _updates(data[i], type_, ind[i][j], smooth_over[i],
                               fig, ax, figsize,
                               colours[i][j], xlim, ylim, env_type,
                               alpha, xscale, yscale, base, xlabel, ylabel,
                               runs, keep_shape, bootstrap, significance)

            # Create legend label and lines
            agent_name = names[i]
            label = agent_name
            legend_labels.append(label)

    ax.legend(legend_labels)

    return fig, ax


def _updates(data, type_, ind, smooth_over, fig=None, ax=None,
             figsize=(12, 6), colour=None, xlim=None, ylim=None,
             env_type="continuing", alpha=0.2, xscale="log",
             yscale="linear", base=10, xlabel=None, ylabel=None, runs=False,
             keep_shape=False, bootstrap=False, significance=0.1):
    if colour is None:
        colour = _get_default_colours([ind])[0]

    # # Setup figure
    # fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel,
    #                      xscale, yscale, base)

    hyperparams = data["experiment_data"][ind]["agent_hyperparams"]
    if "batch_size" in hyperparams:
        batch_size = hyperparams["batch_size"]
    else:
        batch_size = 1

    # Calculate updates to plot along x-axis
    steps_per_episode = data["experiment"]["environment"]["steps_per_episode"]
    total_steps = data["experiment"]["environment"]["total_timesteps"]

    # Include the first eval -- before any updates have been done
    start = 0
    stop = (total_steps + steps_per_episode) * batch_size
    step = steps_per_episode * batch_size
    updates = np.arange(start, stop, step)

    # Get the data to plot depending on the env type
    if env_type == "continuing":
        rewards, err, all_ = \
            rewards_per_update_continuing(data, ind, smooth_over, type_,
                                          keep_shape, bootstrap, significance)
    else:
        rewards, err, all_ = \
            rewards_per_update_episodic(data, ind, smooth_over, updates, type_,
                                        keep_shape, bootstrap, significance)

    # If the original shape is not kept, then the updates need to be adjusted
    # for smoothing
    if not keep_shape:
        updates = updates[:len(updates) - smooth_over + 1]

    # Plot the mean
    ax.step(updates, rewards, where="post", color=colour)
    if not runs:
        if not bootstrap:
            # Standard error
            ax.fill_between(updates, rewards-err, rewards+err, alpha=alpha,
                            step="post", color=colour)
        else:
            # Bootstrap confidence intervals
            ax.fill_between(updates, err[0], err[1], alpha=alpha, step="post",
                            color=colour)

    # Plot each run
    if runs:
        for i in range(all_.shape[0]):
            ax.step(updates, all_[i, :], linestyle="--", alpha=alpha,
                    color=colour, where="post")

    return fig, ax


def rewards_per_update_continuing(data, ind, smooth_over, type_="train",
                                  keep_shape=False, bootstrap=False,
                                  significance=0.1):
    returns = []
    for run in data["experiment_data"][ind]["runs"]:
        if type_ == "train":
            # The episodic returns for each run
            run_returns = list(run["train_episode_rewards"])

            # Plot the first eval episode, since this is the only data for
            # 0 updates
            first_ep = run["eval_episode_rewards"][0].Vmean()
            run_returns.insert(0, first_ep)

            returns.append(run_returns)

        else:
            returns.append(run["eval_episode_rewards"].mean(axis=-1))

    # Get the steps per episode so we can plot average reward
    steps_per_ep = data["experiment"]["environment"]["steps_per_episode"]
    returns = np.array(returns) / steps_per_ep

    returns = exp.smooth(returns, smooth_over, keep_shape=keep_shape)

    # Calculate standard error
    if not bootstrap:
        runs = len(data["experiment_data"][ind]["runs"])
        stderr = np.std(np.array(returns), axis=0) / np.sqrt(runs)
        return returns.mean(axis=0), stderr, returns

    # Calculate bootstrap confidence interval
    else:
        conf = exp.bootstrap_conf(returns, significance)
        return returns.mean(axis=0), conf, returns


def rewards_per_update_episodic(data, ind, smooth_over, updates, type_="train",
                                keep_shape=False, bootstrap=False,
                                significance=0.1):
    returns = []
    for run in data["experiment_data"][ind]["runs"]:
        if type_ == "train":
            # Determine the size of batches used in updates
            hyperparams = data["experiment_data"][ind]["agent_hyperparams"]
            if "batch_size" in hyperparams:
                batch_size = hyperparams["batch_size"]
            else:
                batch_size = 1

            # Calculate the cumulative number of updates over all previous
            # episodes
            run_updates = np.cumsum(run["train_episode_steps"]) * batch_size

            # For each cumulative number of updates plotted on the x-axis,
            # get the performance of the episode with number of updates closest
            # to this value
            run_returns = []
            for i in range(1, len(updates)):
                ep_gt_lower_bound = np.where(run_updates >= updates[i])[0][0]
                episode = run["train_episode_rewards"][ep_gt_lower_bound]
                run_returns.append(episode)

            # Plot the first eval episode, since this is the only data for
            # 0 updates
            first_ep = run["eval_episode_rewards"][0].mean()
            run_returns.insert(0, first_ep)
            returns.append(run_returns)

        elif type_ == "eval":
            returns.append(run["eval_episode_rewards"].mean(axis=-1))

    # # Smooth data
    returns = np.array(returns)
    returns = exp.smooth(returns, smooth_over, keep_shape=keep_shape)

    # Calculate standard error
    if not bootstrap:
        runs = len(data["experiment_data"][ind]["runs"])
        stderr = np.std(np.array(returns), axis=0) / np.sqrt(runs)
        return returns.mean(axis=0), stderr, returns

    # Calculate bootstrap confidence interval
    else:
        conf = exp.bootstrap_conf(returns, significance)
        return returns.mean(axis=0), conf, returns


def mean_with_bootstrap_conf(data, type_, ind, smooth_over, names,
                             fig=None, ax=None, figsize=(12, 6),
                             xlim=None, ylim=None, alpha=0.2,
                             colours=None, env_type="continuing",
                             significance=0.01, keep_shape=False,
                             xlabel=None, ylabel=None):
    """
    Plots the average training or evaluation return over all runs with
    confidence intervals.

    Given a list of data dictionaries of the form returned by main.py, this
    function will plot each episodic return for the list of hyperparameter
    settings ind each data dictionary. The ind argument is a list, where each
    element is a list of hyperparameter settings to plot for the data
    dictionary at the same index as this list. For example, if ind[i] = [1, 2],
    then plots will be generated for the data dictionary at location i
    in the data argument for hyperparameter settings ind[i] = [1, 2].
    The smooth_over argument tells how many previous data points to smooth
    over

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    names : list of str
        The name of the agents, used for the legend
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of list of str
        The colours to use for each hyperparameter settings plot for each data
        dictionary
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot
    """
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel)

    # Set the colours to be default if not set
    if colours is None:
        colours = _get_default_colours(ind)

    # Track the total timesteps per hyperparam setting over all episodes and
    # the cumulative timesteps per episode per data dictionary (timesteps
    # should be consistent between all hp settings in a single data dict)
    total_timesteps = []
    cumulative_timesteps = []

    for i in range(len(data)):
        if type_ == "train":
            cumulative_timesteps.append(exp.get_cumulative_timesteps(data[i]
                                        ["experiment_data"][ind[i][0]]["runs"]
                                        [0]["train_episode_steps"]))
        elif type_ == "eval":
            cumulative_timesteps.append(data[i]["experiment_data"][ind[i][0]]
                                        ["runs"][0]["timesteps_at_eval"])
        else:
            raise ValueError("type_ must be one of 'train', 'eval'")
        total_timesteps.append(cumulative_timesteps[-1][-1])

    # Find the minimum of total trained-for timesteps. Each plot will only
    # be plotted on the x-axis until this value
    min_timesteps = min(total_timesteps)

    # For each data dictionary, find the minimum index where the timestep at
    # that index is >=  minimum timestep
    ind_ge_min_timesteps = []
    for cumulative_timesteps_per_data in cumulative_timesteps:
        final_ind = np.where(cumulative_timesteps_per_data >=
                             min_timesteps)[0][0]
        # Since indexing will stop right before the minimum, increment it
        ind_ge_min_timesteps.append(final_ind + 1)

    # Plot all data for all HP settings, only up until the minimum index
    # fig, ax = None, None
    if env_type == "continuing":
        plot_fn = _plot_mean_with_conf_continuing
    else:
        plot_fn = _plot_mean_with_conf_episodic

    for i in range(len(data)):
        fig, ax = \
            plot_fn(data=data[i], type_=type_,
                    ind=ind[i], smooth_over=smooth_over[i], name=names[i],
                    fig=fig, ax=ax, figsize=figsize, xlim=xlim, ylim=ylim,
                    last_ind=ind_ge_min_timesteps[i], alpha=alpha,
                    colours=colours[i], significance=significance,
                    keep_shape=keep_shape)

    return fig, ax


def _plot_mean_with_conf_continuing(data, type_, ind, smooth_over, fig=None,
                                    ax=None, figsize=(12, 6), name="",
                                    last_ind=-1, xlabel="Episodes",
                                    ylabel="Average Return", xlim=None,
                                    ylim=None, alpha=0.2, colours=None,
                                    significance=0.01, keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a single
    data dictionary on a continuing environment. Bootstrap confidence intervals
    are plotted as shaded regions.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    timestep_multiply : int, optional
        A value to multiply each timstep by, by default 1. This is useful if
        your agent does multiple updates per timestep and you want to plot
        performance vs. number of updates.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    # This should be the exact same as the episodic version except without
    # reducing the episodes. Follow the same structure as the episodic function
    # and the continuing function with standard error.
    raise NotImplementedError


def _plot_mean_with_conf_episodic(data, type_, ind, smooth_over, fig=None,
                                  ax=None, figsize=(12, 6), name="",
                                  last_ind=-1, xlabel="Episodes",
                                  ylabel="Average Return", xlim=None,
                                  ylim=None, alpha=0.2, colours=None,
                                  significance=0.01, keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a single
    data dictionary on an episodic environment. Bootstrap confidence intervals
    are plotted as shaded regions.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    if colours is not None and len(colours) != len(ind):
        raise ValueError("must have one colour for each hyperparameter " +
                         "setting")

    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    # if ax is None and fig is None:
    #     fig = plt.figure(figsize=figsize)
    #     ax = fig.add_subplot()

    # if xlim is not None:
    #     ax.set_xlim(xlim)
    # if ylim is not None:
    #     ax.set_ylim(ylim)

    conf_level = "{:.2f}".format(1-significance)
    title = f"""Average {type_.title()} Return per Run with {conf_level}
                 Confidence Intervals"""
    fig, ax = _setup_fig(fig, ax, figsize, title, xlim, ylim, xlabel, ylabel)

    # Plot with bootstrap confidence interval
    for i in range(len(ind)):
        data = exp.reduce_episodes(data, ind[i], type_=type_)
        _, mean, conf = exp.get_mean_err(data, type_, ind[i], smooth_over,
                                         exp.bootstrap_conf,
                                         err_args={
                                            "significance": significance,
                                         },
                                         keep_shape=keep_shape)

        mean = mean[:last_ind]
        conf = conf[:, :last_ind]

        episodes = np.arange(mean.shape[0])

        # Plot based on colours
        label = name
        print(mean.shape, conf.shape, episodes.shape)
        _plot_shaded(ax, episodes, mean, conf, colours[i], label, alpha)

    ax.legend()
    conf_level = "{:.2f}".format(1-significance)
    ax.set_title(f"""Average {type_.title()} Return per Run with {conf_level}
                 Confidence Intervals""")
    # ax.set_ylabel(ylabel)
    # ax.set_xlabel(xlabel)

    fig.show()
    return fig, ax


@deprecated(reason="""Use exp.get_mean_err with exp.bootstrap_conf""")
def get_mean_bootstrap_conf(data, type_, ind, smooth_over,
                            env_type="continuing", significance=0.01,
                            keep_shape=False):
    """
    Gets the mean return per episode along with the bootstrap confidence
    interval on an episode-by-episode basis over all runs.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : int
        The hyperparameter setting to plot
    smooth_over : int
        The number of previous returns to smooth over. If less than 2, then
        no smoothing is performed
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    3-tuple of array_like
        The timesteps at each return to plot, the mean return to plot, and
        the confidence interval for each episode to plot as a shaded region
    """
    timesteps = None  # So the linter doesn't have a temper tantrum

    # Determine the timesteps to plot at
    if type_ == "eval":
        timesteps = \
            data["experiment_data"][ind]["runs"][0]["timesteps_at_eval"]

    elif type_ == "train":
        timesteps_per_ep = \
            data["experiment_data"][ind]["runs"][0]["train_episode_steps"]
        timesteps = exp.get_cumulative_timesteps(timesteps_per_ep)

    # Get the mean over all episodes per evaluation step (for online
    # returns, this axis will have length 1 so we squeeze it)
    returns = exp.get_returns(data, type_, ind, env_type=env_type)
    returns = returns.mean(axis=-1)

    # Smooth over returns
    returns = exp.smooth(returns, smooth_over, keep_shape=keep_shape)

    # Get the standard deviation of mean episodes per evaluation
    # step over all runs
    conf = exp.bootstrap_conf(returns, significance)

    # Get the mean over all runs
    mean = returns.mean(axis=0)

    # If we do not keep the shape, then due to smoothing we lose some
    # episodes
    if not keep_shape:
        end = len(timesteps) - smooth_over + 1
        timesteps = timesteps[:end]

    return timesteps, mean, conf


def plot_mean_with_runs(data, type_, ind, smooth_over, names, colours=None,
                        figsize=(12, 6), xlim=None, ylim=None, alpha=0.2,
                        plot_avg=True, env_type="continuing",
                        keep_shape=False, fig=None, ax=None, detrend=None):
    """
    Plots the mean return over all runs and the return for each run for a list
    of data dictionaries and hyperparameter indices

    Plots both the mean return per episode (over runs) as well as the return
    for each individual run (including "mini-runs", if a set of concurrent
    episodes were run for all runs, e.g. multiple evaluation episodes per
    run at set intervals)

    Note that this function takes in a list of data dictionaries and will
    plot the runs for each ind (which is a list of lists, where each super-list
    refers to a data dictionary and each sub-list refers to the indices for
    that data dictionary to plot).

    Example
    -------
    plot_mean_with_runs([sac_data, linear_data], "train", [[3439], [38, 3]],
    smooth_over=[5, 2], names=["SAC", "LAC"], figsize=(12, 6), alpha=0.2,
    plot_avg=True, env_type="episodic")

    will plot hyperparameter index 3439 for the sac_data, smoothing over the
    last 5 episodes, and the label will have the term "SAC" in it; also plots
    the mean and each individual run on the linear_data for hyperparameter
    settings 38 and 3, smoothing over the last 2 episodes for each and with
    the term "LAC" in the labels.

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    figsize : tuple(int, int)
        The size of the figure to plot
    names : list of str
        The name of the agents, used for the legend
    colours : list of list of str, optional
        The colours to use for each hyperparameter settings plot for each data
        dictionary, by default None
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha to use for plots of the runs, by default 0.1
    plot_avg : bool, optional
        If concurrent episodes are executed in each run (e.g. multiple
        evaluation episodes are run at set intervals), then whether to plot the
        performance of each separately or the average performance over all
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    fig : plt.Figure
        The figure to plot on
    ax : plt.Axes
        The axis to plot on
    detrend : str, optional
        How to detrend the runs for each data, by default None. One of
        "difference", "linear", or "mean"

    Returns
    -------
    tuple of plt.figure, plt.Axes
        The figure and axis plotted on
    """
    # Set the colours to be default if not set
    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    # fig = plt.figure(figsize=figsize)
    # ax = fig.add_subplot()
    if env_type == "continuing":
        xlabel = "Timesteps"
        ylabel = "Average Reward"
    else:
        xlabel = "Episodes"
        ylabel = "Return"
    title = f"Mean Return with Runs"

    fig, ax = _setup_fig(fig, ax, figsize, title, xlim, ylim, xlabel, ylabel)

    # Plot for each data dictionary given
    legend_lines = []
    legend_labels = []
    for i in range(len(data)):
        for _ in range(len(ind[i])):
            fig, ax, labels, lines = \
                _plot_mean_with_runs(data[i], type_, ind[i], smooth_over[i],
                                     names[i], colours[i], figsize, xlim, ylim,
                                     alpha, plot_avg, env_type, fig, ax,
                                     keep_shape, detrend)

            legend_lines.extend(lines)
            legend_labels.extend(labels)

    ax.legend(legend_lines, legend_labels)
    fig.show()

    return fig, ax


def _plot_mean_with_runs(data, type_, ind, smooth_over, name, colours=None,
                         figsize=(12, 6), xlim=None, ylim=None, alpha=0.2,
                         plot_avg=True, env_type="continuing", fig=None,
                         ax=None, keep_shape=False, detrend=None):
    """
    Plots the mean return over all runs and the return for each run for a
    single data dictionary and for each in a list of hyperparameter settings.

    Similar to plot_mean_with_runs, except that this function takes in only
    a single data dictionary.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot for
        each agent. For example [1, 2] means that the agent plots will use
        hyperparameter settings indices 1 and 2.
    smooth_over : int
        The number of previous data points to smooth over for the agent's
        plot. Note that this is *not* the number of timesteps to smooth over,
        but rather the number of data points to smooth over. For example,
        if you save the return every 1,000 timesteps, then setting this value
        to 15 will smooth over the last 15 readings, or 15,000 timesteps.
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str
        The name of the agents, used for the legend
    colours : list of list of str, optional
        The colours to use for each hyperparameter settings plot for each data
        dictionary, by default None
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha to use for plots of the runs, by default 0.1
    plot_avg : bool, optional
        If concurrent episodes are executed in each run (e.g. multiple
        evaluation episodes are run at set intervals), then whether to plot the
        performance of each separately or the average performance over all
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    detrend : str, optional
        How to detrend the runs for each data, by default None. One of
        "difference", "linear", or "mean"

    Returns
    -------
    tuple of plt.figure, plt.Axes, list of str, list of mpl.Lines2D
        The figure and axis plotted on as well as the list of strings to use
        as labels and the list of lines to include in the legend
    """
    # Set up figure and axis
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim)

    if colours is None:
        colours = _get_default_colours(ind)

    # Store the info to keep in the legend
    legend_labels = []
    legend_lines = []

    # Plot each selected hyperparameter setting in the data dictionary
    for j in range(len(ind)):
        fig, ax, labels, lines = \
            _plot_mean_with_runs_single_hp(data, type_, ind[j], smooth_over,
                                           name, colours[j], figsize, xlim,
                                           ylim, alpha, plot_avg, env_type,
                                           fig, ax, keep_shape, detrend)
        legend_labels.extend(labels)
        legend_lines.extend(lines)

    return fig, ax, legend_labels, legend_lines


def _plot_mean_with_runs_single_hp(data, type_, ind, smooth_over, name,
                                   colour=None, figsize=(12, 6), xlim=None,
                                   ylim=None, alpha=0.2, plot_avg=True,
                                   env_type="continuing", fig=None, ax=None,
                                   keep_shape=False, detrend=None):
    """
    Plots the mean return over all runs and the return for each run for a
    single data dictionary and a single hyperparameter setting.

    Similar to _plot_mean_with_runs, except that this function takes in only
    a single data dictionary and a single hyperparameter setting.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : int
        The hyperparameter settings indices to plot for the agent. For example
        5 means that the agent plots will use hyperparameter settings index 5.
    smooth_over : int
        The number of previous data points to smooth over for the agent's
        plot. Note that this is *not* the number of timesteps to smooth over,
        but rather the number of data points to smooth over. For example,
        if you save the return every 1,000 timesteps, then setting this value
        to 15 will smooth over the last 15 readings, or 15,000 timesteps.
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str
        The name of the agents, used for the legend
    colours : list of list of str, optional
        The colours to use for each hyperparameter settings plot for each data
        dictionary, by default None
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha to use for plots of the runs, by default 0.1
    plot_avg : bool, optional
        If concurrent episodes are executed in each run (e.g. multiple
        evaluation episodes are run at set intervals), then whether to plot the
        performance of each separately or the average performance over all
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    detrend : str, optional
        How to detrend the runs for each data, by default None. One of
        "difference", "linear", or "mean"

    Returns
    -------
    tuple of plt.figure, plt.Axes, list of str, list of mpl.Lines2D
        The figure and axis plotted on as well as the list of strings to use
        as labels and the list of lines to include in the legend
    """
    # Set up figure and axis
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim)

    if colour is None:
        colour = _get_default_colours([ind])[0]

    # Determine the timesteps to plot at
    if type_ == "eval":
        timesteps = \
            data["experiment_data"][ind]["runs"][0]["timesteps_at_eval"]

    elif type_ == "train":
        timesteps_per_ep = \
            data["experiment_data"][ind]["runs"][0]["train_episode_steps"]
        timesteps = exp.get_cumulative_timesteps(timesteps_per_ep)

    # Plot the average reward
    if env_type == "continuing":
        episode_steps = data["experiment"]["environment"]["steps_per_episode"]

    # Get returns
    all_returns = exp.get_returns(data, type_, ind, env_type)

    # If concurrent episodes are run in each run then average them if
    # appropriate
    if type_ == "eval" and plot_avg:
        all_returns = all_returns.mean(axis=-1)
    elif type_ == "eval" and not plot_avg:
        print(all_returns.shape)
        all_returns = np.concatenate(all_returns, axis=1)
        all_returns = all_returns.transpose()
    elif type_ == "train":
        all_returns = np.squeeze(all_returns)

    # Smooth returns
    all_returns = exp.smooth(all_returns, smooth_over, keep_shape)

    # Plot the average reward
    if env_type == "continuing":
        episode_steps = data["experiment"]["environment"]
        episode_steps = episode_steps["steps_per_episode"]
        all_returns /= episode_steps

    # Detrend the data if appropriate
    if detrend == "linear":
        all_returns = exp.detrend_linear(all_returns, 1)
    elif detrend == "mean":
        all_returns = exp.detrend_linear(all_returns, 1, "constant")
    elif detrend == "difference":
        all_returns = exp.detrend_difference(all_returns, 1)

    # Determine whether to plot episodes or timesteps on the x-axis, which is
    # dependent on the environment type
    if env_type == "episodic":
        xvalues = np.arange(all_returns.shape[1])  # episodes
    else:
        xvalues = timesteps[:all_returns.shape[1]]

    # Plot each run
    for run in range(all_returns.shape[0]):
        ax.plot(xvalues, all_returns[run], color=colour, linestyle="--",
                alpha=alpha)

    # Plot the mean
    mean = all_returns.mean(axis=0)
    ax.plot(xvalues, mean, color=colour)

    # Store legend identifiers for the run
    legend_labels = []
    legend_lines = []
    legend_labels.append("Individual Runs")
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="--",
                               alpha=alpha))

    # Set up the legend variables for the mean over all runs
    label = name
    legend_labels.append(label)
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="-"))

    return fig, ax, legend_labels, legend_lines


def mean_with_stderr(data, type_, ind, smooth_over, names,
                     fig=None, ax=None, figsize=(12, 6),
                     xlim=None, ylim=None, alpha=0.2,
                     colours=None, env_type="continuing",
                     keep_shape=False, xlabel="", ylabel=""):
    """
    Plots the average training or evaluation return over all runs with standard
    error.

    Given a list of data dictionaries of the form returned by main.py, this
    function will plot each episodic return for the list of hyperparameter
    settings ind each data dictionary. The ind argument is a list, where each
    element is a list of hyperparameter settings to plot for the data
    dictionary at the same index as this list. For example, if ind[i] = [1, 2],
    then plots will be generated for the data dictionary at location i
    in the data argument for hyperparameter settings ind[i] = [1, 2].
    The smooth_over argument tells how many previous data points to smooth
    over

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    names : list of str
        The name of the agents, used for the legend
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of list of str
        The colours to use for each hyperparameter settings plot for each data
        dictionary
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot
    """
    # Set the colours to be default if not set
    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    title = f"Average {type_.title()} Return per Run with Standard Error"
    fig, ax = _setup_fig(fig, ax, figsize, xlim=xlim, ylim=ylim, xlabel=xlabel,
                         ylabel=ylabel, title=title)
    # if ax is None and fig is None:
    #     fig = plt.figure(figsize=figsize)
    #     ax = fig.add_subplot()

    # if xlim is not None:
    #     ax.set_xlim(xlim)
    # if ylim is not None:
    #     ax.set_ylim(ylim)

    # Track the total timesteps per hyperparam setting over all episodes and
    # the cumulative timesteps per episode per data dictionary (timesteps
    # should be consistent between all hp settings in a single data dict)
    total_timesteps = []
    cumulative_timesteps = []

    for i in range(len(data)):
        if type_ == "train":
            cumulative_timesteps.append(exp.get_cumulative_timesteps(data[i]
                                        ["experiment_data"][ind[i][0]]["runs"]
                                        [0]["train_episode_steps"]))
        elif type_ == "eval":
            cumulative_timesteps.append(data[i]["experiment_data"][ind[i][0]]
                                        ["runs"][0]["timesteps_at_eval"])
        else:
            raise ValueError("type_ must be one of 'train', 'eval'")
        total_timesteps.append(cumulative_timesteps[-1][-1])

    # Find the minimum of total trained-for timesteps. Each plot will only
    # be plotted on the x-axis until this value
    min_timesteps = min(total_timesteps)

    # For each data dictionary, find the minimum index where the timestep at
    # that index is >=  minimum timestep
    ind_ge_min_timesteps = []
    for cumulative_timesteps_per_data in cumulative_timesteps:
        final_ind = np.where(cumulative_timesteps_per_data >=
                             min_timesteps)[0][0]
        # Since indexing will stop right before the minimum, increment it
        ind_ge_min_timesteps.append(final_ind + 1)

    # Plot all data for all HP settings, only up until the minimum index
    # fig, ax = None, None
    plot_fn = _plot_mean_with_stderr_continuing if env_type == "continuing" \
        else _plot_mean_with_stderr_episodic

    for i in range(len(data)):
        fig, ax = \
            plot_fn(data=data[i], type_=type_,
                    ind=ind[i], smooth_over=smooth_over[i], name=names[i],
                    fig=fig, ax=ax, figsize=figsize, xlim=xlim, ylim=ylim,
                    last_ind=ind_ge_min_timesteps[i], alpha=alpha,
                    colours=colours[i], keep_shape=keep_shape)

    return fig, ax


def _plot_mean_with_stderr_continuing(data, type_, ind, smooth_over, fig=None,
                                      ax=None, figsize=(12, 6), xlim=None,
                                      ylim=None, xlabel=None, ylabel=None,
                                      name="", last_ind=-1,
                                      timestep_multiply=None, alpha=0.2,
                                      colours=None,
                                      keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a single
    data dictionary on a continuing environment. Standard error
    is plotted as shaded regions.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    timestep_multiply : array_like of float, optional
        A value to multiply each timstep by, by default None. This is useful if
        your agent does multiple updates per timestep and you want to plot
        performance vs. number of updates.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    if colours is not None and len(colours) != len(ind):
        raise ValueError("must have one colour for each hyperparameter " +
                         "setting")

    if timestep_multiply is None:
        timestep_multiply = [1] * len(ind)

    if ax is not None and fig is None:
        raise ValueError("must pass figure when passing axis")

    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    if ax is None and fig is None:
        title = f"Average {type_.title()} Return per Run with Standard Error"
        fig, ax = _setup_fig(fig, ax, figsize, xlim=xlim, ylim=ylim,
                             xlabel=xlabel, ylabel=ylabel, title=title)
    #     fig = plt.figure(figsize=figsize)
    #     ax = fig.add_subplot()

    # if xlim is not None:
    #     ax.set_xlim(xlim)
    # if ylim is not None:
    #     ax.set_ylim(ylim)

    episode_length = data["experiment"]["environment"]["steps_per_episode"]

    # Plot with the standard error
    for i in range(len(ind)):
        timesteps, mean, std = exp.get_mean_err(data, type_, ind[i],
                                                smooth_over, exp.stderr,
                                                keep_shape=keep_shape)
        timesteps = np.array(timesteps[:last_ind]) * timestep_multiply[i]
        mean = mean[:last_ind] / episode_length
        std = std[:last_ind] / episode_length

        # Plot based on colours
        label = name
        if colours is not None:
            _plot_shaded(ax, timesteps, mean, std, colours[i], label, alpha)
        else:
            _plot_shaded(ax, timesteps, mean, std, None, label, alpha)

    ax.legend()

    fig.show()
    return fig, ax


def _plot_mean_with_stderr_episodic(data, type_, ind, smooth_over, fig=None,
                                    ax=None, figsize=(12, 6), name="",
                                    last_ind=-1, xlabel="Episodes",
                                    ylabel="Average Return", xlim=None,
                                    ylim=None, alpha=0.2, colours=None,
                                    keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a
    single data dictionary on an episodic environment. Plots shaded retions
    as standard error.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    if colours is not None and len(colours) != len(ind):
        raise ValueError("must have one colour for each hyperparameter " +
                         "setting")

    if ax is not None and fig is None:
        raise ValueError("must pass figure when passing axis")

    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    if ax is None and fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    # Plot with the standard error
    for i in range(len(ind)):
        data = exp.reduce_episodes(data, ind[i], type_=type_)

        # data has consistent # of episodes, so treat as env_type="continuing"
        _, mean, std = exp.get_mean_err(data, type_, ind[i], smooth_over,
                                        exp.stderr, keep_shape=keep_shape)
        mean = mean[:last_ind]
        std = std[:last_ind]
        episodes = np.arange(mean.shape[0])

        # Plot based on colours
        label = name
        if colours is not None:
            _plot_shaded(ax, episodes, mean, std, colours[i], label, alpha)
        else:
            _plot_shaded(ax, episodes, mean, std, None, label, alpha)

    ax.legend()
    ax.set_title(f"Average {type_.title()} Return per Run with Standard Error")
    ax.set_ylabel(ylabel)
    ax.set_xlabel(xlabel)

    fig.show()
    return fig, ax


def hyperparam_sensitivity_plot(dir, type_, hp_name, combo, env_config,
                                agent_config, figsize=(12, 6), fig=None,
                                ax=None, annotate=False, colour=None,
                                env_type="continuing", conf=False,
                                significance=0.1):
    """
    Plots the hyperparameter sensitivity for hp_name, where the combo argument
    determines which hyperparameters are held constant.

    Given hyperparameters a, b, and c, let's say we want to get all
    hyperparameter settings indices where a varies, and b and c are constant.
    if a, b, and c can each be 1 or 2, then there are four ways that a can
    vary with b and c remaining constant:

        [
            ((a=1, b=1, c=1), (a=2, b=1, c=1)),         combo = 0
            ((a=1, b=2, c=1), (a=2, b=2, c=1)),         combo = 1
            ((a=1, b=1, c=2), (a=2, b=1, c=2)),         combo = 2
            ((a=1, b=2, c=2), (a=2, b=2, c=2))          combo = 3
        ]

    The combo argument indexes into this list of hyperparameter settings

    Parameters
    ----------
    dir : str
        The directory containing all data dictionaries generated from running
        main.py, separated into one file per HP index
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    combo : int
        Determines the values of the constant hyperparameters. Given that
        only one hyperparameter may vary, there are many different sets
        having this hyperparameter varying with all others remaining constant
        since each constant hyperparameter may take on many values. This
        argument cycles through all sets of hyperparameter settings indices
        that have only one hyperparameter varying and all others constant.
    env_config : dict
        The environment configuration file as a Python dictionary
    agent_config : dict
        The agent configuration file as a Python dictionary
    figsize : tuple(int, int)
        The size of the figure
    fig : plt.figure, optional
        The figure to plot on, by default None. If None, a new figure is
        created
    ax : plt.Axes, optional
        The axis to plot on, by default None. If None, a new axis is created.
    annotate : bool, optional
        Whether or not to annoate each point drawn with its y-value
    colour : str
        The colour to use in the plot, by default None. If None, then default
        values will be used.

    Returns
    -------
    tuple of plt.figure, plt.Axes
        The figure and axis that was plotted on

    Raises
    ------
    ValueError
        If only one of fig, ax are specified
    """
    if (fig is None and ax is not None) or (ax is None and fig is not None):
        raise ValueError("fig and ax must both be specified if one is "
                         "specified")

    if not conf:
        stderr_fn = exp.get_mean_returns_with_stderr_hp_varying
        hp_values, mean_returns, error = \
            stderr_fn(dir, type_, hp_name, combo, env_config, agent_config,
                      after=0, env_type=env_type)
    else:
        conf_fn = exp.get_mean_returns_with_conf_hp_varying
        hp_values, mean_returns, error = \
            conf_fn(dir, type_, hp_name, combo, env_config, agent_config,
                    after=0, env_type=env_type, significance=significance)

    # Set up the figure
    if fig is None and ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    ax.set_title(f"{hp_name.title()} Hyperparameter Sensitivity Plot " +
                 f"over {type_.title()} Data")
    ax.set_xlabel("Hyperparameter Values")

    if env_type == "episodic":
        ax.set_ylabel("Average Return")
    elif env_type == "continuing":
        ax.set_ylabel("Average Reward")

    if colour is None:
        colour = DEFAULT_COLOURS[0]

    if env_type == "continuing":
        ep_steps = env_config["steps_per_episode"]
        mean_returns /= ep_steps
        error /= ep_steps

    _plot_shaded(ax, hp_values, mean_returns, error, colour, "",
                 alpha=0.25)

    if annotate:
        for i, coords in enumerate(zip(hp_values, mean_returns)):
            ax.annotate(coords[i], xy=coords)

    return fig, ax


def hyperparam_sensitivity_plot_by_settings_index(dir, type_, hp_name, ind,
                                                  env_config, agent_config,
                                                  figsize=(12, 6),
                                                  annotate=False,
                                                  env_type="continuing",
                                                  conf=False,
                                                  significance=0.1):
    """
    Plot hyperparameter sensitivity by hyperparameter settings index.

    Plots the hyperparameter sensitivity plot where all
    constant hyperparameters are equal to those defined by the
    argument hyperparameter settings index. The only hyperparameter that
    is allowed to vary is the one corresponding to hp_name.

    Parameters
    ----------
    dir : str
        The directory containing all data dictionaries generated by running
        main.py
    type_ : str
        Whether to plot the training or evaluation data, one of
        'train', 'eval'
    hp_name : str
        The name of the hyperparameter to plot the sensitivity of
    config : dict
        A Python dictionary of the agent configuration file used to run
        the experiment
    ind : int
        The hyperparameter settings index that should define the
        values of all constant hyperparameters. Only the hyperparameter
        defined by hp_name is allowed to vary, all others must be equal to
        the values of those defined by the settings index ind.
    env_config : dict
        The environment configuration file as a Python dictionary
    agent_config : dict
        The agent configuration file as a Python dictionary
    figsize : tuple(int, int)
        The size of the figure
    annotate : bool, optional
        Whether or not to annoate each point drawn with its y-value

    Returns
    -------
    tuple of plt.figure, plt.Axes
        The figure and axis that was plotted on
    """
    # Get all the combinations of hyperparameters with only hp_name varying
    hp_name_varying_only_settings = \
        exp.get_varying_single_hyperparam(env_config, agent_config, hp_name)

    # Get the index (into the list of tuples with only hp_name varying) of
    # the tuple which has hp_name varying only, but all other constant
    # hyperparameters equal to those defined by hyperparam setting number ind.
    # So hp_combo_constant_hps_equal_ind is the hyperparam combination with
    # all constant hyperparams equal to those in settings number ind.
    hp_combo_constant_hps_equal_ind = \
        next(filter(lambda x: ind in x, hp_name_varying_only_settings))
    combo_index = hp_name_varying_only_settings.index(
        hp_combo_constant_hps_equal_ind)

    print(hp_name_varying_only_settings[combo_index])

    # Plot the hyperparameter sensitivity plot using all constant
    # hyperparameters equal to those defined in the settings
    # number ind
    return hyperparam_sensitivity_plot(dir, type_, hp_name, combo_index,
                                       env_config, agent_config,
                                       figsize=figsize, annotate=annotate,
                                       env_type=env_type, conf=conf,
                                       significance=significance)


@deprecated(reason="""We no longer store the state seen at each timestep
            as this was becoming very memory intensive. We are storing the
            states seen at each timestep as tile coded counts, similar to the
            states used in updates""")
def plot_states_visited(data, figsize=(15, 10), max_=0.1, xlabel="", ylabel="",
                        title="", hp_ind=-1, rasterize=True,
                        env_type="continuing"):
    """
    Plots the state visitation distribution during training

    Given a data dictionary, plots the state visitation distribution of a
    single hyperparameter setting.

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : float, optional
        The cutoff threshold for updates, by default 1.0. Any states that have
        updates larger than this will be capped at this value. Note that
        update counts are normalized between [0, 1].
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    title : str, optional
        The title of the plot, by default ""
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'; by default
        "continuing"

    Returns
    -------
    plt.figure, plt.Axes, array_like
        The figure and axis plotted on and the array of state counts
    """
    # Choose the hyperparameter setting if not st
    if hp_ind < 0:
        hp_ind = next(iter(data["experiment_data"]))

    # Set up the TileCoder
    min_state_values = data["experiment"]["environment"]["state_low"]
    max_state_values = data["experiment"]["environment"]["state_high"]
    state_bins = data["experiment"]["environment"]["state_bins"]
    state_coder = TileCoder(low=min_state_values, high=max_state_values,
                            bins=state_bins, num_tilings=1, include_bias=False,
                            offset=False, normalize=False)

    states_visited = np.zeros(np.prod(state_bins))

    # If the type is continuing, then all data fits nicely into a single matrix
    # and we can easily calculate everything using array operations
    if env_type == "continuing":
        size = data["experiment_data"][hp_ind]["runs"][0]["train_states"].shape
        for run in data["experiment_data"][hp_ind]["runs"]:
            states = run["train_states"]
            states = states.reshape(size[0] * size[1], size[2])
            states_visited += state_coder(states).mean(axis=0)

    # If episodic, then since runs/episodes have different lengths, we instead
    # have an array, where each row /col is a different length (list of states)
    # since each run/episode has a different length. So, we have to deal with
    # this case differently, looping over all runs, then all states per
    # episode per run
    elif env_type == "episodic":
        for run in data["experiment_data"][hp_ind]["runs"]:
            for episode_states in run["train_states"]:
                # For each episode, get the states and tile code, we go through
                # episodes separately since each has its own length
                states = np.stack(episode_states)
                states_visited += state_coder(states).mean(axis=0)

    fig, ax = _state_heatmap(states_visited, data["experiment"]["environment"],
                             figsize, max_, xlabel, ylabel, rasterize, title),
    return fig, ax, states_visited


def state_heatmap(data, figsize=(15, 10), max_=1.0, xlabel="", ylabel="",
                  updates=True, use_next_state=False, hp_ind=-1,
                  rasterize=True, time_ind=-1, use_interval_only=False,
                  title="", cbar=True, fig=None, ax=None):
    """
    Plot a heatmap of the states used in each update

    This function will only plot a 2D heatmap of states. The first state
    coordinate will refer to the rows of the heatmap, while the second will
    refer to the columns. The minimum value of the first state coordinate
    will be the top of the heatmap, while the max will be the bottom.
    Similarly, the maximum of the second coordinate will be the right of
    the plot, while the minimum will be the left.

    This function assumes a vector that has as each element, the number of
    updates performed to that corresponding state. If the state space is
    continuous, simply use a high resolution tile coder and discretize the
    state space.

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : float, optional
        The cutoff threshold for updates, by default 1.0. Any states that have
        updates larger than this will be capped at this value. Note that
        update counts are normalized between [0, 1].
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    updates : bool, optional
        Whether to plot the states used in updates or the states visited, by
        default True.
    use_next_state : bool, optional
        Whether to plot the state or next state in the SARSA update,
        by default False
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.
    time_ind : int, optional
        In each experiment, the state update heatmap is stored at multiple
        timesteps at certain intervals. This parameter controls which of the
        update heatmaps to plot, with 0 being the first and -1 being the last
        heatmap saved, by default -1.
    use_interval_only : bool, optional
        Whether to plot only the states visited between the time_ind reading
        and the previous reading. That is, whether or not to plot only the
        states used in updates or visited between the argument time_ind and
        the previous, but none before. By default False.
    title : str, optional
        The title of the figure, by default ""

    Returns
    -------
    plt.figure, plt.Axes, array_like
        The figure and axis plotted on and the array of state counts

    Raises
    ------
    ValueError
        If attempting to plot the next states, but not plotting states used in
        updates
    """
    if fig is None and ax is not None:
        raise ValueError("fig and ax must both be specified")

    if use_next_state and not updates:
        raise ValueError("""can only plot next states when plotting states
                         used in updates, set 'updates=True'""")

    # Choose the hyperparameter setting if not st
    if hp_ind < 0:
        hp_ind = next(iter(data["experiment_data"]))

    # Get the number of updates per each state per run
    if updates:
        state_type = "update_states" if not use_next_state else \
            "update_next_states"
    else:
        state_type = "train_states"

    # Get the mean states used in each update or visited at each timestep
    # over all runs
    states = _get_mean_state_bin_count(data, state_type, hp_ind, time_ind)

    # Calculate the states seen only for the current interval
    if use_interval_only:
        if time_ind == 0:
            prev_states = np.zeros(states.shape)
        else:
            prev_states = _get_mean_state_bin_count(data, state_type, hp_ind,
                                                    time_ind-1)
        states -= prev_states

    # Plot
    fig, ax = _state_heatmap(states, data["experiment"]["environment"],
                             figsize, max_, xlabel, ylabel, title, rasterize,
                             cbar, fig=fig, ax=ax)
    return fig, ax, states


def _get_mean_state_bin_count(data, state_type, hp_ind, time_ind):
    # Get the states used in each update or visited at each timestep
    states = []
    for run in data["experiment_data"][hp_ind]["runs"]:
        states.append(run[state_type][time_ind])

    # Get the average number of updates per state over all runs
    return np.stack(states).mean(axis=0)


def _state_heatmap(visited_states, env_config, figsize=(15, 10), max_=0.1,
                   xlabel="", ylabel="", title="", rasterize=True, cbar=True,
                   fig=None, ax=None):
    """
    Plots a heatmap of states, where states are separated into bins

    This function is designed to work with the TileCoder object. It takes in
    a flattened array, where each index in the array refers to a sequential
    bin of states. The element at each position is the number of times that
    state was updated or encountered.

    Parameters
    ----------
    visited_states : array_like
        The flattened array of counts, where each index in the array
        corresponds to the index of the bin
    env_config : dict
        The environment configuration file for the environment that the
        experiment was run on
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : float, optional
        The cutoff threshold for updates, by default 1.0. Any states that have
        updates larger than this will be capped at this value. Note that
        update counts are normalized between [0, 1].
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    title : str, optional
        The title of the figure, by default ""
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.

    Returns
    -------
    plt.figure, plt.Axes, array_like
        The figure and axis plotted on and the array of state counts
    """
    if fig is None and ax is not None:
        raise ValueError("fig and ax must both be specified")

    if fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()
    ax.set_title(title)
    ax.set(xlabel=xlabel, ylabel=ylabel)

    max_visited = visited_states.max()
    min_visited = visited_states.min()
    if max_visited != min_visited:
        visited_states = (visited_states - min_visited) / \
            (max_visited - min_visited)

    # Threshold the updates
    ind = np.where(visited_states > max_)[0]
    visited_states[ind] = max_

    # Determine the bins per state dimension
    if len(env_config["state_bins"]) == 1:
        state_bins = (1, env_config["state_bins"][0])
    else:
        state_bins = env_config["state_bins"]
    visited_states = visited_states.reshape(state_bins)

    # Plot the heatmap
    ax = sns.heatmap(visited_states, ax=ax, linewidths=0, rasterized=rasterize,
                     cbar=cbar, vmin=0, vmax=1.0)

    # Set the axis tick locations
    num_ticks = 3
    ax.xaxis.set_major_locator(ticker.LinearLocator(num_ticks))
    ax.yaxis.set_major_locator(ticker.LinearLocator(num_ticks))

    # Label axis ticks
    xtick, ytick = get_ticks(num_ticks, num_ticks, env_config, type_="state")

    ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    fig.show()
    return fig, ax


def action_heatmap_with_barchart(data, figsize=(20, 7), max_=-1,
                                 xlabel="", ylabel="", hp_ind=-1,
                                 rasterize=True, ratio=(2, 1), time_ind=-1):
    """
    Plot the action heatmap with a barchart above it

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : int, optional
        The cutoff threshold for updates, by default 10. Any states that have
        updates larger than this will be capped at this value.
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.
    ratio : tuple, optional
        The ratio of bar char size to heatmap size, by default (2, 1)
    time_ind : int, optional
        In each experiment, the state update heatmap is stored at multiple
        timesteps at certain intervals. This parameter controls which of the
        update heatmaps to plot, with 0 being the first and -1 being the last
        heatmap saved, by default -1.

    Returns
    -------
    plt.figure, plt.Axes, plt.Axes
        The figure, barchart axis, and heatmap axis plotte on
    """
    # Set up the plot
    fig = plt.figure(figsize=figsize)
    gs = gridspec.GridSpec(*ratio, height_ratios=[4, 1])
    bar_ax = fig.add_subplot(gs[0])
    heat_ax = fig.add_subplot(gs[1])

    # Get the number updates per action bin
    if hp_ind < 0:
        hp_ind = next(iter(data["experiment_data"]))
    update_actions = []
    for run in data["experiment_data"][hp_ind]["runs"]:
        update_actions.append(run["update_actions"][time_ind])

    # Get the average number of upates per state over all runs
    update_actions = np.stack(update_actions).mean(axis=0)

    # Plot the barchart
    bins = np.arange(update_actions.shape[0])
    bar_ax.bar(bins, update_actions)

    # Set the axis limits for the barchart
    bar_ax.set_xlim((bins[0], bins[-1]))
    ymax = update_actions.max()
    bar_ax.set_ylim((0, ymax))

    # Threshold the updates
    if max_ > -1:
        ind = np.where(update_actions > max_)[0]
        update_actions[ind] = max_

    # Determine the bins per action dimension
    env_config = data["experiment"]["environment"]
    if len(env_config["action_bins"]) == 1:
        action_bins = (1, env_config["action_bins"][0])
    else:
        action_bins = env_config["action_bins"]
    update_actions = update_actions.reshape(action_bins)

    # Plot the heatmap
    heat_ax = sns.heatmap(update_actions, ax=heat_ax, linewidths=0,
                          rasterized=rasterize,
                          cbar_kws={"orientation": "horizontal"})
    heat_ax.set(xlabel=xlabel, ylabel=ylabel)

    # Set the heatmap axis tick locations
    num_ticks = 3
    heat_ax.xaxis.set_major_locator(ticker.LinearLocator(num_ticks))
    heat_ax.yaxis.set_major_locator(ticker.LinearLocator(0))

    # Label heatmap axis ticks
    xtick, ytick = get_ticks(num_ticks, 0, env_config, type_="action")

    # Place the heatmaps ticks
    heat_ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    heat_ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    fig.show()
    return fig, bar_ax, heat_ax


def get_ticks(x_num_ticks, y_num_ticks, env_config, type_="state"):
    """
    Factor to get the functions to label the x and y ticks

    Parameters
    ----------
    num_ticks : int
        The total number of ticks along the axis
    env_config : dict
        The environment configuration file as a dict
    type_ : str, optional
        The type of plot to get the ticks of the heatmap for, one of 'state',
        'action', by default 'state'

    Returns
    -------
    2-tuple of func
        The function to generate the x tick labels and the function to
        generate the y tick labels

    Raises
    -------
    NotImplementedError
        If trying to get tick labels for an unsupported environment
    """
    if env_config["env_name"] == "Pendulum-v1":
        return _2D_ticks(x_num_ticks, y_num_ticks, env_config, type_)
    elif env_config["env_name"] == "MountainCarContinuous-v0":
        return _2D_ticks(x_num_ticks, y_num_ticks, env_config, type_)
    else:
        raise NotImplementedError()


def _2D_ticks(x_num_ticks, y_num_ticks, env_config, type_):
    """
    Generates the tick labels for the heatmap on a 2D-state environment

    Parameters
    ----------
    num_ticks : int
        The total number of ticks along the axis
    env_config : dict
        The environment configuration file as a dict
    type_ : str
        The type of plot to get the ticks of the heatmap for, one of 'state',
        'action'

    Returns
    -------
    2-tuple of func
        The function to generate the x tick labels and the function to
        generate the y tick labels
    """
    if type_ == "state":
        ymin, xmin = env_config["state_low"]
        ymax, xmax = env_config["state_high"]
    else:
        ymin, xmin = 0.0, env_config["action_low"][0]
        ymax, xmax = 0.0, env_config["action_high"][0]

    def xtick(x, pos):
        if pos == 0:
            return "{:.2f}".format(xmin)
        elif pos == x_num_ticks - 1:
            return "{:.2f}".format(xmax)
        elif pos == x_num_ticks - (x_num_ticks // 2):
            return "{:.1f}".format((xmax + xmin) / 2)
        else:
            return ""

    def ytick(x, pos):
        if pos == 0:
            return "{:.2f}".format(ymin)
        elif pos == y_num_ticks - 1:
            return "{:.2f}".format(ymax)
        elif pos == y_num_ticks - (y_num_ticks // 2):
            return "{:.1f}".format((ymax + ymin) / 2)
        else:
            return ""

    return xtick, ytick


def state_update_surface_plot(data, figsize=(20, 15),
                              type_="surface", max_=100, smooth=1.5,
                              view=(30, -45), linewidth=0, colour=None,
                              time_ind=-1):
    """
    Plot the surface of states used in each update

    Note that unlike the heatmap, the updates plotted for the surface are not
    normalized betwee 0 and 1.

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    type_ : str, optional
        The type of surface to plot, by default "surface". One of 'surface',
        'wireframe', or 'bar'
    max_ : int, optional
        The cutoff threshold for updates, by default 100. Any states that have
        updates larger than this will be capped at this value.
    smooth : float, optional
        The amount to smoothen the plot, by default 1.5. This value corresponds
        to the standard deviation of a Gaussian convolution kernel applied to
        the surface.
    view : 2-tuple of int, optional
        The view in 3D space, by default (30, -45)
    linewidth : int, optional
        The width of the lines separating cells, by default 0
    colour : str
        The colour to use for the plot, by default None. If None, uses the
        default colours.
    time_ind : int, optional
        In each experiment, the state update heatmap is stored at multiple
        timesteps at certain intervals. This parameter controls which of the
        update heatmaps to plot, with 0 being the first and -1 being the last
        heatmap saved, by default -1.

    Returns
    -------
    plt.figure, plt.Axes3D
        The figure and axis plotten on
    """
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(projection="3d")

    if colour is None:
        colour = DEFAULT_COLOURS[0]

    hp = next(iter(data["experiment_data"]))
    update_states = []
    for run in data["experiment_data"][hp]["runs"]:
        update_states.append(run["update_states"][time_ind])

    update_states = np.stack(update_states).mean(axis=0)
    ind = np.where(update_states > max_)[0]
    update_states[ind] = max_
    update_states = update_states.reshape((100, 100))

    update_states = ndimage.filters.gaussian_filter(update_states,
                                                    sigma=smooth)

    x = np.linspace(-np.pi, np.pi, 100)
    y = np.linspace(-8, 8, 100)
    x, y = np.meshgrid(x, y)

    ax.set_zlim3d(0, max_)
    ax.view_init(view[0], view[1])
    ax.grid(None)

    if type_ == "surface":
        ax.plot_surface(x, y, update_states, antialiased=False,
                        linewidth=linewidth, color=colour)
    elif type_ == "wireframe":
        ax.plot_wireframe(x, y, update_states, antialiased=False,
                          linewidth=linewidth, color=colour)
    elif type_ == "bar":
        # Get the min/max x/y values
        env_config = data["experiment"]["environment"]
        xmin, ymin = env_config["state_low"]
        xmax, ymax = env_config["state_high"]

        # Determine the width, depth, and height of barx
        dx = [(xmax - xmin) / 100] * len(x.ravel())
        dy = [(ymax - ymin) / 100] * len(y.ravel())
        dz = update_states.ravel()

        z = [0] * len(y.ravel())
        ax.bar3d(x.ravel(), y.ravel(), z, dx=dx, dy=dy, dz=dz, color=colour)

    fig.show()
    return fig, ax


def return_distribution(data, type_, hp_ind, bins, figsize=(12, 6), xlim=None,
                        ylim=None, after=0, before=-1):
    """
    Plots the distribution of returns on either an episodic or continuing
    environment

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    type_ : str, optional
        The type of surface to plot, by default "surface". One of 'surface',
        'wireframe', or 'bar'
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    bins : Iterable, int
        The bins to use for the plot. If an Iterable, then each value in the
        Iterable is considered as a cutoff for bins. If an integer, separates
        the returns into that many bins
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    xlim : 2-tuple of float, optional
        The cutoff points for the x-axis to plot between, by default None
    ylim : 2-tuple of float, optional
        The cutoff points for the y-axis to plot between, by default None

    Returns
    -------
    plt.figure, plt.Axes3D
        The figure and axis plotten on
    """
    # Get the episode returns for each run
    run_returns = []
    return_type = type_ + "_episode_rewards"
    for run in data["experiment_data"][hp_ind]["runs"]:
        run_returns.append(np.mean(run[return_type][after:before]))

    title = f"Learning Curve Distribution - HP Settings {hp_ind}"
    return _return_histogram(run_returns, bins, figsize, title, xlim, ylim)


def _return_histogram(run_returns, bins, figsize, title, xlim, ylim, kde=True):
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot()

    ax.set_title(title)
    ax.set_xlabel("Average Return Per Run")
    ax.set_ylabel("Relative Frequency")
    _ = sns.histplot(run_returns, bins=bins, kde=kde)

    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    # Plot relative frequency on the y-axis
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos:
                                 "{:.2f}".format(x / len(run_returns))))

    fig.show()
    return fig, ax


def inds_vs_hp_heatmap(dir_, type_, hp1, hp2, inds, env_config, agent_config,
                       figsize=(12, 8), annotate=True, fmt=".2g",
                       rasterize=False, cmap=CMAP, linewidths=0,
                       env_type="continuing", default_value=-1000, after=0,
                       before=None, fig=None, ax=None, vmin=None, vmax=None):
    """
    Plots the mean returns of one hyperparameter vs another.

    Given a list of hyperparameter indices and two hyperparameter names,
    plots the mean returns generated by varying the second hyperparameter,
    and leaving the first hyperparameter fixed as each hyperparameter index's
    value. The plot is a heatmap. In this way, each row of the heatmap will
    be a certain HP setting with a fixed HP, while the columns will be the
    mean returns given the second HP varying for each HP index.

    Parameters
    ----------
    dir_ : str
        The directory which contains all data dictionaries
    type_ : str
        The type of returns to consider, one of 'train' or 'eval'
    hp1 : str
        The name of the hyperparameter setting to vary on the x-axis
    hp2 : str
        The name of the hyperparameter to plot on the y-axis for each
        hyperparameter settings index plotted
    inds : Iterable of int
        An iterable of hyperparameter indices used to plot
    env_config : dict
        The environment configuration file used to run the experiment
    agent_config : dict
        The agent configuration file used to run the experiment
    figsize : tuple, optional
        The size of the figure, by default (12, 8)
    annotate : bool, optional
        Whether each box in the heatmap should be annotated with its value,
        by default True
    fmt : str, optional
        If annotating, the format to use, by default ".2g"
    rasterize : bool, optional
        Whether or not the final heatmap should be rasterized, by default False
    cmap : str, optional
        The colourmap to use to plot, by default CMAP
    linewidths : int, optional
        The width of the lines used to separate boxes in the heatmap, by
        default 0
    env_type : str, optional
        Whether the environment is continuing or episodic, by default
        "continuing". One of 'episodic', 'continuing'
    default_value : int, optional
        The default value to use for heatmap boxes when no valid second
        hyperparameter can be used for the first hyperparameter, by default
        -1000. This happens, for example in hyperparameter settings that would
        have batch size > replay capacity.
    after : int, optional
        Only consider returns after this episode, by default 0
    before : int, optional
        Only consider returns before this episode, by default None
    fig : plt.Figure, optional
        The figure to plot on, by default None. If None, a new Figure is
        created.
    ax : plt.Axes, optional
        The axis to plot on, by default None. If None, a new Axes is created.
    vmin : float, optional
        Minimum value to anchor the colourmap to, otherwise it is inferred
        from the data.
    vmax : float, optional
        Maximum value to anchor the colourmap to, otherwise it is inferred
        from the data.

    Returns
    -------
    plt.Figure, plt.Axes
        The figure and axis plotted on

    Raises
    ------
    ValueError
        If the figure is not specified but the axis is
        If the type of return considered is not one of "train" or "eval"
        If an illegal environment type is given
    """
    if fig is None and ax is not None:
        raise ValueError("when ax is specified, fig should also be specified")

    if type_ not in ("train", "eval"):
        raise ValueError("type_ must be one of 'train' or 'eval'")

    if env_type not in ("episodic", "continuing"):
        raise ValueError("unknown environment type")

    values = []
    for ind in inds:
        # Get all the combinations of hyperparameters with only hp_name varying
        hp1_varying_only_settings = \
            exp.get_varying_single_hyperparam(env_config, agent_config, hp1)

        # Get the index (into the list of tuples with only hp_name varying) of
        # the tuple which has hp_name varying only, but all other constant
        # hyperparameters equal to those defined by hyperparam setting ind.
        # So hp_combo_constant_hps_equal_ind is the hyperparam combination with
        # all constant hyperparams equal to those in settings number ind.
        hp_combo_constant_hps_equal_ind = \
            next(filter(lambda x: ind in x, hp1_varying_only_settings))
        combo_index = hp1_varying_only_settings.index(
            hp_combo_constant_hps_equal_ind)

        # Get the mean returns for each hyperparameter setting
        mean_returns = []
        for elem in hp1_varying_only_settings[combo_index]:
            if elem is None:
                # No hyperparameter setting for this combo of hyperparams
                # e.g. batch size > buffer size
                # Use the default value for the return
                mean_returns.append(default_value)
                continue

            # Open the data file and compute the mean returns
            agent = agent_config["agent_name"]
            env = env_config["env_name"]
            dir_ = dir_.rstrip("/")
            file_ = dir_ + f"/{env}_{agent}_hp-{elem}.pkl"
            with open(file_, "rb") as in_file:
                data = pickle.load(in_file)

                # Calculate and save the mean return for this HP setting
                mean_return = exp.get_returns(data, type_, elem, env_type)
                if before is not None:
                    mean_return = mean_return[:, after:before].mean()
                else:
                    mean_return = mean_return[:, after:].mean()
                mean_returns.append(mean_return)

        values.append(mean_returns)
    all_mean_returns = np.stack(values)

    # Create the figure and axis
    if fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    # Plot the heatmap and set the labels
    ax = sns.heatmap(all_mean_returns, ax=ax, rasterized=rasterize,
                     annot=annotate, fmt=fmt, cmap=cmap, linewidths=linewidths,
                     vmin=vmin, vmax=vmax)
    ax.set(xlabel=hp1, ylabel=hp2)

    # Define the tick generation function
    def get_ticks(agent_config, hp1, hp2, inds):
        yvalues = []
        for ind in inds:
            params = exp.get_sweep_parameters(agent_config["parameters"],
                                              {"gamma": 0}, ind)[0]
            yvalues.append(params[hp2])

        max_x_ticks = len(agent_config["parameters"][hp1])

        def xticks(x, pos):
            x_tick = agent_config["parameters"][hp1][pos % max_x_ticks]
            return "{:.3g}".format(x_tick)

        def yticks(x, pos):
            y_tick = yvalues[pos]
            return "{:.3g}".format(y_tick)

        return xticks, yticks

    # Place the heatmaps ticks
    xtick, ytick = get_ticks(agent_config, hp1, hp2, inds)
    ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    return fig, ax


def inds_heatmap(dir_, type_, hp1, hp2, inds, env_config, agent_config,
                 figsize=(12, 8), annotate=True, fmt=".2g", rasterize=False,
                 cmap=CMAP, linewidths=0, env_type="continuing",
                 default_value=-1000, after=0, before=None, fig=None, ax=None,
                 vmin=None, vmax=None, div=1):
    """
    Plots the mean returns of each hyperparameter in a 2D matrix.

    Given a 2D array of hyperparameter settings indices, plots the mean return
    of each hyperparameter settings index as a heatmap. This function assumes
    that each row in the matrix corresponds to a different value of one
    hyperparameter, and each column corresponds to the values of a different
    hyperparameter, and these hyperparameters are plotted as the x labels and
    y labels.

    Note that any hyperparameter setting less than 0 is ignored, so these
    indices can be used to plot the default value in the heatmap.

    Parameters
    ----------
    dir_ : str
        The directory which contains all data dictionaries
    type_ : str
        The type of returns to consider, one of 'train' or 'eval'
    hp1 : str
        The name of the hyperparameter setting to vary on the x-axis
    hp2 : str
        The name of the hyperparameter to plot on the y-axis for each
        hyperparameter settings index plotted
    inds : 2D array_like of int
        Matrix of hyperparameter indices to plot
    env_config : dict
        The environment configuration file used to run the experiment
    agent_config : dict
        The agent configuration file used to run the experiment
    figsize : tuple, optional
        The size of the figure, by default (12, 8)
    annotate : bool, optional
        Whether each box in the heatmap should be annotated with its value,
        by default True
    fmt : str, optional
        If annotating, the format to use, by default ".2g"
    rasterize : bool, optional
        Whether or not the final heatmap should be rasterized, by default False
    cmap : str, optional
        The colourmap to use to plot, by default CMAP
    linewidths : int, optional
        The width of the lines used to separate boxes in the heatmap, by
        default 0
    env_type : str, optional
        Whether the environment is continuing or episodic, by default
        "continuing". One of 'episodic', 'continuing'
    default_value : int, optional
        The default value to use for heatmap boxes when no valid second
        hyperparameter can be used for the first hyperparameter, by default
        -1000. This happens, for example in hyperparameter settings that would
        have batch size > replay capacity.
    after : int, optional
        Only consider returns after this episode, by default 0
    before : int, optional
        Only consider returns before this episode, by default None
    fig : plt.Figure, optional
        The figure to plot on, by default None. If None, a new Figure is
        created.
    ax : plt.Axes, optional
        The axis to plot on, by default None. If None, a new Axes is created.
    vmin : float, optional
        Minimum value to anchor the colourmap to, otherwise it is inferred
        from the data.
    vmax : float, optional
        Maximum value to anchor the colourmap to, otherwise it is inferred
        from the data.

    Returns
    -------
    plt.Figure, plt.Axes
        The figure and axis plotted on

    Raises
    ------
    ValueError
        If the figure is not specified but the axis is
        If the type of return considered is not one of "train" or "eval"
        If an illegal environment type is given
    """
    # Ensure the hyperparameter settings have consistent hp values in each
    # row and column
    _check_index_matrix(agent_config, env_config, inds, hp2)
    _check_index_matrix(agent_config, env_config, inds.transpose(), hp1)

    if fig is None and ax is not None:
        raise ValueError("when ax is specified, fig should also be specified")

    if type_ not in ("train", "eval"):
        raise ValueError("type_ must be one of 'train' or 'eval'")

    if env_type not in ("episodic", "continuing"):
        raise ValueError("unknown environment type")

    values = np.zeros(inds.shape)
    for row in range(inds.shape[0]):
        for ind in range(inds[row].shape[0]):
            if inds[row, ind] < 0:
                values[row, ind] = default_value
                continue

            # Open the data file and compute the mean returns
            agent = agent_config["agent_name"]
            env = env_config["env_name"]
            dir_ = dir_.rstrip("/")
            index = inds[row, ind]
            file_ = dir_ + f"/{env}_{agent}_hp-{index}.pkl"
            with open(file_, "rb") as in_file:
                data = pickle.load(in_file)

                # Calculate and save the mean return for this HP setting
                mean_return = exp.get_avg_returns(data, type_, inds[row, ind],
                                                  after, before)

            values[row, ind] = mean_return.mean() / div

    # Create the figure and axis
    if fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    # Plot the heatmap and set the labels
    ax = sns.heatmap(values, ax=ax, rasterized=rasterize,
                     annot=annotate, fmt=fmt, cmap=cmap, linewidths=linewidths,
                     vmin=vmin, vmax=vmax)
    ax.set(xlabel=hp1, ylabel=hp2)

    # Define the tick generation function
    def get_ticks(agent_config, hp1, hp2, inds):
        yvalues = []
        inds = inds.ravel()
        for ind in range(inds.size):
            params = exp.get_sweep_parameters(agent_config["parameters"],
                                              {"gamma": 0}, inds[ind])[0]
            if not params[hp2] in yvalues:
                yvalues.append(params[hp2])

        max_x_ticks = len(agent_config["parameters"][hp1])

        def xticks(x, pos):
            x_tick = agent_config["parameters"][hp1][pos % max_x_ticks]
            return "{:.3g}".format(x_tick)

        def yticks(x, pos):
            y_tick = yvalues[pos]
            return "{:.3g}".format(y_tick)

        return xticks, yticks

    # Place the heatmaps ticks
    xtick, ytick = get_ticks(agent_config, hp1, hp2, inds)
    ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    return fig, ax


def _check_index_matrix(agent_config, env_config, inds, hp):
    """
    Check the columns of a matrix of hyperparameters indices for consistent
    hyperparameter values.

    Checks each row of a matrix of hyperparameter settings to ensure that each
    hyperparameter setting in the row has the same value for a certain
    hyperparameter. This is needed to ensure consistency when plotting a
    matrix of hyperparameter settings as a heatmap.

    Note that any hyperparameter setting less than 0 is ignored, so these
    indices can be used to plot the default value in the heatmap.

    Parameters
    ----------
    agent_config : dict
        The agent configuration file used to run the experiment
    env_config : dict
        The environment configuration file used to run the experiment
    inds : 2D array_like of int
        Matrix of hyperparameter indices to check to ensure that each row
        has hyperparameter settings that have the same hyperparameter value
        for a specific hyperparameter
    hp : str
        The hyperparameter to ensure is consistent between hyperparameter
        settings on each row.

    Raises
    ------
    IndexError
        If the hyperparameter indices do not refer to settings with the same
        value for the argument hyperparameter.
    ValueError
        If the array is not 2D.
    """
    if len(inds.shape) != 2:
        raise ValueError("inds must be a matrix")

    parameters = agent_config["parameters"]
    for row in range(inds.shape[0]):
        # The value that all rows should have as HP value
        hp_reference = None
        for col in range(inds.shape[1]):
            index = inds[row, col]
            if index < 0:
                # No hyperparameter setting in this position
                continue

            # Get the value of the hyperparameter for the index in the current
            # position
            hp_value = exp.get_sweep_parameters(parameters, env_config, index)
            hp_value = hp_value[0][hp]

            if hp_reference is None:
                # This is the first HP value encountered, use it as a reference
                # value for all other HPs in the row
                hp_reference = hp_value
                continue

            if hp_value != hp_reference:
                # Illegal HP value
                raise IndexError("hyperparameter indices do not share the " +
                                 f"same {hp} value")


@deprecated(reason="no longer storing states this way")
def absolute_state_over_time(data, after=0, figsize=(12, 6), state_ind=0,
                             colour=None, alpha=0.2, title="", xlabel="",
                             ylabel="", xlim=None, ylim=None):
    """
    Plot the average absolute value of a state feature per episode

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    after : int, optional
        Consider the state features only after this timestep for each
        episode, by default 0
    figsize : tuple, optional
        The size of the figure, by default (12, 6)
    state_ind : int, optional
        The index of the state feature to plot, by default 0
    colour : str, optional
        The colour for the plot, by default None
    alpha : float, optional
        The alpha value for the individual runs, by default 0.1
    title : str, optional
        The title for the plot, by default ""
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    xlim : 2-tuple of float, optional
        The limits for the x axis, by default None
    ylim : 2-tuple of float, optional
        The limits for the y axis, by default None

    Returns
    -------
    plt.figure, plt.Axes
        The figure and Axes plotted on
    """
    if colour is None:
        colour = DEFAULT_COLOURS[0]

    hp = next(iter(data["experiment_data"]))
    states = []
    for run in data["experiment_data"][hp]["runs"]:
        states.append(run["train_states"])

    states = np.stack(states)[:, :, after:]
    states = np.abs(states)

    # Calculate the average performance, averaging over runs and then
    # timesteps to get the average angle per episode
    avg = states.mean(axis=0)
    avg = avg.mean(axis=1)

    # Get the average angle per episode for each run
    states = states.mean(axis=2)

    # Get some prelimiary variables for the legend
    legend_labels = []
    legend_lines = []
    agent_name = data["experiment"]["agent"]["agent_name"]

    # Get the mean over all runs and timesteps for each episode
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot()

    # Set axis limits
    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    # Plot the average angle per episode for all runs
    ax.plot(avg[:, state_ind], color=colour, linestyle="-")
    label = f"mean {agent_name} - Hyperparameter Setting {hp}"
    legend_labels.append(label)
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="-"))

    # Plot average average absolute angle for each episode
    # (averaged over timesteps) for each separate run
    for i in range(states.shape[0]):
        ax.plot(states[i, :, state_ind], linestyle="--", color=colour,
                alpha=alpha)
    legend_labels.append("individual runs")
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="--",
                               alpha=alpha))

    ax.legend(legend_lines, legend_labels)

    ax.set_title(title)
    ax.set_ylabel(ylabel)
    ax.set_xlabel(xlabel)

    return fig, ax


def _get_default_colours(iter_):
    """
    Recursively turns elements of an Iterable into strings representing
    colours.

    This function will turn each element of an Iterable into strings that
    represent colours, recursively. If the elements of an Iterable are
    also Iterable, then this function will recursively descend all the way
    through every Iterable until it finds an Iterable with non-Iterable
    elements. These elements will be replaced by strings that represent
    colours. In effect, this function keeps the data structure, but replaces
    non-Iterable elements by strings representing colours. Note that this
    funcion assumes that all elements of an Iterable are of the same type,
    and so it only checks if the first element of an Iterable object is
    Iterable or not to stop the recursion.

    Parameters
    ----------
    iter_ : collections.Iterable
        The top-level Iterable object to turn into an Iterable of strings of
        colours, recursively.

    Returns
    -------
    list of list of ... of strings
        A data structure that has the same architecture as the input Iterable
        but with all non-Iterable elements replaced by strings.
    """
    colours = []

    # Calculate the number of lists at the current level to go through
    paths = range(len(iter_))

    # Return a list of colours if the elements of the list are not lists
    if not isinstance(iter_[0], Iterable):
        global OFFSET
        col = [DEFAULT_COLOURS[(OFFSET + i) % len(DEFAULT_COLOURS)]
               for i in paths]
        OFFSET += len(paths)
        return col

    # For each list at the current level, get the colours corresponding to
    # this level
    for i in paths:
        colours.append(_get_default_colours(iter_[i]))

    return colours


def _plot_shaded(ax, x, y, region, colour, label, alpha):
    """
    Plots a curve with a shaded region.

    Parameters
    ----------
    ax : plt.Axes
        The axis to plot on
    x : Iterable
        The points on the x-axis
    y : Iterable
        The points on the y-axis
    region : list or array_like
        The region to shade about the y points. The shaded region will be
        y +/- region. If region is a list or 1D np.ndarray, then the region
        is used both for the + and - portions. If region is a 2D np.ndarray,
        then the first row will be used as the lower bound (-) and the
        second row will be used for the upper bound (+). That is, the region
        between (lower_bound, upper_bound) will be shaded, and there will be
        no subtraction/adding of the y-values.
    colour : str
        The colour to plot with
    label : str
        The label to use for the plot
    alpha : float
        The alpha value for the shaded region
    """
    if colour is None:
        colour = DEFAULT_COLOURS[0]

    ax.plot(x, y, color=colour, label=label)
    if type(region) == list:
        ax.fill_between(x, y-region, y+region, alpha=alpha, color=colour)
    elif type(region) == np.ndarray and len(region.shape) == 1:
        ax.fill_between(x, y-region, y+region, alpha=alpha, color=colour)
    elif type(region) == np.ndarray and len(region.shape) == 2:
        ax.fill_between(x, region[0, :], region[1, :], alpha=alpha,
                        color=colour)


def _setup_fig(fig, ax, figsize=None, title=None, xlim=None, ylim=None,
               xlabel=None, ylabel=None, xscale=None, yscale=None, xbase=None,
               ybase=None):
    if fig is None:
        if ax is not None:
            raise ValueError("Must specify figure when axis given")
        if figsize is not None:
            fig = plt.figure(figsize=figsize)
        else:
            fig = plt.figure()

    if ax is None:
        ax = fig.add_subplot()

    if title is not None:
        ax.set_title(title)

    if xlabel is not None:
        ax.set_xlabel(xlabel)

    if ylabel is not None:
        ax.set_ylabel(ylabel)

    if xlim is not None:
        ax.set_xlim(xlim)

    if ylim is not None:
        ax.set_ylim(ylim)

    if xscale is not None:
        if xbase is not None:
            ax.set_xscale(xscale, base=xbase)
        else:
            ax.set_xscale(xscale)

    if yscale is not None:
        if ybase is not None:
            ax.set_yscale(yscale, base=ybase)
        else:
            ax.set_yscale(yscale)

    return fig, ax


def reset():
    """
    Resets the colours offset
    """
    global OFFSET
    OFFSET = 0


if __name__ == "__main__":
    print(_get_default_colours([[79], [58]]))

<filename>utils/ERLplot_utils.py
#!/usr/bin/env python3

# Import modules
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib import ticker, gridspec
from deprecated import deprecated
import utils.ERLexperiment_utils as exp
import numpy as np
from scipy import ndimage
import seaborn as sns
from collections import Iterable
from utils.tile_coder import TileCoder
import matplotlib as mpl
import pickle


# TODO: Add documentation for plot_updates and plot_updates_
# TODO: Update action plots
# TODO: state_update_heatmap should do 2D heatmap of multi-dimensional
#       states, just give the two dimensions to plot in a heatmap. same for the
#       surface plot!!
# TODO: separate state_update_heatmap into two functions, one of which will
#   simply take a vector of update counts and create a heatmap of it given the
#   dimensions to reduce it to (e.g. reshape(100, 100))

# TODO: make _pendulum_ticks more general to label any number of ticks

# TODO: Plot mean with runs episodic
# TODO: Plot mean with stderr episodic for multiple data files
# TODO: Change action plot to histogram
# TODO: Make the state surface plot use either state updates or states visited
#       Like the new scheme in state heatmap plot
# TODO: State surface plot should take similar params to heatmap
# TODO: Improve documentation
# TODO: plot_mean_with_runs() should plot avg rewards for continuing


# Global variables
CMAP = "rocket"
DEFAULT_COLOURS = list(sns.color_palette(CMAP, 6).as_hex())
plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=sns.color_palette(CMAP))
# sns.set_theme(palette=CMAP)
OFFSET = 0  # The offset to start in DEFAULT_COLOURS


def updates(data, type_, ind, smooth_over, names, fig=None, ax=None,
            figsize=(12, 6), last_ind=-1, colours=None, xlim=None, ylim=None,
            env_type="continuing", label="", alpha=0.2, xscale="log",
            yscale="linear", base=10, xlabel=None, ylabel=None,
            runs=False, keep_shape=False, bootstrap=False, significance=0.1):
    """
    Plots the performance vs number of samples used in updates

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    names : list of str
        The name of the agents, used for the legend
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    colours : list of list of str
        The colours to use for each hyperparameter settings plot for each data
        dictionary
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    xscale : str, optional
        The scale at which to plot the x axis, by default "log"
    yscale : str, optional
        The scale at which to plot the y axis, by default "linear"
    base : int, optional
        The base of the scale of the x axis, by default 10
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    title : str, optional
        The title of the plot, by default ""
    runs : bool, optional
        Whether or not to plot the individual runs or error bars, by default
        False. If True, plots individual runs instead of error bars.
    keep_shape : bool, optional
        Whether the shape of the input data should be kept after smoothing, by
        default False. If True, then the first smooth_over episodes are
        smoothed and kept. If False, then the first smooth_over episodes are
        replaced by their fully smoothed value.
    bootstrap : bool, optional
        Whether or plot standard error or bootstrap confidence intervals, by
        default False. If True, bootstrap confidence intervals are plotted
        instead of standard error.
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.Figure, plt.Axes
        The figure and axes plotted on
    """
    if colours is None:
        colours = _get_default_colours(ind)

    # Setup figure
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel,
                         xscale, yscale, base)

    legend_labels = []
    for i in range(len(data)):
        for j in range(len(ind[i])):
            fig, ax = _updates(data[i], type_, ind[i][j], smooth_over[i],
                               fig, ax, figsize,
                               colours[i][j], xlim, ylim, env_type,
                               alpha, xscale, yscale, base, xlabel, ylabel,
                               runs, keep_shape, bootstrap, significance)

            # Create legend label and lines
            agent_name = names[i]
            label = agent_name
            legend_labels.append(label)

    ax.legend(legend_labels)

    return fig, ax


def _updates(data, type_, ind, smooth_over, fig=None, ax=None,
             figsize=(12, 6), colour=None, xlim=None, ylim=None,
             env_type="continuing", alpha=0.2, xscale="log",
             yscale="linear", base=10, xlabel=None, ylabel=None, runs=False,
             keep_shape=False, bootstrap=False, significance=0.1):
    if colour is None:
        colour = _get_default_colours([ind])[0]

    # # Setup figure
    # fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel,
    #                      xscale, yscale, base)

    hyperparams = data["experiment_data"][ind]["agent_hyperparams"]
    if "batch_size" in hyperparams:
        batch_size = hyperparams["batch_size"]
    else:
        batch_size = 1

    # Calculate updates to plot along x-axis
    steps_per_episode = data["experiment"]["environment"]["steps_per_episode"]
    total_steps = data["experiment"]["environment"]["total_timesteps"]

    # Include the first eval -- before any updates have been done
    start = 0
    stop = (total_steps + steps_per_episode) * batch_size
    step = steps_per_episode * batch_size
    updates = np.arange(start, stop, step)

    # Get the data to plot depending on the env type
    if env_type == "continuing":
        rewards, err, all_ = \
            rewards_per_update_continuing(data, ind, smooth_over, type_,
                                          keep_shape, bootstrap, significance)
    else:
        rewards, err, all_ = \
            rewards_per_update_episodic(data, ind, smooth_over, updates, type_,
                                        keep_shape, bootstrap, significance)

    # If the original shape is not kept, then the updates need to be adjusted
    # for smoothing
    if not keep_shape:
        updates = updates[:len(updates) - smooth_over + 1]

    # Plot the mean
    ax.step(updates, rewards, where="post", color=colour)
    if not runs:
        if not bootstrap:
            # Standard error
            ax.fill_between(updates, rewards-err, rewards+err, alpha=alpha,
                            step="post", color=colour)
        else:
            # Bootstrap confidence intervals
            ax.fill_between(updates, err[0], err[1], alpha=alpha, step="post",
                            color=colour)

    # Plot each run
    if runs:
        for i in range(all_.shape[0]):
            ax.step(updates, all_[i, :], linestyle="--", alpha=alpha,
                    color=colour, where="post")

    return fig, ax


def rewards_per_update_continuing(data, ind, smooth_over, type_="train",
                                  keep_shape=False, bootstrap=False,
                                  significance=0.1):
    returns = []
    for run in data["experiment_data"][ind]["runs"]:
        if type_ == "train":
            # The episodic returns for each run
            run_returns = list(run["train_episode_rewards"])

            # Plot the first eval episode, since this is the only data for
            # 0 updates
            first_ep = run["eval_episode_rewards"][0].Vmean()
            run_returns.insert(0, first_ep)

            returns.append(run_returns)

        else:
            returns.append(run["eval_episode_rewards"].mean(axis=-1))

    # Get the steps per episode so we can plot average reward
    steps_per_ep = data["experiment"]["environment"]["steps_per_episode"]
    returns = np.array(returns) / steps_per_ep

    returns = exp.smooth(returns, smooth_over, keep_shape=keep_shape)

    # Calculate standard error
    if not bootstrap:
        runs = len(data["experiment_data"][ind]["runs"])
        stderr = np.std(np.array(returns), axis=0) / np.sqrt(runs)
        return returns.mean(axis=0), stderr, returns

    # Calculate bootstrap confidence interval
    else:
        conf = exp.bootstrap_conf(returns, significance)
        return returns.mean(axis=0), conf, returns


def rewards_per_update_episodic(data, ind, smooth_over, updates, type_="train",
                                keep_shape=False, bootstrap=False,
                                significance=0.1):
    returns = []
    for run in data["experiment_data"][ind]["runs"]:
        if type_ == "train":
            # Determine the size of batches used in updates
            hyperparams = data["experiment_data"][ind]["agent_hyperparams"]
            if "batch_size" in hyperparams:
                batch_size = hyperparams["batch_size"]
            else:
                batch_size = 1

            # Calculate the cumulative number of updates over all previous
            # episodes
            run_updates = np.cumsum(run["train_episode_steps"]) * batch_size

            # For each cumulative number of updates plotted on the x-axis,
            # get the performance of the episode with number of updates closest
            # to this value
            run_returns = []
            for i in range(1, len(updates)):
                ep_gt_lower_bound = np.where(run_updates >= updates[i])[0][0]
                episode = run["train_episode_rewards"][ep_gt_lower_bound]
                run_returns.append(episode)

            # Plot the first eval episode, since this is the only data for
            # 0 updates
            first_ep = run["eval_episode_rewards"][0].mean()
            run_returns.insert(0, first_ep)
            returns.append(run_returns)

        elif type_ == "eval":
            returns.append(run["eval_episode_rewards"].mean(axis=-1))

    # # Smooth data
    returns = np.array(returns)
    returns = exp.smooth(returns, smooth_over, keep_shape=keep_shape)

    # Calculate standard error
    if not bootstrap:
        runs = len(data["experiment_data"][ind]["runs"])
        stderr = np.std(np.array(returns), axis=0) / np.sqrt(runs)
        return returns.mean(axis=0), stderr, returns

    # Calculate bootstrap confidence interval
    else:
        conf = exp.bootstrap_conf(returns, significance)
        return returns.mean(axis=0), conf, returns


def mean_with_bootstrap_conf(data, type_, ind, smooth_over, names,
                             fig=None, ax=None, figsize=(12, 6),
                             xlim=None, ylim=None, alpha=0.2,
                             colours=None, env_type="continuing",
                             significance=0.01, keep_shape=False,
                             xlabel=None, ylabel=None):
    """
    Plots the average training or evaluation return over all runs with
    confidence intervals.

    Given a list of data dictionaries of the form returned by main.py, this
    function will plot each episodic return for the list of hyperparameter
    settings ind each data dictionary. The ind argument is a list, where each
    element is a list of hyperparameter settings to plot for the data
    dictionary at the same index as this list. For example, if ind[i] = [1, 2],
    then plots will be generated for the data dictionary at location i
    in the data argument for hyperparameter settings ind[i] = [1, 2].
    The smooth_over argument tells how many previous data points to smooth
    over

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    names : list of str
        The name of the agents, used for the legend
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of list of str
        The colours to use for each hyperparameter settings plot for each data
        dictionary
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot
    """
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel)

    # Set the colours to be default if not set
    if colours is None:
        colours = _get_default_colours(ind)

    # Track the total timesteps per hyperparam setting over all episodes and
    # the cumulative timesteps per episode per data dictionary (timesteps
    # should be consistent between all hp settings in a single data dict)
    total_timesteps = []
    cumulative_timesteps = []

    for i in range(len(data)):
        if type_ == "train":
            cumulative_timesteps.append(exp.get_cumulative_timesteps(data[i]
                                        ["experiment_data"][ind[i][0]]["runs"]
                                        [0]["train_episode_steps"]))
        elif type_ == "eval":
            cumulative_timesteps.append(data[i]["experiment_data"][ind[i][0]]
                                        ["runs"][0]["timesteps_at_eval"])
        else:
            raise ValueError("type_ must be one of 'train', 'eval'")
        total_timesteps.append(cumulative_timesteps[-1][-1])

    # Find the minimum of total trained-for timesteps. Each plot will only
    # be plotted on the x-axis until this value
    min_timesteps = min(total_timesteps)

    # For each data dictionary, find the minimum index where the timestep at
    # that index is >=  minimum timestep
    ind_ge_min_timesteps = []
    for cumulative_timesteps_per_data in cumulative_timesteps:
        final_ind = np.where(cumulative_timesteps_per_data >=
                             min_timesteps)[0][0]
        # Since indexing will stop right before the minimum, increment it
        ind_ge_min_timesteps.append(final_ind + 1)

    # Plot all data for all HP settings, only up until the minimum index
    # fig, ax = None, None
    if env_type == "continuing":
        plot_fn = _plot_mean_with_conf_continuing
    else:
        plot_fn = _plot_mean_with_conf_episodic

    for i in range(len(data)):
        fig, ax = \
            plot_fn(data=data[i], type_=type_,
                    ind=ind[i], smooth_over=smooth_over[i], name=names[i],
                    fig=fig, ax=ax, figsize=figsize, xlim=xlim, ylim=ylim,
                    last_ind=ind_ge_min_timesteps[i], alpha=alpha,
                    colours=colours[i], significance=significance,
                    keep_shape=keep_shape)

    return fig, ax


def _plot_mean_with_conf_continuing(data, type_, ind, smooth_over, fig=None,
                                    ax=None, figsize=(12, 6), name="",
                                    last_ind=-1, xlabel="Episodes",
                                    ylabel="Average Return", xlim=None,
                                    ylim=None, alpha=0.2, colours=None,
                                    significance=0.01, keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a single
    data dictionary on a continuing environment. Bootstrap confidence intervals
    are plotted as shaded regions.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    timestep_multiply : int, optional
        A value to multiply each timstep by, by default 1. This is useful if
        your agent does multiple updates per timestep and you want to plot
        performance vs. number of updates.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    # This should be the exact same as the episodic version except without
    # reducing the episodes. Follow the same structure as the episodic function
    # and the continuing function with standard error.
    raise NotImplementedError


def _plot_mean_with_conf_episodic(data, type_, ind, smooth_over, fig=None,
                                  ax=None, figsize=(12, 6), name="",
                                  last_ind=-1, xlabel="Episodes",
                                  ylabel="Average Return", xlim=None,
                                  ylim=None, alpha=0.2, colours=None,
                                  significance=0.01, keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a single
    data dictionary on an episodic environment. Bootstrap confidence intervals
    are plotted as shaded regions.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    if colours is not None and len(colours) != len(ind):
        raise ValueError("must have one colour for each hyperparameter " +
                         "setting")

    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    # if ax is None and fig is None:
    #     fig = plt.figure(figsize=figsize)
    #     ax = fig.add_subplot()

    # if xlim is not None:
    #     ax.set_xlim(xlim)
    # if ylim is not None:
    #     ax.set_ylim(ylim)

    conf_level = "{:.2f}".format(1-significance)
    title = f"""Average {type_.title()} Return per Run with {conf_level}
                 Confidence Intervals"""
    fig, ax = _setup_fig(fig, ax, figsize, title, xlim, ylim, xlabel, ylabel)

    # Plot with bootstrap confidence interval
    for i in range(len(ind)):
        data = exp.reduce_episodes(data, ind[i], type_=type_)
        _, mean, conf = exp.get_mean_err(data, type_, ind[i], smooth_over,
                                         exp.bootstrap_conf,
                                         err_args={
                                            "significance": significance,
                                         },
                                         keep_shape=keep_shape)

        mean = mean[:last_ind]
        conf = conf[:, :last_ind]

        episodes = np.arange(mean.shape[0])

        # Plot based on colours
        label = name
        print(mean.shape, conf.shape, episodes.shape)
        _plot_shaded(ax, episodes, mean, conf, colours[i], label, alpha)

    ax.legend()
    conf_level = "{:.2f}".format(1-significance)
    ax.set_title(f"""Average {type_.title()} Return per Run with {conf_level}
                 Confidence Intervals""")
    # ax.set_ylabel(ylabel)
    # ax.set_xlabel(xlabel)

    fig.show()
    return fig, ax


@deprecated(reason="""Use exp.get_mean_err with exp.bootstrap_conf""")
def get_mean_bootstrap_conf(data, type_, ind, smooth_over,
                            env_type="continuing", significance=0.01,
                            keep_shape=False):
    """
    Gets the mean return per episode along with the bootstrap confidence
    interval on an episode-by-episode basis over all runs.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : int
        The hyperparameter setting to plot
    smooth_over : int
        The number of previous returns to smooth over. If less than 2, then
        no smoothing is performed
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    significance : float, optional
        The significance level for the confidence interval, by default 0.01

    Returns
    -------
    3-tuple of array_like
        The timesteps at each return to plot, the mean return to plot, and
        the confidence interval for each episode to plot as a shaded region
    """
    timesteps = None  # So the linter doesn't have a temper tantrum

    # Determine the timesteps to plot at
    if type_ == "eval":
        timesteps = \
            data["experiment_data"][ind]["runs"][0]["timesteps_at_eval"]

    elif type_ == "train":
        timesteps_per_ep = \
            data["experiment_data"][ind]["runs"][0]["train_episode_steps"]
        timesteps = exp.get_cumulative_timesteps(timesteps_per_ep)

    # Get the mean over all episodes per evaluation step (for online
    # returns, this axis will have length 1 so we squeeze it)
    returns = exp.get_returns(data, type_, ind, env_type=env_type)
    returns = returns.mean(axis=-1)

    # Smooth over returns
    returns = exp.smooth(returns, smooth_over, keep_shape=keep_shape)

    # Get the standard deviation of mean episodes per evaluation
    # step over all runs
    conf = exp.bootstrap_conf(returns, significance)

    # Get the mean over all runs
    mean = returns.mean(axis=0)

    # If we do not keep the shape, then due to smoothing we lose some
    # episodes
    if not keep_shape:
        end = len(timesteps) - smooth_over + 1
        timesteps = timesteps[:end]

    return timesteps, mean, conf


def plot_mean_with_runs(data, type_, ind, smooth_over, names, colours=None,
                        figsize=(12, 6), xlim=None, ylim=None, alpha=0.2,
                        plot_avg=True, env_type="continuing",
                        keep_shape=False, fig=None, ax=None, detrend=None):
    """
    Plots the mean return over all runs and the return for each run for a list
    of data dictionaries and hyperparameter indices

    Plots both the mean return per episode (over runs) as well as the return
    for each individual run (including "mini-runs", if a set of concurrent
    episodes were run for all runs, e.g. multiple evaluation episodes per
    run at set intervals)

    Note that this function takes in a list of data dictionaries and will
    plot the runs for each ind (which is a list of lists, where each super-list
    refers to a data dictionary and each sub-list refers to the indices for
    that data dictionary to plot).

    Example
    -------
    plot_mean_with_runs([sac_data, linear_data], "train", [[3439], [38, 3]],
    smooth_over=[5, 2], names=["SAC", "LAC"], figsize=(12, 6), alpha=0.2,
    plot_avg=True, env_type="episodic")

    will plot hyperparameter index 3439 for the sac_data, smoothing over the
    last 5 episodes, and the label will have the term "SAC" in it; also plots
    the mean and each individual run on the linear_data for hyperparameter
    settings 38 and 3, smoothing over the last 2 episodes for each and with
    the term "LAC" in the labels.

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    figsize : tuple(int, int)
        The size of the figure to plot
    names : list of str
        The name of the agents, used for the legend
    colours : list of list of str, optional
        The colours to use for each hyperparameter settings plot for each data
        dictionary, by default None
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha to use for plots of the runs, by default 0.1
    plot_avg : bool, optional
        If concurrent episodes are executed in each run (e.g. multiple
        evaluation episodes are run at set intervals), then whether to plot the
        performance of each separately or the average performance over all
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    fig : plt.Figure
        The figure to plot on
    ax : plt.Axes
        The axis to plot on
    detrend : str, optional
        How to detrend the runs for each data, by default None. One of
        "difference", "linear", or "mean"

    Returns
    -------
    tuple of plt.figure, plt.Axes
        The figure and axis plotted on
    """
    # Set the colours to be default if not set
    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    # fig = plt.figure(figsize=figsize)
    # ax = fig.add_subplot()
    if env_type == "continuing":
        xlabel = "Timesteps"
        ylabel = "Average Reward"
    else:
        xlabel = "Episodes"
        ylabel = "Return"
    title = f"Mean Return with Runs"

    fig, ax = _setup_fig(fig, ax, figsize, title, xlim, ylim, xlabel, ylabel)

    # Plot for each data dictionary given
    legend_lines = []
    legend_labels = []
    for i in range(len(data)):
        for _ in range(len(ind[i])):
            fig, ax, labels, lines = \
                _plot_mean_with_runs(data[i], type_, ind[i], smooth_over[i],
                                     names[i], colours[i], figsize, xlim, ylim,
                                     alpha, plot_avg, env_type, fig, ax,
                                     keep_shape, detrend)

            legend_lines.extend(lines)
            legend_labels.extend(labels)

    ax.legend(legend_lines, legend_labels)
    fig.show()

    return fig, ax


def _plot_mean_with_runs(data, type_, ind, smooth_over, name, colours=None,
                         figsize=(12, 6), xlim=None, ylim=None, alpha=0.2,
                         plot_avg=True, env_type="continuing", fig=None,
                         ax=None, keep_shape=False, detrend=None):
    """
    Plots the mean return over all runs and the return for each run for a
    single data dictionary and for each in a list of hyperparameter settings.

    Similar to plot_mean_with_runs, except that this function takes in only
    a single data dictionary.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot for
        each agent. For example [1, 2] means that the agent plots will use
        hyperparameter settings indices 1 and 2.
    smooth_over : int
        The number of previous data points to smooth over for the agent's
        plot. Note that this is *not* the number of timesteps to smooth over,
        but rather the number of data points to smooth over. For example,
        if you save the return every 1,000 timesteps, then setting this value
        to 15 will smooth over the last 15 readings, or 15,000 timesteps.
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str
        The name of the agents, used for the legend
    colours : list of list of str, optional
        The colours to use for each hyperparameter settings plot for each data
        dictionary, by default None
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha to use for plots of the runs, by default 0.1
    plot_avg : bool, optional
        If concurrent episodes are executed in each run (e.g. multiple
        evaluation episodes are run at set intervals), then whether to plot the
        performance of each separately or the average performance over all
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    detrend : str, optional
        How to detrend the runs for each data, by default None. One of
        "difference", "linear", or "mean"

    Returns
    -------
    tuple of plt.figure, plt.Axes, list of str, list of mpl.Lines2D
        The figure and axis plotted on as well as the list of strings to use
        as labels and the list of lines to include in the legend
    """
    # Set up figure and axis
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim)

    if colours is None:
        colours = _get_default_colours(ind)

    # Store the info to keep in the legend
    legend_labels = []
    legend_lines = []

    # Plot each selected hyperparameter setting in the data dictionary
    for j in range(len(ind)):
        fig, ax, labels, lines = \
            _plot_mean_with_runs_single_hp(data, type_, ind[j], smooth_over,
                                           name, colours[j], figsize, xlim,
                                           ylim, alpha, plot_avg, env_type,
                                           fig, ax, keep_shape, detrend)
        legend_labels.extend(labels)
        legend_lines.extend(lines)

    return fig, ax, legend_labels, legend_lines


def _plot_mean_with_runs_single_hp(data, type_, ind, smooth_over, name,
                                   colour=None, figsize=(12, 6), xlim=None,
                                   ylim=None, alpha=0.2, plot_avg=True,
                                   env_type="continuing", fig=None, ax=None,
                                   keep_shape=False, detrend=None):
    """
    Plots the mean return over all runs and the return for each run for a
    single data dictionary and a single hyperparameter setting.

    Similar to _plot_mean_with_runs, except that this function takes in only
    a single data dictionary and a single hyperparameter setting.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : int
        The hyperparameter settings indices to plot for the agent. For example
        5 means that the agent plots will use hyperparameter settings index 5.
    smooth_over : int
        The number of previous data points to smooth over for the agent's
        plot. Note that this is *not* the number of timesteps to smooth over,
        but rather the number of data points to smooth over. For example,
        if you save the return every 1,000 timesteps, then setting this value
        to 15 will smooth over the last 15 readings, or 15,000 timesteps.
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str
        The name of the agents, used for the legend
    colours : list of list of str, optional
        The colours to use for each hyperparameter settings plot for each data
        dictionary, by default None
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha to use for plots of the runs, by default 0.1
    plot_avg : bool, optional
        If concurrent episodes are executed in each run (e.g. multiple
        evaluation episodes are run at set intervals), then whether to plot the
        performance of each separately or the average performance over all
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'
    detrend : str, optional
        How to detrend the runs for each data, by default None. One of
        "difference", "linear", or "mean"

    Returns
    -------
    tuple of plt.figure, plt.Axes, list of str, list of mpl.Lines2D
        The figure and axis plotted on as well as the list of strings to use
        as labels and the list of lines to include in the legend
    """
    # Set up figure and axis
    fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim)

    if colour is None:
        colour = _get_default_colours([ind])[0]

    # Determine the timesteps to plot at
    if type_ == "eval":
        timesteps = \
            data["experiment_data"][ind]["runs"][0]["timesteps_at_eval"]

    elif type_ == "train":
        timesteps_per_ep = \
            data["experiment_data"][ind]["runs"][0]["train_episode_steps"]
        timesteps = exp.get_cumulative_timesteps(timesteps_per_ep)

    # Plot the average reward
    if env_type == "continuing":
        episode_steps = data["experiment"]["environment"]["steps_per_episode"]

    # Get returns
    all_returns = exp.get_returns(data, type_, ind, env_type)

    # If concurrent episodes are run in each run then average them if
    # appropriate
    if type_ == "eval" and plot_avg:
        all_returns = all_returns.mean(axis=-1)
    elif type_ == "eval" and not plot_avg:
        print(all_returns.shape)
        all_returns = np.concatenate(all_returns, axis=1)
        all_returns = all_returns.transpose()
    elif type_ == "train":
        all_returns = np.squeeze(all_returns)

    # Smooth returns
    all_returns = exp.smooth(all_returns, smooth_over, keep_shape)

    # Plot the average reward
    if env_type == "continuing":
        episode_steps = data["experiment"]["environment"]
        episode_steps = episode_steps["steps_per_episode"]
        all_returns /= episode_steps

    # Detrend the data if appropriate
    if detrend == "linear":
        all_returns = exp.detrend_linear(all_returns, 1)
    elif detrend == "mean":
        all_returns = exp.detrend_linear(all_returns, 1, "constant")
    elif detrend == "difference":
        all_returns = exp.detrend_difference(all_returns, 1)

    # Determine whether to plot episodes or timesteps on the x-axis, which is
    # dependent on the environment type
    if env_type == "episodic":
        xvalues = np.arange(all_returns.shape[1])  # episodes
    else:
        xvalues = timesteps[:all_returns.shape[1]]

    # Plot each run
    for run in range(all_returns.shape[0]):
        ax.plot(xvalues, all_returns[run], color=colour, linestyle="--",
                alpha=alpha)

    # Plot the mean
    mean = all_returns.mean(axis=0)
    ax.plot(xvalues, mean, color=colour)

    # Store legend identifiers for the run
    legend_labels = []
    legend_lines = []
    legend_labels.append("Individual Runs")
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="--",
                               alpha=alpha))

    # Set up the legend variables for the mean over all runs
    label = name
    legend_labels.append(label)
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="-"))

    return fig, ax, legend_labels, legend_lines


def mean_with_stderr(data, type_, ind, smooth_over, names,
                     fig=None, ax=None, figsize=(12, 6),
                     xlim=None, ylim=None, alpha=0.2,
                     colours=None, env_type="continuing",
                     keep_shape=False, xlabel="", ylabel=""):
    """
    Plots the average training or evaluation return over all runs with standard
    error.

    Given a list of data dictionaries of the form returned by main.py, this
    function will plot each episodic return for the list of hyperparameter
    settings ind each data dictionary. The ind argument is a list, where each
    element is a list of hyperparameter settings to plot for the data
    dictionary at the same index as this list. For example, if ind[i] = [1, 2],
    then plots will be generated for the data dictionary at location i
    in the data argument for hyperparameter settings ind[i] = [1, 2].
    The smooth_over argument tells how many previous data points to smooth
    over

    Parameters
    ----------
    data : list of dict
        The Python data dictionaries generated from running main.py for the
        agents
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of iter of int
        The list of lists of hyperparameter settings indices to plot for
        each agent. For example [[1, 2], [3, 4]] means that the first agent
        plots will use hyperparameter settings indices 1 and 2, while the
        second will use 3 and 4.
    smooth_over : list of int
        The number of previous data points to smooth over for the agent's
        plot for each data dictionary. Note that this is *not* the number of
        timesteps to smooth over, but rather the number of data points to
        smooth over. For example, if you save the return every 1,000
        timesteps, then setting this value to 15 will smooth over the last
        15 readings, or 15,000 timesteps. For example, [1, 2] will mean that
        the plots using the first data dictionary will smooth over the past 1
        data points, while the second will smooth over the passed 2 data
        points for each hyperparameter setting.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    names : list of str
        The name of the agents, used for the legend
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of list of str
        The colours to use for each hyperparameter settings plot for each data
        dictionary
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot
    """
    # Set the colours to be default if not set
    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    title = f"Average {type_.title()} Return per Run with Standard Error"
    fig, ax = _setup_fig(fig, ax, figsize, xlim=xlim, ylim=ylim, xlabel=xlabel,
                         ylabel=ylabel, title=title)
    # if ax is None and fig is None:
    #     fig = plt.figure(figsize=figsize)
    #     ax = fig.add_subplot()

    # if xlim is not None:
    #     ax.set_xlim(xlim)
    # if ylim is not None:
    #     ax.set_ylim(ylim)

    # Track the total timesteps per hyperparam setting over all episodes and
    # the cumulative timesteps per episode per data dictionary (timesteps
    # should be consistent between all hp settings in a single data dict)
    total_timesteps = []
    cumulative_timesteps = []

    for i in range(len(data)):
        if type_ == "train":
            cumulative_timesteps.append(exp.get_cumulative_timesteps(data[i]
                                        ["experiment_data"][ind[i][0]]["runs"]
                                        [0]["train_episode_steps"]))
        elif type_ == "eval":
            cumulative_timesteps.append(data[i]["experiment_data"][ind[i][0]]
                                        ["runs"][0]["timesteps_at_eval"])
        else:
            raise ValueError("type_ must be one of 'train', 'eval'")
        total_timesteps.append(cumulative_timesteps[-1][-1])

    # Find the minimum of total trained-for timesteps. Each plot will only
    # be plotted on the x-axis until this value
    min_timesteps = min(total_timesteps)

    # For each data dictionary, find the minimum index where the timestep at
    # that index is >=  minimum timestep
    ind_ge_min_timesteps = []
    for cumulative_timesteps_per_data in cumulative_timesteps:
        final_ind = np.where(cumulative_timesteps_per_data >=
                             min_timesteps)[0][0]
        # Since indexing will stop right before the minimum, increment it
        ind_ge_min_timesteps.append(final_ind + 1)

    # Plot all data for all HP settings, only up until the minimum index
    # fig, ax = None, None
    plot_fn = _plot_mean_with_stderr_continuing if env_type == "continuing" \
        else _plot_mean_with_stderr_episodic

    for i in range(len(data)):
        fig, ax = \
            plot_fn(data=data[i], type_=type_,
                    ind=ind[i], smooth_over=smooth_over[i], name=names[i],
                    fig=fig, ax=ax, figsize=figsize, xlim=xlim, ylim=ylim,
                    last_ind=ind_ge_min_timesteps[i], alpha=alpha,
                    colours=colours[i], keep_shape=keep_shape)

    return fig, ax


def _plot_mean_with_stderr_continuing(data, type_, ind, smooth_over, fig=None,
                                      ax=None, figsize=(12, 6), xlim=None,
                                      ylim=None, xlabel=None, ylabel=None,
                                      name="", last_ind=-1,
                                      timestep_multiply=None, alpha=0.2,
                                      colours=None,
                                      keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a single
    data dictionary on a continuing environment. Standard error
    is plotted as shaded regions.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    timestep_multiply : array_like of float, optional
        A value to multiply each timstep by, by default None. This is useful if
        your agent does multiple updates per timestep and you want to plot
        performance vs. number of updates.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    if colours is not None and len(colours) != len(ind):
        raise ValueError("must have one colour for each hyperparameter " +
                         "setting")

    if timestep_multiply is None:
        timestep_multiply = [1] * len(ind)

    if ax is not None and fig is None:
        raise ValueError("must pass figure when passing axis")

    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    if ax is None and fig is None:
        title = f"Average {type_.title()} Return per Run with Standard Error"
        fig, ax = _setup_fig(fig, ax, figsize, xlim=xlim, ylim=ylim,
                             xlabel=xlabel, ylabel=ylabel, title=title)
    #     fig = plt.figure(figsize=figsize)
    #     ax = fig.add_subplot()

    # if xlim is not None:
    #     ax.set_xlim(xlim)
    # if ylim is not None:
    #     ax.set_ylim(ylim)

    episode_length = data["experiment"]["environment"]["steps_per_episode"]

    # Plot with the standard error
    for i in range(len(ind)):
        timesteps, mean, std = exp.get_mean_err(data, type_, ind[i],
                                                smooth_over, exp.stderr,
                                                keep_shape=keep_shape)
        timesteps = np.array(timesteps[:last_ind]) * timestep_multiply[i]
        mean = mean[:last_ind] / episode_length
        std = std[:last_ind] / episode_length

        # Plot based on colours
        label = name
        if colours is not None:
            _plot_shaded(ax, timesteps, mean, std, colours[i], label, alpha)
        else:
            _plot_shaded(ax, timesteps, mean, std, None, label, alpha)

    ax.legend()

    fig.show()
    return fig, ax


def _plot_mean_with_stderr_episodic(data, type_, ind, smooth_over, fig=None,
                                    ax=None, figsize=(12, 6), name="",
                                    last_ind=-1, xlabel="Episodes",
                                    ylabel="Average Return", xlim=None,
                                    ylim=None, alpha=0.2, colours=None,
                                    keep_shape=False):
    """
    Plots the average training or evaluation return over all runs for a
    single data dictionary on an episodic environment. Plots shaded retions
    as standard error.

    Parameters
    ----------
    data : dict
        The Python data dictionary generated from running main.py
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    ind : iter of int
        The list of hyperparameter settings indices to plot
    smooth_over : int
        The number of previous data points to smooth over. Note that this
        is *not* the number of timesteps to smooth over, but rather the number
        of data points to smooth over. For example, if you save the return
        every 1,000 timesteps, then setting this value to 15 will smooth
        over the last 15 readings, or 15,000 timesteps.
    fig : plt.figure
        The figure to plot on, by default None. If None, creates a new figure
    ax : plt.Axes
        The axis to plot on, by default None, If None, creates a new axis
    figsize : tuple(int, int)
        The size of the figure to plot
    name : str, optional
        The name of the agent, used for the legend
    last_ind : int, optional
        The index of the last element to plot in the returns list,
        by default -1. This is useful if you want to plot many things on the
        same axis, but all of which have a different number of elements. This
        way, we can plot the first last_ind elements of each returns for each
        agent.
    xlim : float, optional
        The x limit for the plot, by default None
    ylim : float, optional
        The y limit for the plot, by default None
    alpha : float, optional
        The alpha channel for the plot, by default 0.1
    colours : list of str
        The colours to use for each plot of each hyperparameter setting
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'. By default
        'continuing'

    Returns
    -------
    plt.figure, plt.Axes
        The figure and axes of the plot

    Raises
    ------
    ValueError
        When an axis is passed but no figure is passed
        When an appropriate number of colours is not specified to cover all
        hyperparameter settings
    """
    if colours is not None and len(colours) != len(ind):
        raise ValueError("must have one colour for each hyperparameter " +
                         "setting")

    if ax is not None and fig is None:
        raise ValueError("must pass figure when passing axis")

    if colours is None:
        colours = _get_default_colours(ind)

    # Set up figure
    if ax is None and fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    # Plot with the standard error
    for i in range(len(ind)):
        data = exp.reduce_episodes(data, ind[i], type_=type_)

        # data has consistent # of episodes, so treat as env_type="continuing"
        _, mean, std = exp.get_mean_err(data, type_, ind[i], smooth_over,
                                        exp.stderr, keep_shape=keep_shape)
        mean = mean[:last_ind]
        std = std[:last_ind]
        episodes = np.arange(mean.shape[0])

        # Plot based on colours
        label = name
        if colours is not None:
            _plot_shaded(ax, episodes, mean, std, colours[i], label, alpha)
        else:
            _plot_shaded(ax, episodes, mean, std, None, label, alpha)

    ax.legend()
    ax.set_title(f"Average {type_.title()} Return per Run with Standard Error")
    ax.set_ylabel(ylabel)
    ax.set_xlabel(xlabel)

    fig.show()
    return fig, ax


def hyperparam_sensitivity_plot(dir, type_, hp_name, combo, env_config,
                                agent_config, figsize=(12, 6), fig=None,
                                ax=None, annotate=False, colour=None,
                                env_type="continuing", conf=False,
                                significance=0.1):
    """
    Plots the hyperparameter sensitivity for hp_name, where the combo argument
    determines which hyperparameters are held constant.

    Given hyperparameters a, b, and c, let's say we want to get all
    hyperparameter settings indices where a varies, and b and c are constant.
    if a, b, and c can each be 1 or 2, then there are four ways that a can
    vary with b and c remaining constant:

        [
            ((a=1, b=1, c=1), (a=2, b=1, c=1)),         combo = 0
            ((a=1, b=2, c=1), (a=2, b=2, c=1)),         combo = 1
            ((a=1, b=1, c=2), (a=2, b=1, c=2)),         combo = 2
            ((a=1, b=2, c=2), (a=2, b=2, c=2))          combo = 3
        ]

    The combo argument indexes into this list of hyperparameter settings

    Parameters
    ----------
    dir : str
        The directory containing all data dictionaries generated from running
        main.py, separated into one file per HP index
    type_ : str
        Which type of data to plot, one of "eval" or "train"
    combo : int
        Determines the values of the constant hyperparameters. Given that
        only one hyperparameter may vary, there are many different sets
        having this hyperparameter varying with all others remaining constant
        since each constant hyperparameter may take on many values. This
        argument cycles through all sets of hyperparameter settings indices
        that have only one hyperparameter varying and all others constant.
    env_config : dict
        The environment configuration file as a Python dictionary
    agent_config : dict
        The agent configuration file as a Python dictionary
    figsize : tuple(int, int)
        The size of the figure
    fig : plt.figure, optional
        The figure to plot on, by default None. If None, a new figure is
        created
    ax : plt.Axes, optional
        The axis to plot on, by default None. If None, a new axis is created.
    annotate : bool, optional
        Whether or not to annoate each point drawn with its y-value
    colour : str
        The colour to use in the plot, by default None. If None, then default
        values will be used.

    Returns
    -------
    tuple of plt.figure, plt.Axes
        The figure and axis that was plotted on

    Raises
    ------
    ValueError
        If only one of fig, ax are specified
    """
    if (fig is None and ax is not None) or (ax is None and fig is not None):
        raise ValueError("fig and ax must both be specified if one is "
                         "specified")

    if not conf:
        stderr_fn = exp.get_mean_returns_with_stderr_hp_varying
        hp_values, mean_returns, error = \
            stderr_fn(dir, type_, hp_name, combo, env_config, agent_config,
                      after=0, env_type=env_type)
    else:
        conf_fn = exp.get_mean_returns_with_conf_hp_varying
        hp_values, mean_returns, error = \
            conf_fn(dir, type_, hp_name, combo, env_config, agent_config,
                    after=0, env_type=env_type, significance=significance)

    # Set up the figure
    if fig is None and ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    ax.set_title(f"{hp_name.title()} Hyperparameter Sensitivity Plot " +
                 f"over {type_.title()} Data")
    ax.set_xlabel("Hyperparameter Values")

    if env_type == "episodic":
        ax.set_ylabel("Average Return")
    elif env_type == "continuing":
        ax.set_ylabel("Average Reward")

    if colour is None:
        colour = DEFAULT_COLOURS[0]

    if env_type == "continuing":
        ep_steps = env_config["steps_per_episode"]
        mean_returns /= ep_steps
        error /= ep_steps

    _plot_shaded(ax, hp_values, mean_returns, error, colour, "",
                 alpha=0.25)

    if annotate:
        for i, coords in enumerate(zip(hp_values, mean_returns)):
            ax.annotate(coords[i], xy=coords)

    return fig, ax


def hyperparam_sensitivity_plot_by_settings_index(dir, type_, hp_name, ind,
                                                  env_config, agent_config,
                                                  figsize=(12, 6),
                                                  annotate=False,
                                                  env_type="continuing",
                                                  conf=False,
                                                  significance=0.1):
    """
    Plot hyperparameter sensitivity by hyperparameter settings index.

    Plots the hyperparameter sensitivity plot where all
    constant hyperparameters are equal to those defined by the
    argument hyperparameter settings index. The only hyperparameter that
    is allowed to vary is the one corresponding to hp_name.

    Parameters
    ----------
    dir : str
        The directory containing all data dictionaries generated by running
        main.py
    type_ : str
        Whether to plot the training or evaluation data, one of
        'train', 'eval'
    hp_name : str
        The name of the hyperparameter to plot the sensitivity of
    config : dict
        A Python dictionary of the agent configuration file used to run
        the experiment
    ind : int
        The hyperparameter settings index that should define the
        values of all constant hyperparameters. Only the hyperparameter
        defined by hp_name is allowed to vary, all others must be equal to
        the values of those defined by the settings index ind.
    env_config : dict
        The environment configuration file as a Python dictionary
    agent_config : dict
        The agent configuration file as a Python dictionary
    figsize : tuple(int, int)
        The size of the figure
    annotate : bool, optional
        Whether or not to annoate each point drawn with its y-value

    Returns
    -------
    tuple of plt.figure, plt.Axes
        The figure and axis that was plotted on
    """
    # Get all the combinations of hyperparameters with only hp_name varying
    hp_name_varying_only_settings = \
        exp.get_varying_single_hyperparam(env_config, agent_config, hp_name)

    # Get the index (into the list of tuples with only hp_name varying) of
    # the tuple which has hp_name varying only, but all other constant
    # hyperparameters equal to those defined by hyperparam setting number ind.
    # So hp_combo_constant_hps_equal_ind is the hyperparam combination with
    # all constant hyperparams equal to those in settings number ind.
    hp_combo_constant_hps_equal_ind = \
        next(filter(lambda x: ind in x, hp_name_varying_only_settings))
    combo_index = hp_name_varying_only_settings.index(
        hp_combo_constant_hps_equal_ind)

    print(hp_name_varying_only_settings[combo_index])

    # Plot the hyperparameter sensitivity plot using all constant
    # hyperparameters equal to those defined in the settings
    # number ind
    return hyperparam_sensitivity_plot(dir, type_, hp_name, combo_index,
                                       env_config, agent_config,
                                       figsize=figsize, annotate=annotate,
                                       env_type=env_type, conf=conf,
                                       significance=significance)


@deprecated(reason="""We no longer store the state seen at each timestep
            as this was becoming very memory intensive. We are storing the
            states seen at each timestep as tile coded counts, similar to the
            states used in updates""")
def plot_states_visited(data, figsize=(15, 10), max_=0.1, xlabel="", ylabel="",
                        title="", hp_ind=-1, rasterize=True,
                        env_type="continuing"):
    """
    Plots the state visitation distribution during training

    Given a data dictionary, plots the state visitation distribution of a
    single hyperparameter setting.

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : float, optional
        The cutoff threshold for updates, by default 1.0. Any states that have
        updates larger than this will be capped at this value. Note that
        update counts are normalized between [0, 1].
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    title : str, optional
        The title of the plot, by default ""
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.
    env_type : str, optional
        The type of environment, one of 'continuing', 'episodic'; by default
        "continuing"

    Returns
    -------
    plt.figure, plt.Axes, array_like
        The figure and axis plotted on and the array of state counts
    """
    # Choose the hyperparameter setting if not st
    if hp_ind < 0:
        hp_ind = next(iter(data["experiment_data"]))

    # Set up the TileCoder
    min_state_values = data["experiment"]["environment"]["state_low"]
    max_state_values = data["experiment"]["environment"]["state_high"]
    state_bins = data["experiment"]["environment"]["state_bins"]
    state_coder = TileCoder(low=min_state_values, high=max_state_values,
                            bins=state_bins, num_tilings=1, include_bias=False,
                            offset=False, normalize=False)

    states_visited = np.zeros(np.prod(state_bins))

    # If the type is continuing, then all data fits nicely into a single matrix
    # and we can easily calculate everything using array operations
    if env_type == "continuing":
        size = data["experiment_data"][hp_ind]["runs"][0]["train_states"].shape
        for run in data["experiment_data"][hp_ind]["runs"]:
            states = run["train_states"]
            states = states.reshape(size[0] * size[1], size[2])
            states_visited += state_coder(states).mean(axis=0)

    # If episodic, then since runs/episodes have different lengths, we instead
    # have an array, where each row /col is a different length (list of states)
    # since each run/episode has a different length. So, we have to deal with
    # this case differently, looping over all runs, then all states per
    # episode per run
    elif env_type == "episodic":
        for run in data["experiment_data"][hp_ind]["runs"]:
            for episode_states in run["train_states"]:
                # For each episode, get the states and tile code, we go through
                # episodes separately since each has its own length
                states = np.stack(episode_states)
                states_visited += state_coder(states).mean(axis=0)

    fig, ax = _state_heatmap(states_visited, data["experiment"]["environment"],
                             figsize, max_, xlabel, ylabel, rasterize, title),
    return fig, ax, states_visited


def state_heatmap(data, figsize=(15, 10), max_=1.0, xlabel="", ylabel="",
                  updates=True, use_next_state=False, hp_ind=-1,
                  rasterize=True, time_ind=-1, use_interval_only=False,
                  title="", cbar=True, fig=None, ax=None):
    """
    Plot a heatmap of the states used in each update

    This function will only plot a 2D heatmap of states. The first state
    coordinate will refer to the rows of the heatmap, while the second will
    refer to the columns. The minimum value of the first state coordinate
    will be the top of the heatmap, while the max will be the bottom.
    Similarly, the maximum of the second coordinate will be the right of
    the plot, while the minimum will be the left.

    This function assumes a vector that has as each element, the number of
    updates performed to that corresponding state. If the state space is
    continuous, simply use a high resolution tile coder and discretize the
    state space.

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : float, optional
        The cutoff threshold for updates, by default 1.0. Any states that have
        updates larger than this will be capped at this value. Note that
        update counts are normalized between [0, 1].
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    updates : bool, optional
        Whether to plot the states used in updates or the states visited, by
        default True.
    use_next_state : bool, optional
        Whether to plot the state or next state in the SARSA update,
        by default False
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.
    time_ind : int, optional
        In each experiment, the state update heatmap is stored at multiple
        timesteps at certain intervals. This parameter controls which of the
        update heatmaps to plot, with 0 being the first and -1 being the last
        heatmap saved, by default -1.
    use_interval_only : bool, optional
        Whether to plot only the states visited between the time_ind reading
        and the previous reading. That is, whether or not to plot only the
        states used in updates or visited between the argument time_ind and
        the previous, but none before. By default False.
    title : str, optional
        The title of the figure, by default ""

    Returns
    -------
    plt.figure, plt.Axes, array_like
        The figure and axis plotted on and the array of state counts

    Raises
    ------
    ValueError
        If attempting to plot the next states, but not plotting states used in
        updates
    """
    if fig is None and ax is not None:
        raise ValueError("fig and ax must both be specified")

    if use_next_state and not updates:
        raise ValueError("""can only plot next states when plotting states
                         used in updates, set 'updates=True'""")

    # Choose the hyperparameter setting if not st
    if hp_ind < 0:
        hp_ind = next(iter(data["experiment_data"]))

    # Get the number of updates per each state per run
    if updates:
        state_type = "update_states" if not use_next_state else \
            "update_next_states"
    else:
        state_type = "train_states"

    # Get the mean states used in each update or visited at each timestep
    # over all runs
    states = _get_mean_state_bin_count(data, state_type, hp_ind, time_ind)

    # Calculate the states seen only for the current interval
    if use_interval_only:
        if time_ind == 0:
            prev_states = np.zeros(states.shape)
        else:
            prev_states = _get_mean_state_bin_count(data, state_type, hp_ind,
                                                    time_ind-1)
        states -= prev_states

    # Plot
    fig, ax = _state_heatmap(states, data["experiment"]["environment"],
                             figsize, max_, xlabel, ylabel, title, rasterize,
                             cbar, fig=fig, ax=ax)
    return fig, ax, states


def _get_mean_state_bin_count(data, state_type, hp_ind, time_ind):
    # Get the states used in each update or visited at each timestep
    states = []
    for run in data["experiment_data"][hp_ind]["runs"]:
        states.append(run[state_type][time_ind])

    # Get the average number of updates per state over all runs
    return np.stack(states).mean(axis=0)


def _state_heatmap(visited_states, env_config, figsize=(15, 10), max_=0.1,
                   xlabel="", ylabel="", title="", rasterize=True, cbar=True,
                   fig=None, ax=None):
    """
    Plots a heatmap of states, where states are separated into bins

    This function is designed to work with the TileCoder object. It takes in
    a flattened array, where each index in the array refers to a sequential
    bin of states. The element at each position is the number of times that
    state was updated or encountered.

    Parameters
    ----------
    visited_states : array_like
        The flattened array of counts, where each index in the array
        corresponds to the index of the bin
    env_config : dict
        The environment configuration file for the environment that the
        experiment was run on
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : float, optional
        The cutoff threshold for updates, by default 1.0. Any states that have
        updates larger than this will be capped at this value. Note that
        update counts are normalized between [0, 1].
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    title : str, optional
        The title of the figure, by default ""
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.

    Returns
    -------
    plt.figure, plt.Axes, array_like
        The figure and axis plotted on and the array of state counts
    """
    if fig is None and ax is not None:
        raise ValueError("fig and ax must both be specified")

    if fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()
    ax.set_title(title)
    ax.set(xlabel=xlabel, ylabel=ylabel)

    max_visited = visited_states.max()
    min_visited = visited_states.min()
    if max_visited != min_visited:
        visited_states = (visited_states - min_visited) / \
            (max_visited - min_visited)

    # Threshold the updates
    ind = np.where(visited_states > max_)[0]
    visited_states[ind] = max_

    # Determine the bins per state dimension
    if len(env_config["state_bins"]) == 1:
        state_bins = (1, env_config["state_bins"][0])
    else:
        state_bins = env_config["state_bins"]
    visited_states = visited_states.reshape(state_bins)

    # Plot the heatmap
    ax = sns.heatmap(visited_states, ax=ax, linewidths=0, rasterized=rasterize,
                     cbar=cbar, vmin=0, vmax=1.0)

    # Set the axis tick locations
    num_ticks = 3
    ax.xaxis.set_major_locator(ticker.LinearLocator(num_ticks))
    ax.yaxis.set_major_locator(ticker.LinearLocator(num_ticks))

    # Label axis ticks
    xtick, ytick = get_ticks(num_ticks, num_ticks, env_config, type_="state")

    ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    fig.show()
    return fig, ax


def action_heatmap_with_barchart(data, figsize=(20, 7), max_=-1,
                                 xlabel="", ylabel="", hp_ind=-1,
                                 rasterize=True, ratio=(2, 1), time_ind=-1):
    """
    Plot the action heatmap with a barchart above it

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    max_ : int, optional
        The cutoff threshold for updates, by default 10. Any states that have
        updates larger than this will be capped at this value.
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    rasterize : bool, optional
        Whether or not to rasterize the graphics, by default True. Useful
        to remove gridlines between cells when saving as a PDF or SVG; in
        this case, set to False.
    ratio : tuple, optional
        The ratio of bar char size to heatmap size, by default (2, 1)
    time_ind : int, optional
        In each experiment, the state update heatmap is stored at multiple
        timesteps at certain intervals. This parameter controls which of the
        update heatmaps to plot, with 0 being the first and -1 being the last
        heatmap saved, by default -1.

    Returns
    -------
    plt.figure, plt.Axes, plt.Axes
        The figure, barchart axis, and heatmap axis plotte on
    """
    # Set up the plot
    fig = plt.figure(figsize=figsize)
    gs = gridspec.GridSpec(*ratio, height_ratios=[4, 1])
    bar_ax = fig.add_subplot(gs[0])
    heat_ax = fig.add_subplot(gs[1])

    # Get the number updates per action bin
    if hp_ind < 0:
        hp_ind = next(iter(data["experiment_data"]))
    update_actions = []
    for run in data["experiment_data"][hp_ind]["runs"]:
        update_actions.append(run["update_actions"][time_ind])

    # Get the average number of upates per state over all runs
    update_actions = np.stack(update_actions).mean(axis=0)

    # Plot the barchart
    bins = np.arange(update_actions.shape[0])
    bar_ax.bar(bins, update_actions)

    # Set the axis limits for the barchart
    bar_ax.set_xlim((bins[0], bins[-1]))
    ymax = update_actions.max()
    bar_ax.set_ylim((0, ymax))

    # Threshold the updates
    if max_ > -1:
        ind = np.where(update_actions > max_)[0]
        update_actions[ind] = max_

    # Determine the bins per action dimension
    env_config = data["experiment"]["environment"]
    if len(env_config["action_bins"]) == 1:
        action_bins = (1, env_config["action_bins"][0])
    else:
        action_bins = env_config["action_bins"]
    update_actions = update_actions.reshape(action_bins)

    # Plot the heatmap
    heat_ax = sns.heatmap(update_actions, ax=heat_ax, linewidths=0,
                          rasterized=rasterize,
                          cbar_kws={"orientation": "horizontal"})
    heat_ax.set(xlabel=xlabel, ylabel=ylabel)

    # Set the heatmap axis tick locations
    num_ticks = 3
    heat_ax.xaxis.set_major_locator(ticker.LinearLocator(num_ticks))
    heat_ax.yaxis.set_major_locator(ticker.LinearLocator(0))

    # Label heatmap axis ticks
    xtick, ytick = get_ticks(num_ticks, 0, env_config, type_="action")

    # Place the heatmaps ticks
    heat_ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    heat_ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    fig.show()
    return fig, bar_ax, heat_ax


def get_ticks(x_num_ticks, y_num_ticks, env_config, type_="state"):
    """
    Factor to get the functions to label the x and y ticks

    Parameters
    ----------
    num_ticks : int
        The total number of ticks along the axis
    env_config : dict
        The environment configuration file as a dict
    type_ : str, optional
        The type of plot to get the ticks of the heatmap for, one of 'state',
        'action', by default 'state'

    Returns
    -------
    2-tuple of func
        The function to generate the x tick labels and the function to
        generate the y tick labels

    Raises
    -------
    NotImplementedError
        If trying to get tick labels for an unsupported environment
    """
    if env_config["env_name"] == "Pendulum-v1":
        return _2D_ticks(x_num_ticks, y_num_ticks, env_config, type_)
    elif env_config["env_name"] == "MountainCarContinuous-v0":
        return _2D_ticks(x_num_ticks, y_num_ticks, env_config, type_)
    else:
        raise NotImplementedError()


def _2D_ticks(x_num_ticks, y_num_ticks, env_config, type_):
    """
    Generates the tick labels for the heatmap on a 2D-state environment

    Parameters
    ----------
    num_ticks : int
        The total number of ticks along the axis
    env_config : dict
        The environment configuration file as a dict
    type_ : str
        The type of plot to get the ticks of the heatmap for, one of 'state',
        'action'

    Returns
    -------
    2-tuple of func
        The function to generate the x tick labels and the function to
        generate the y tick labels
    """
    if type_ == "state":
        ymin, xmin = env_config["state_low"]
        ymax, xmax = env_config["state_high"]
    else:
        ymin, xmin = 0.0, env_config["action_low"][0]
        ymax, xmax = 0.0, env_config["action_high"][0]

    def xtick(x, pos):
        if pos == 0:
            return "{:.2f}".format(xmin)
        elif pos == x_num_ticks - 1:
            return "{:.2f}".format(xmax)
        elif pos == x_num_ticks - (x_num_ticks // 2):
            return "{:.1f}".format((xmax + xmin) / 2)
        else:
            return ""

    def ytick(x, pos):
        if pos == 0:
            return "{:.2f}".format(ymin)
        elif pos == y_num_ticks - 1:
            return "{:.2f}".format(ymax)
        elif pos == y_num_ticks - (y_num_ticks // 2):
            return "{:.1f}".format((ymax + ymin) / 2)
        else:
            return ""

    return xtick, ytick


def state_update_surface_plot(data, figsize=(20, 15),
                              type_="surface", max_=100, smooth=1.5,
                              view=(30, -45), linewidth=0, colour=None,
                              time_ind=-1):
    """
    Plot the surface of states used in each update

    Note that unlike the heatmap, the updates plotted for the surface are not
    normalized betwee 0 and 1.

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    type_ : str, optional
        The type of surface to plot, by default "surface". One of 'surface',
        'wireframe', or 'bar'
    max_ : int, optional
        The cutoff threshold for updates, by default 100. Any states that have
        updates larger than this will be capped at this value.
    smooth : float, optional
        The amount to smoothen the plot, by default 1.5. This value corresponds
        to the standard deviation of a Gaussian convolution kernel applied to
        the surface.
    view : 2-tuple of int, optional
        The view in 3D space, by default (30, -45)
    linewidth : int, optional
        The width of the lines separating cells, by default 0
    colour : str
        The colour to use for the plot, by default None. If None, uses the
        default colours.
    time_ind : int, optional
        In each experiment, the state update heatmap is stored at multiple
        timesteps at certain intervals. This parameter controls which of the
        update heatmaps to plot, with 0 being the first and -1 being the last
        heatmap saved, by default -1.

    Returns
    -------
    plt.figure, plt.Axes3D
        The figure and axis plotten on
    """
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(projection="3d")

    if colour is None:
        colour = DEFAULT_COLOURS[0]

    hp = next(iter(data["experiment_data"]))
    update_states = []
    for run in data["experiment_data"][hp]["runs"]:
        update_states.append(run["update_states"][time_ind])

    update_states = np.stack(update_states).mean(axis=0)
    ind = np.where(update_states > max_)[0]
    update_states[ind] = max_
    update_states = update_states.reshape((100, 100))

    update_states = ndimage.filters.gaussian_filter(update_states,
                                                    sigma=smooth)

    x = np.linspace(-np.pi, np.pi, 100)
    y = np.linspace(-8, 8, 100)
    x, y = np.meshgrid(x, y)

    ax.set_zlim3d(0, max_)
    ax.view_init(view[0], view[1])
    ax.grid(None)

    if type_ == "surface":
        ax.plot_surface(x, y, update_states, antialiased=False,
                        linewidth=linewidth, color=colour)
    elif type_ == "wireframe":
        ax.plot_wireframe(x, y, update_states, antialiased=False,
                          linewidth=linewidth, color=colour)
    elif type_ == "bar":
        # Get the min/max x/y values
        env_config = data["experiment"]["environment"]
        xmin, ymin = env_config["state_low"]
        xmax, ymax = env_config["state_high"]

        # Determine the width, depth, and height of barx
        dx = [(xmax - xmin) / 100] * len(x.ravel())
        dy = [(ymax - ymin) / 100] * len(y.ravel())
        dz = update_states.ravel()

        z = [0] * len(y.ravel())
        ax.bar3d(x.ravel(), y.ravel(), z, dx=dx, dy=dy, dz=dz, color=colour)

    fig.show()
    return fig, ax


def return_distribution(data, type_, hp_ind, bins, figsize=(12, 6), xlim=None,
                        ylim=None, after=0, before=-1):
    """
    Plots the distribution of returns on either an episodic or continuing
    environment

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    type_ : str, optional
        The type of surface to plot, by default "surface". One of 'surface',
        'wireframe', or 'bar'
    hp_ind : int, optional
        The hyperparameter settings index in the data dictionary to use for
        the plot, by default -1. If less than 0, then the first hyperparameter
        setting in the dictionary is used.
    bins : Iterable, int
        The bins to use for the plot. If an Iterable, then each value in the
        Iterable is considered as a cutoff for bins. If an integer, separates
        the returns into that many bins
    figsize : tuple, optional
        The size of the figure to plot, by default (15, 10)
    xlim : 2-tuple of float, optional
        The cutoff points for the x-axis to plot between, by default None
    ylim : 2-tuple of float, optional
        The cutoff points for the y-axis to plot between, by default None

    Returns
    -------
    plt.figure, plt.Axes3D
        The figure and axis plotten on
    """
    # Get the episode returns for each run
    run_returns = []
    return_type = type_ + "_episode_rewards"
    for run in data["experiment_data"][hp_ind]["runs"]:
        run_returns.append(np.mean(run[return_type][after:before]))

    title = f"Learning Curve Distribution - HP Settings {hp_ind}"
    return _return_histogram(run_returns, bins, figsize, title, xlim, ylim)


def _return_histogram(run_returns, bins, figsize, title, xlim, ylim, kde=True):
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot()

    ax.set_title(title)
    ax.set_xlabel("Average Return Per Run")
    ax.set_ylabel("Relative Frequency")
    _ = sns.histplot(run_returns, bins=bins, kde=kde)

    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    # Plot relative frequency on the y-axis
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos:
                                 "{:.2f}".format(x / len(run_returns))))

    fig.show()
    return fig, ax


def inds_vs_hp_heatmap(dir_, type_, hp1, hp2, inds, env_config, agent_config,
                       figsize=(12, 8), annotate=True, fmt=".2g",
                       rasterize=False, cmap=CMAP, linewidths=0,
                       env_type="continuing", default_value=-1000, after=0,
                       before=None, fig=None, ax=None, vmin=None, vmax=None):
    """
    Plots the mean returns of one hyperparameter vs another.

    Given a list of hyperparameter indices and two hyperparameter names,
    plots the mean returns generated by varying the second hyperparameter,
    and leaving the first hyperparameter fixed as each hyperparameter index's
    value. The plot is a heatmap. In this way, each row of the heatmap will
    be a certain HP setting with a fixed HP, while the columns will be the
    mean returns given the second HP varying for each HP index.

    Parameters
    ----------
    dir_ : str
        The directory which contains all data dictionaries
    type_ : str
        The type of returns to consider, one of 'train' or 'eval'
    hp1 : str
        The name of the hyperparameter setting to vary on the x-axis
    hp2 : str
        The name of the hyperparameter to plot on the y-axis for each
        hyperparameter settings index plotted
    inds : Iterable of int
        An iterable of hyperparameter indices used to plot
    env_config : dict
        The environment configuration file used to run the experiment
    agent_config : dict
        The agent configuration file used to run the experiment
    figsize : tuple, optional
        The size of the figure, by default (12, 8)
    annotate : bool, optional
        Whether each box in the heatmap should be annotated with its value,
        by default True
    fmt : str, optional
        If annotating, the format to use, by default ".2g"
    rasterize : bool, optional
        Whether or not the final heatmap should be rasterized, by default False
    cmap : str, optional
        The colourmap to use to plot, by default CMAP
    linewidths : int, optional
        The width of the lines used to separate boxes in the heatmap, by
        default 0
    env_type : str, optional
        Whether the environment is continuing or episodic, by default
        "continuing". One of 'episodic', 'continuing'
    default_value : int, optional
        The default value to use for heatmap boxes when no valid second
        hyperparameter can be used for the first hyperparameter, by default
        -1000. This happens, for example in hyperparameter settings that would
        have batch size > replay capacity.
    after : int, optional
        Only consider returns after this episode, by default 0
    before : int, optional
        Only consider returns before this episode, by default None
    fig : plt.Figure, optional
        The figure to plot on, by default None. If None, a new Figure is
        created.
    ax : plt.Axes, optional
        The axis to plot on, by default None. If None, a new Axes is created.
    vmin : float, optional
        Minimum value to anchor the colourmap to, otherwise it is inferred
        from the data.
    vmax : float, optional
        Maximum value to anchor the colourmap to, otherwise it is inferred
        from the data.

    Returns
    -------
    plt.Figure, plt.Axes
        The figure and axis plotted on

    Raises
    ------
    ValueError
        If the figure is not specified but the axis is
        If the type of return considered is not one of "train" or "eval"
        If an illegal environment type is given
    """
    if fig is None and ax is not None:
        raise ValueError("when ax is specified, fig should also be specified")

    if type_ not in ("train", "eval"):
        raise ValueError("type_ must be one of 'train' or 'eval'")

    if env_type not in ("episodic", "continuing"):
        raise ValueError("unknown environment type")

    values = []
    for ind in inds:
        # Get all the combinations of hyperparameters with only hp_name varying
        hp1_varying_only_settings = \
            exp.get_varying_single_hyperparam(env_config, agent_config, hp1)

        # Get the index (into the list of tuples with only hp_name varying) of
        # the tuple which has hp_name varying only, but all other constant
        # hyperparameters equal to those defined by hyperparam setting ind.
        # So hp_combo_constant_hps_equal_ind is the hyperparam combination with
        # all constant hyperparams equal to those in settings number ind.
        hp_combo_constant_hps_equal_ind = \
            next(filter(lambda x: ind in x, hp1_varying_only_settings))
        combo_index = hp1_varying_only_settings.index(
            hp_combo_constant_hps_equal_ind)

        # Get the mean returns for each hyperparameter setting
        mean_returns = []
        for elem in hp1_varying_only_settings[combo_index]:
            if elem is None:
                # No hyperparameter setting for this combo of hyperparams
                # e.g. batch size > buffer size
                # Use the default value for the return
                mean_returns.append(default_value)
                continue

            # Open the data file and compute the mean returns
            agent = agent_config["agent_name"]
            env = env_config["env_name"]
            dir_ = dir_.rstrip("/")
            file_ = dir_ + f"/{env}_{agent}_hp-{elem}.pkl"
            with open(file_, "rb") as in_file:
                data = pickle.load(in_file)

                # Calculate and save the mean return for this HP setting
                mean_return = exp.get_returns(data, type_, elem, env_type)
                if before is not None:
                    mean_return = mean_return[:, after:before].mean()
                else:
                    mean_return = mean_return[:, after:].mean()
                mean_returns.append(mean_return)

        values.append(mean_returns)
    all_mean_returns = np.stack(values)

    # Create the figure and axis
    if fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    # Plot the heatmap and set the labels
    ax = sns.heatmap(all_mean_returns, ax=ax, rasterized=rasterize,
                     annot=annotate, fmt=fmt, cmap=cmap, linewidths=linewidths,
                     vmin=vmin, vmax=vmax)
    ax.set(xlabel=hp1, ylabel=hp2)

    # Define the tick generation function
    def get_ticks(agent_config, hp1, hp2, inds):
        yvalues = []
        for ind in inds:
            params = exp.get_sweep_parameters(agent_config["parameters"],
                                              {"gamma": 0}, ind)[0]
            yvalues.append(params[hp2])

        max_x_ticks = len(agent_config["parameters"][hp1])

        def xticks(x, pos):
            x_tick = agent_config["parameters"][hp1][pos % max_x_ticks]
            return "{:.3g}".format(x_tick)

        def yticks(x, pos):
            y_tick = yvalues[pos]
            return "{:.3g}".format(y_tick)

        return xticks, yticks

    # Place the heatmaps ticks
    xtick, ytick = get_ticks(agent_config, hp1, hp2, inds)
    ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    return fig, ax


def inds_heatmap(dir_, type_, hp1, hp2, inds, env_config, agent_config,
                 figsize=(12, 8), annotate=True, fmt=".2g", rasterize=False,
                 cmap=CMAP, linewidths=0, env_type="continuing",
                 default_value=-1000, after=0, before=None, fig=None, ax=None,
                 vmin=None, vmax=None, div=1):
    """
    Plots the mean returns of each hyperparameter in a 2D matrix.

    Given a 2D array of hyperparameter settings indices, plots the mean return
    of each hyperparameter settings index as a heatmap. This function assumes
    that each row in the matrix corresponds to a different value of one
    hyperparameter, and each column corresponds to the values of a different
    hyperparameter, and these hyperparameters are plotted as the x labels and
    y labels.

    Note that any hyperparameter setting less than 0 is ignored, so these
    indices can be used to plot the default value in the heatmap.

    Parameters
    ----------
    dir_ : str
        The directory which contains all data dictionaries
    type_ : str
        The type of returns to consider, one of 'train' or 'eval'
    hp1 : str
        The name of the hyperparameter setting to vary on the x-axis
    hp2 : str
        The name of the hyperparameter to plot on the y-axis for each
        hyperparameter settings index plotted
    inds : 2D array_like of int
        Matrix of hyperparameter indices to plot
    env_config : dict
        The environment configuration file used to run the experiment
    agent_config : dict
        The agent configuration file used to run the experiment
    figsize : tuple, optional
        The size of the figure, by default (12, 8)
    annotate : bool, optional
        Whether each box in the heatmap should be annotated with its value,
        by default True
    fmt : str, optional
        If annotating, the format to use, by default ".2g"
    rasterize : bool, optional
        Whether or not the final heatmap should be rasterized, by default False
    cmap : str, optional
        The colourmap to use to plot, by default CMAP
    linewidths : int, optional
        The width of the lines used to separate boxes in the heatmap, by
        default 0
    env_type : str, optional
        Whether the environment is continuing or episodic, by default
        "continuing". One of 'episodic', 'continuing'
    default_value : int, optional
        The default value to use for heatmap boxes when no valid second
        hyperparameter can be used for the first hyperparameter, by default
        -1000. This happens, for example in hyperparameter settings that would
        have batch size > replay capacity.
    after : int, optional
        Only consider returns after this episode, by default 0
    before : int, optional
        Only consider returns before this episode, by default None
    fig : plt.Figure, optional
        The figure to plot on, by default None. If None, a new Figure is
        created.
    ax : plt.Axes, optional
        The axis to plot on, by default None. If None, a new Axes is created.
    vmin : float, optional
        Minimum value to anchor the colourmap to, otherwise it is inferred
        from the data.
    vmax : float, optional
        Maximum value to anchor the colourmap to, otherwise it is inferred
        from the data.

    Returns
    -------
    plt.Figure, plt.Axes
        The figure and axis plotted on

    Raises
    ------
    ValueError
        If the figure is not specified but the axis is
        If the type of return considered is not one of "train" or "eval"
        If an illegal environment type is given
    """
    # Ensure the hyperparameter settings have consistent hp values in each
    # row and column
    _check_index_matrix(agent_config, env_config, inds, hp2)
    _check_index_matrix(agent_config, env_config, inds.transpose(), hp1)

    if fig is None and ax is not None:
        raise ValueError("when ax is specified, fig should also be specified")

    if type_ not in ("train", "eval"):
        raise ValueError("type_ must be one of 'train' or 'eval'")

    if env_type not in ("episodic", "continuing"):
        raise ValueError("unknown environment type")

    values = np.zeros(inds.shape)
    for row in range(inds.shape[0]):
        for ind in range(inds[row].shape[0]):
            if inds[row, ind] < 0:
                values[row, ind] = default_value
                continue

            # Open the data file and compute the mean returns
            agent = agent_config["agent_name"]
            env = env_config["env_name"]
            dir_ = dir_.rstrip("/")
            index = inds[row, ind]
            file_ = dir_ + f"/{env}_{agent}_hp-{index}.pkl"
            with open(file_, "rb") as in_file:
                data = pickle.load(in_file)

                # Calculate and save the mean return for this HP setting
                mean_return = exp.get_avg_returns(data, type_, inds[row, ind],
                                                  after, before)

            values[row, ind] = mean_return.mean() / div

    # Create the figure and axis
    if fig is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot()

    # Plot the heatmap and set the labels
    ax = sns.heatmap(values, ax=ax, rasterized=rasterize,
                     annot=annotate, fmt=fmt, cmap=cmap, linewidths=linewidths,
                     vmin=vmin, vmax=vmax)
    ax.set(xlabel=hp1, ylabel=hp2)

    # Define the tick generation function
    def get_ticks(agent_config, hp1, hp2, inds):
        yvalues = []
        inds = inds.ravel()
        for ind in range(inds.size):
            params = exp.get_sweep_parameters(agent_config["parameters"],
                                              {"gamma": 0}, inds[ind])[0]
            if not params[hp2] in yvalues:
                yvalues.append(params[hp2])

        max_x_ticks = len(agent_config["parameters"][hp1])

        def xticks(x, pos):
            x_tick = agent_config["parameters"][hp1][pos % max_x_ticks]
            return "{:.3g}".format(x_tick)

        def yticks(x, pos):
            y_tick = yvalues[pos]
            return "{:.3g}".format(y_tick)

        return xticks, yticks

    # Place the heatmaps ticks
    xtick, ytick = get_ticks(agent_config, hp1, hp2, inds)
    ax.xaxis.set_major_formatter(ticker.FuncFormatter(xtick))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(ytick))

    return fig, ax


def _check_index_matrix(agent_config, env_config, inds, hp):
    """
    Check the columns of a matrix of hyperparameters indices for consistent
    hyperparameter values.

    Checks each row of a matrix of hyperparameter settings to ensure that each
    hyperparameter setting in the row has the same value for a certain
    hyperparameter. This is needed to ensure consistency when plotting a
    matrix of hyperparameter settings as a heatmap.

    Note that any hyperparameter setting less than 0 is ignored, so these
    indices can be used to plot the default value in the heatmap.

    Parameters
    ----------
    agent_config : dict
        The agent configuration file used to run the experiment
    env_config : dict
        The environment configuration file used to run the experiment
    inds : 2D array_like of int
        Matrix of hyperparameter indices to check to ensure that each row
        has hyperparameter settings that have the same hyperparameter value
        for a specific hyperparameter
    hp : str
        The hyperparameter to ensure is consistent between hyperparameter
        settings on each row.

    Raises
    ------
    IndexError
        If the hyperparameter indices do not refer to settings with the same
        value for the argument hyperparameter.
    ValueError
        If the array is not 2D.
    """
    if len(inds.shape) != 2:
        raise ValueError("inds must be a matrix")

    parameters = agent_config["parameters"]
    for row in range(inds.shape[0]):
        # The value that all rows should have as HP value
        hp_reference = None
        for col in range(inds.shape[1]):
            index = inds[row, col]
            if index < 0:
                # No hyperparameter setting in this position
                continue

            # Get the value of the hyperparameter for the index in the current
            # position
            hp_value = exp.get_sweep_parameters(parameters, env_config, index)
            hp_value = hp_value[0][hp]

            if hp_reference is None:
                # This is the first HP value encountered, use it as a reference
                # value for all other HPs in the row
                hp_reference = hp_value
                continue

            if hp_value != hp_reference:
                # Illegal HP value
                raise IndexError("hyperparameter indices do not share the " +
                                 f"same {hp} value")


@deprecated(reason="no longer storing states this way")
def absolute_state_over_time(data, after=0, figsize=(12, 6), state_ind=0,
                             colour=None, alpha=0.2, title="", xlabel="",
                             ylabel="", xlim=None, ylim=None):
    """
    Plot the average absolute value of a state feature per episode

    Parameters
    ----------
    data : dict
        The data dictionary containing the runs of a single hyperparameter
        setting
    after : int, optional
        Consider the state features only after this timestep for each
        episode, by default 0
    figsize : tuple, optional
        The size of the figure, by default (12, 6)
    state_ind : int, optional
        The index of the state feature to plot, by default 0
    colour : str, optional
        The colour for the plot, by default None
    alpha : float, optional
        The alpha value for the individual runs, by default 0.1
    title : str, optional
        The title for the plot, by default ""
    xlabel : str, optional
        The label for the x axis, by default ""
    ylabel : str, optional
        The label for the y axis, by default ""
    xlim : 2-tuple of float, optional
        The limits for the x axis, by default None
    ylim : 2-tuple of float, optional
        The limits for the y axis, by default None

    Returns
    -------
    plt.figure, plt.Axes
        The figure and Axes plotted on
    """
    if colour is None:
        colour = DEFAULT_COLOURS[0]

    hp = next(iter(data["experiment_data"]))
    states = []
    for run in data["experiment_data"][hp]["runs"]:
        states.append(run["train_states"])

    states = np.stack(states)[:, :, after:]
    states = np.abs(states)

    # Calculate the average performance, averaging over runs and then
    # timesteps to get the average angle per episode
    avg = states.mean(axis=0)
    avg = avg.mean(axis=1)

    # Get the average angle per episode for each run
    states = states.mean(axis=2)

    # Get some prelimiary variables for the legend
    legend_labels = []
    legend_lines = []
    agent_name = data["experiment"]["agent"]["agent_name"]

    # Get the mean over all runs and timesteps for each episode
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot()

    # Set axis limits
    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    # Plot the average angle per episode for all runs
    ax.plot(avg[:, state_ind], color=colour, linestyle="-")
    label = f"mean {agent_name} - Hyperparameter Setting {hp}"
    legend_labels.append(label)
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="-"))

    # Plot average average absolute angle for each episode
    # (averaged over timesteps) for each separate run
    for i in range(states.shape[0]):
        ax.plot(states[i, :, state_ind], linestyle="--", color=colour,
                alpha=alpha)
    legend_labels.append("individual runs")
    legend_lines.append(Line2D([0], [0], color=colour, linestyle="--",
                               alpha=alpha))

    ax.legend(legend_lines, legend_labels)

    ax.set_title(title)
    ax.set_ylabel(ylabel)
    ax.set_xlabel(xlabel)

    return fig, ax


def _get_default_colours(iter_):
    """
    Recursively turns elements of an Iterable into strings representing
    colours.

    This function will turn each element of an Iterable into strings that
    represent colours, recursively. If the elements of an Iterable are
    also Iterable, then this function will recursively descend all the way
    through every Iterable until it finds an Iterable with non-Iterable
    elements. These elements will be replaced by strings that represent
    colours. In effect, this function keeps the data structure, but replaces
    non-Iterable elements by strings representing colours. Note that this
    funcion assumes that all elements of an Iterable are of the same type,
    and so it only checks if the first element of an Iterable object is
    Iterable or not to stop the recursion.

    Parameters
    ----------
    iter_ : collections.Iterable
        The top-level Iterable object to turn into an Iterable of strings of
        colours, recursively.

    Returns
    -------
    list of list of ... of strings
        A data structure that has the same architecture as the input Iterable
        but with all non-Iterable elements replaced by strings.
    """
    colours = []

    # Calculate the number of lists at the current level to go through
    paths = range(len(iter_))

    # Return a list of colours if the elements of the list are not lists
    if not isinstance(iter_[0], Iterable):
        global OFFSET
        col = [DEFAULT_COLOURS[(OFFSET + i) % len(DEFAULT_COLOURS)]
               for i in paths]
        OFFSET += len(paths)
        return col

    # For each list at the current level, get the colours corresponding to
    # this level
    for i in paths:
        colours.append(_get_default_colours(iter_[i]))

    return colours


def _plot_shaded(ax, x, y, region, colour, label, alpha):
    """
    Plots a curve with a shaded region.

    Parameters
    ----------
    ax : plt.Axes
        The axis to plot on
    x : Iterable
        The points on the x-axis
    y : Iterable
        The points on the y-axis
    region : list or array_like
        The region to shade about the y points. The shaded region will be
        y +/- region. If region is a list or 1D np.ndarray, then the region
        is used both for the + and - portions. If region is a 2D np.ndarray,
        then the first row will be used as the lower bound (-) and the
        second row will be used for the upper bound (+). That is, the region
        between (lower_bound, upper_bound) will be shaded, and there will be
        no subtraction/adding of the y-values.
    colour : str
        The colour to plot with
    label : str
        The label to use for the plot
    alpha : float
        The alpha value for the shaded region
    """
    if colour is None:
        colour = DEFAULT_COLOURS[0]

    ax.plot(x, y, color=colour, label=label)
    if type(region) == list:
        ax.fill_between(x, y-region, y+region, alpha=alpha, color=colour)
    elif type(region) == np.ndarray and len(region.shape) == 1:
        ax.fill_between(x, y-region, y+region, alpha=alpha, color=colour)
    elif type(region) == np.ndarray and len(region.shape) == 2:
        ax.fill_between(x, region[0, :], region[1, :], alpha=alpha,
                        color=colour)


def _setup_fig(fig, ax, figsize=None, title=None, xlim=None, ylim=None,
               xlabel=None, ylabel=None, xscale=None, yscale=None, xbase=None,
               ybase=None):
    if fig is None:
        if ax is not None:
            raise ValueError("Must specify figure when axis given")
        if figsize is not None:
            fig = plt.figure(figsize=figsize)
        else:
            fig = plt.figure()

    if ax is None:
        ax = fig.add_subplot()

    if title is not None:
        ax.set_title(title)

    if xlabel is not None:
        ax.set_xlabel(xlabel)

    if ylabel is not None:
        ax.set_ylabel(ylabel)

    if xlim is not None:
        ax.set_xlim(xlim)

    if ylim is not None:
        ax.set_ylim(ylim)

    if xscale is not None:
        if xbase is not None:
            ax.set_xscale(xscale, base=xbase)
        else:
            ax.set_xscale(xscale)

    if yscale is not None:
        if ybase is not None:
            ax.set_yscale(yscale, base=ybase)
        else:
            ax.set_yscale(yscale)

    return fig, ax


def reset():
    """
    Resets the colours offset
    """
    global OFFSET
    OFFSET = 0


if __name__ == "__main__":
    print(_get_default_colours([[79], [58]]))

#!/usr/bin/env python3 # Import modules # TODO: Add documentation for plot_updates and plot_updates_ # TODO: Update action plots # TODO: state_update_heatmap should do 2D heatmap of multi-dimensional #       states, just give the two dimensions to plot in a heatmap. same for the #       surface plot!! # TODO: separate state_update_heatmap into two functions, one of which will #   simply take a vector of update counts and create a heatmap of it given the #   dimensions to reduce it to (e.g. reshape(100, 100)) # TODO: make _pendulum_ticks more general to label any number of ticks # TODO: Plot mean with runs episodic # TODO: Plot mean with stderr episodic for multiple data files # TODO: Change action plot to histogram # TODO: Make the state surface plot use either state updates or states visited #       Like the new scheme in state heatmap plot # TODO: State surface plot should take similar params to heatmap # TODO: Improve documentation # TODO: plot_mean_with_runs() should plot avg rewards for continuing # Global variables # sns.set_theme(palette=CMAP) # The offset to start in DEFAULT_COLOURS Plots the performance vs number of samples used in updates      Parameters     ----------     data : list of dict         The Python data dictionaries generated from running main.py for the         agents     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of iter of int         The list of lists of hyperparameter settings indices to plot for         each agent. For example [[1, 2], [3, 4]] means that the first agent         plots will use hyperparameter settings indices 1 and 2, while the         second will use 3 and 4.     smooth_over : list of int         The number of previous data points to smooth over for the agent's         plot for each data dictionary. Note that this is *not* the number of         timesteps to smooth over, but rather the number of data points to         smooth over. For example, if you save the return every 1,000         timesteps, then setting this value to 15 will smooth over the last         15 readings, or 15,000 timesteps. For example, [1, 2] will mean that         the plots using the first data dictionary will smooth over the past 1         data points, while the second will smooth over the passed 2 data         points for each hyperparameter setting.     names : list of str         The name of the agents, used for the legend     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     colours : list of list of str         The colours to use for each hyperparameter settings plot for each data         dictionary     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     alpha : float, optional         The alpha channel for the plot, by default 0.1     xscale : str, optional         The scale at which to plot the x axis, by default "log"     yscale : str, optional         The scale at which to plot the y axis, by default "linear"     base : int, optional         The base of the scale of the x axis, by default 10     xlabel : str, optional         The label for the x axis, by default ""     ylabel : str, optional         The label for the y axis, by default ""     title : str, optional         The title of the plot, by default ""     runs : bool, optional         Whether or not to plot the individual runs or error bars, by default         False. If True, plots individual runs instead of error bars.     keep_shape : bool, optional         Whether the shape of the input data should be kept after smoothing, by         default False. If True, then the first smooth_over episodes are         smoothed and kept. If False, then the first smooth_over episodes are         replaced by their fully smoothed value.     bootstrap : bool, optional         Whether or plot standard error or bootstrap confidence intervals, by         default False. If True, bootstrap confidence intervals are plotted         instead of standard error.     significance : float, optional         The significance level for the confidence interval, by default 0.01      Returns     -------     plt.Figure, plt.Axes         The figure and axes plotted on # Setup figure # Create legend label and lines # # Setup figure # fig, ax = _setup_fig(fig, ax, figsize, None, xlim, ylim, xlabel, ylabel, #                      xscale, yscale, base) # Calculate updates to plot along x-axis # Include the first eval -- before any updates have been done # Get the data to plot depending on the env type # If the original shape is not kept, then the updates need to be adjusted # for smoothing # Plot the mean # Standard error # Bootstrap confidence intervals # Plot each run # The episodic returns for each run # Plot the first eval episode, since this is the only data for # 0 updates # Get the steps per episode so we can plot average reward # Calculate standard error # Calculate bootstrap confidence interval # Determine the size of batches used in updates # Calculate the cumulative number of updates over all previous # episodes # For each cumulative number of updates plotted on the x-axis, # get the performance of the episode with number of updates closest # to this value # Plot the first eval episode, since this is the only data for # 0 updates # # Smooth data # Calculate standard error # Calculate bootstrap confidence interval Plots the average training or evaluation return over all runs with     confidence intervals.      Given a list of data dictionaries of the form returned by main.py, this     function will plot each episodic return for the list of hyperparameter     settings ind each data dictionary. The ind argument is a list, where each     element is a list of hyperparameter settings to plot for the data     dictionary at the same index as this list. For example, if ind[i] = [1, 2],     then plots will be generated for the data dictionary at location i     in the data argument for hyperparameter settings ind[i] = [1, 2].     The smooth_over argument tells how many previous data points to smooth     over      Parameters     ----------     data : list of dict         The Python data dictionaries generated from running main.py for the         agents     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of iter of int         The list of lists of hyperparameter settings indices to plot for         each agent. For example [[1, 2], [3, 4]] means that the first agent         plots will use hyperparameter settings indices 1 and 2, while the         second will use 3 and 4.     smooth_over : list of int         The number of previous data points to smooth over for the agent's         plot for each data dictionary. Note that this is *not* the number of         timesteps to smooth over, but rather the number of data points to         smooth over. For example, if you save the return every 1,000         timesteps, then setting this value to 15 will smooth over the last         15 readings, or 15,000 timesteps. For example, [1, 2] will mean that         the plots using the first data dictionary will smooth over the past 1         data points, while the second will smooth over the passed 2 data         points for each hyperparameter setting.     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     names : list of str         The name of the agents, used for the legend     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha channel for the plot, by default 0.1     colours : list of list of str         The colours to use for each hyperparameter settings plot for each data         dictionary     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     significance : float, optional         The significance level for the confidence interval, by default 0.01      Returns     -------     plt.figure, plt.Axes         The figure and axes of the plot # Set the colours to be default if not set # Track the total timesteps per hyperparam setting over all episodes and # the cumulative timesteps per episode per data dictionary (timesteps # should be consistent between all hp settings in a single data dict) # Find the minimum of total trained-for timesteps. Each plot will only # be plotted on the x-axis until this value # For each data dictionary, find the minimum index where the timestep at # that index is >=  minimum timestep # Since indexing will stop right before the minimum, increment it # Plot all data for all HP settings, only up until the minimum index # fig, ax = None, None Plots the average training or evaluation return over all runs for a single     data dictionary on a continuing environment. Bootstrap confidence intervals     are plotted as shaded regions.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of int         The list of hyperparameter settings indices to plot     smooth_over : int         The number of previous data points to smooth over. Note that this         is *not* the number of timesteps to smooth over, but rather the number         of data points to smooth over. For example, if you save the return         every 1,000 timesteps, then setting this value to 15 will smooth         over the last 15 readings, or 15,000 timesteps.     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     name : str, optional         The name of the agent, used for the legend     last_ind : int, optional         The index of the last element to plot in the returns list,         by default -1. This is useful if you want to plot many things on the         same axis, but all of which have a different number of elements. This         way, we can plot the first last_ind elements of each returns for each         agent.     timestep_multiply : int, optional         A value to multiply each timstep by, by default 1. This is useful if         your agent does multiple updates per timestep and you want to plot         performance vs. number of updates.     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha channel for the plot, by default 0.1     colours : list of str         The colours to use for each plot of each hyperparameter setting     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     significance : float, optional         The significance level for the confidence interval, by default 0.01      Returns     -------     plt.figure, plt.Axes         The figure and axes of the plot      Raises     ------     ValueError         When an axis is passed but no figure is passed         When an appropriate number of colours is not specified to cover all         hyperparameter settings # This should be the exact same as the episodic version except without # reducing the episodes. Follow the same structure as the episodic function # and the continuing function with standard error. Plots the average training or evaluation return over all runs for a single     data dictionary on an episodic environment. Bootstrap confidence intervals     are plotted as shaded regions.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of int         The list of hyperparameter settings indices to plot     smooth_over : int         The number of previous data points to smooth over. Note that this         is *not* the number of timesteps to smooth over, but rather the number         of data points to smooth over. For example, if you save the return         every 1,000 timesteps, then setting this value to 15 will smooth         over the last 15 readings, or 15,000 timesteps.     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     name : str, optional         The name of the agent, used for the legend     last_ind : int, optional         The index of the last element to plot in the returns list,         by default -1. This is useful if you want to plot many things on the         same axis, but all of which have a different number of elements. This         way, we can plot the first last_ind elements of each returns for each         agent.     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha channel for the plot, by default 0.1     colours : list of str         The colours to use for each plot of each hyperparameter setting     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     significance : float, optional         The significance level for the confidence interval, by default 0.01      Returns     -------     plt.figure, plt.Axes         The figure and axes of the plot      Raises     ------     ValueError         When an axis is passed but no figure is passed         When an appropriate number of colours is not specified to cover all         hyperparameter settings # Set up figure # if ax is None and fig is None: #     fig = plt.figure(figsize=figsize) #     ax = fig.add_subplot() # if xlim is not None: #     ax.set_xlim(xlim) # if ylim is not None: #     ax.set_ylim(ylim) Average {type_.title()} Return per Run with {conf_level}                  Confidence Intervals # Plot with bootstrap confidence interval # Plot based on colours Average {type_.title()} Return per Run with {conf_level}                  Confidence Intervals # ax.set_ylabel(ylabel) # ax.set_xlabel(xlabel) Use exp.get_mean_err with exp.bootstrap_conf Gets the mean return per episode along with the bootstrap confidence     interval on an episode-by-episode basis over all runs.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : int         The hyperparameter setting to plot     smooth_over : int         The number of previous returns to smooth over. If less than 2, then         no smoothing is performed     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     significance : float, optional         The significance level for the confidence interval, by default 0.01      Returns     -------     3-tuple of array_like         The timesteps at each return to plot, the mean return to plot, and         the confidence interval for each episode to plot as a shaded region # So the linter doesn't have a temper tantrum # Determine the timesteps to plot at # Get the mean over all episodes per evaluation step (for online # returns, this axis will have length 1 so we squeeze it) # Smooth over returns # Get the standard deviation of mean episodes per evaluation # step over all runs # Get the mean over all runs # If we do not keep the shape, then due to smoothing we lose some # episodes Plots the mean return over all runs and the return for each run for a list     of data dictionaries and hyperparameter indices      Plots both the mean return per episode (over runs) as well as the return     for each individual run (including "mini-runs", if a set of concurrent     episodes were run for all runs, e.g. multiple evaluation episodes per     run at set intervals)      Note that this function takes in a list of data dictionaries and will     plot the runs for each ind (which is a list of lists, where each super-list     refers to a data dictionary and each sub-list refers to the indices for     that data dictionary to plot).      Example     -------     plot_mean_with_runs([sac_data, linear_data], "train", [[3439], [38, 3]],     smooth_over=[5, 2], names=["SAC", "LAC"], figsize=(12, 6), alpha=0.2,     plot_avg=True, env_type="episodic")      will plot hyperparameter index 3439 for the sac_data, smoothing over the     last 5 episodes, and the label will have the term "SAC" in it; also plots     the mean and each individual run on the linear_data for hyperparameter     settings 38 and 3, smoothing over the last 2 episodes for each and with     the term "LAC" in the labels.      Parameters     ----------     data : list of dict         The Python data dictionaries generated from running main.py for the         agents     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of iter of int         The list of lists of hyperparameter settings indices to plot for         each agent. For example [[1, 2], [3, 4]] means that the first agent         plots will use hyperparameter settings indices 1 and 2, while the         second will use 3 and 4.     smooth_over : list of int         The number of previous data points to smooth over for the agent's         plot for each data dictionary. Note that this is *not* the number of         timesteps to smooth over, but rather the number of data points to         smooth over. For example, if you save the return every 1,000         timesteps, then setting this value to 15 will smooth over the last         15 readings, or 15,000 timesteps. For example, [1, 2] will mean that         the plots using the first data dictionary will smooth over the past 1         data points, while the second will smooth over the passed 2 data         points for each hyperparameter setting.     figsize : tuple(int, int)         The size of the figure to plot     names : list of str         The name of the agents, used for the legend     colours : list of list of str, optional         The colours to use for each hyperparameter settings plot for each data         dictionary, by default None     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha to use for plots of the runs, by default 0.1     plot_avg : bool, optional         If concurrent episodes are executed in each run (e.g. multiple         evaluation episodes are run at set intervals), then whether to plot the         performance of each separately or the average performance over all     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     fig : plt.Figure         The figure to plot on     ax : plt.Axes         The axis to plot on     detrend : str, optional         How to detrend the runs for each data, by default None. One of         "difference", "linear", or "mean"      Returns     -------     tuple of plt.figure, plt.Axes         The figure and axis plotted on # Set the colours to be default if not set # Set up figure # fig = plt.figure(figsize=figsize) # ax = fig.add_subplot() # Plot for each data dictionary given Plots the mean return over all runs and the return for each run for a     single data dictionary and for each in a list of hyperparameter settings.      Similar to plot_mean_with_runs, except that this function takes in only     a single data dictionary.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of int         The list of hyperparameter settings indices to plot for         each agent. For example [1, 2] means that the agent plots will use         hyperparameter settings indices 1 and 2.     smooth_over : int         The number of previous data points to smooth over for the agent's         plot. Note that this is *not* the number of timesteps to smooth over,         but rather the number of data points to smooth over. For example,         if you save the return every 1,000 timesteps, then setting this value         to 15 will smooth over the last 15 readings, or 15,000 timesteps.     figsize : tuple(int, int)         The size of the figure to plot     name : str         The name of the agents, used for the legend     colours : list of list of str, optional         The colours to use for each hyperparameter settings plot for each data         dictionary, by default None     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha to use for plots of the runs, by default 0.1     plot_avg : bool, optional         If concurrent episodes are executed in each run (e.g. multiple         evaluation episodes are run at set intervals), then whether to plot the         performance of each separately or the average performance over all     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     detrend : str, optional         How to detrend the runs for each data, by default None. One of         "difference", "linear", or "mean"      Returns     -------     tuple of plt.figure, plt.Axes, list of str, list of mpl.Lines2D         The figure and axis plotted on as well as the list of strings to use         as labels and the list of lines to include in the legend # Set up figure and axis # Store the info to keep in the legend # Plot each selected hyperparameter setting in the data dictionary Plots the mean return over all runs and the return for each run for a     single data dictionary and a single hyperparameter setting.      Similar to _plot_mean_with_runs, except that this function takes in only     a single data dictionary and a single hyperparameter setting.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : int         The hyperparameter settings indices to plot for the agent. For example         5 means that the agent plots will use hyperparameter settings index 5.     smooth_over : int         The number of previous data points to smooth over for the agent's         plot. Note that this is *not* the number of timesteps to smooth over,         but rather the number of data points to smooth over. For example,         if you save the return every 1,000 timesteps, then setting this value         to 15 will smooth over the last 15 readings, or 15,000 timesteps.     figsize : tuple(int, int)         The size of the figure to plot     name : str         The name of the agents, used for the legend     colours : list of list of str, optional         The colours to use for each hyperparameter settings plot for each data         dictionary, by default None     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha to use for plots of the runs, by default 0.1     plot_avg : bool, optional         If concurrent episodes are executed in each run (e.g. multiple         evaluation episodes are run at set intervals), then whether to plot the         performance of each separately or the average performance over all     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'     detrend : str, optional         How to detrend the runs for each data, by default None. One of         "difference", "linear", or "mean"      Returns     -------     tuple of plt.figure, plt.Axes, list of str, list of mpl.Lines2D         The figure and axis plotted on as well as the list of strings to use         as labels and the list of lines to include in the legend # Set up figure and axis # Determine the timesteps to plot at # Plot the average reward # Get returns # If concurrent episodes are run in each run then average them if # appropriate # Smooth returns # Plot the average reward # Detrend the data if appropriate # Determine whether to plot episodes or timesteps on the x-axis, which is # dependent on the environment type # episodes # Plot each run # Plot the mean # Store legend identifiers for the run # Set up the legend variables for the mean over all runs Plots the average training or evaluation return over all runs with standard     error.      Given a list of data dictionaries of the form returned by main.py, this     function will plot each episodic return for the list of hyperparameter     settings ind each data dictionary. The ind argument is a list, where each     element is a list of hyperparameter settings to plot for the data     dictionary at the same index as this list. For example, if ind[i] = [1, 2],     then plots will be generated for the data dictionary at location i     in the data argument for hyperparameter settings ind[i] = [1, 2].     The smooth_over argument tells how many previous data points to smooth     over      Parameters     ----------     data : list of dict         The Python data dictionaries generated from running main.py for the         agents     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of iter of int         The list of lists of hyperparameter settings indices to plot for         each agent. For example [[1, 2], [3, 4]] means that the first agent         plots will use hyperparameter settings indices 1 and 2, while the         second will use 3 and 4.     smooth_over : list of int         The number of previous data points to smooth over for the agent's         plot for each data dictionary. Note that this is *not* the number of         timesteps to smooth over, but rather the number of data points to         smooth over. For example, if you save the return every 1,000         timesteps, then setting this value to 15 will smooth over the last         15 readings, or 15,000 timesteps. For example, [1, 2] will mean that         the plots using the first data dictionary will smooth over the past 1         data points, while the second will smooth over the passed 2 data         points for each hyperparameter setting.     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     names : list of str         The name of the agents, used for the legend     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha channel for the plot, by default 0.1     colours : list of list of str         The colours to use for each hyperparameter settings plot for each data         dictionary     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'      Returns     -------     plt.figure, plt.Axes         The figure and axes of the plot # Set the colours to be default if not set # Set up figure # if ax is None and fig is None: #     fig = plt.figure(figsize=figsize) #     ax = fig.add_subplot() # if xlim is not None: #     ax.set_xlim(xlim) # if ylim is not None: #     ax.set_ylim(ylim) # Track the total timesteps per hyperparam setting over all episodes and # the cumulative timesteps per episode per data dictionary (timesteps # should be consistent between all hp settings in a single data dict) # Find the minimum of total trained-for timesteps. Each plot will only # be plotted on the x-axis until this value # For each data dictionary, find the minimum index where the timestep at # that index is >=  minimum timestep # Since indexing will stop right before the minimum, increment it # Plot all data for all HP settings, only up until the minimum index # fig, ax = None, None Plots the average training or evaluation return over all runs for a single     data dictionary on a continuing environment. Standard error     is plotted as shaded regions.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of int         The list of hyperparameter settings indices to plot     smooth_over : int         The number of previous data points to smooth over. Note that this         is *not* the number of timesteps to smooth over, but rather the number         of data points to smooth over. For example, if you save the return         every 1,000 timesteps, then setting this value to 15 will smooth         over the last 15 readings, or 15,000 timesteps.     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     name : str, optional         The name of the agent, used for the legend     last_ind : int, optional         The index of the last element to plot in the returns list,         by default -1. This is useful if you want to plot many things on the         same axis, but all of which have a different number of elements. This         way, we can plot the first last_ind elements of each returns for each         agent.     timestep_multiply : array_like of float, optional         A value to multiply each timstep by, by default None. This is useful if         your agent does multiple updates per timestep and you want to plot         performance vs. number of updates.     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha channel for the plot, by default 0.1     colours : list of str         The colours to use for each plot of each hyperparameter setting     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'      Returns     -------     plt.figure, plt.Axes         The figure and axes of the plot      Raises     ------     ValueError         When an axis is passed but no figure is passed         When an appropriate number of colours is not specified to cover all         hyperparameter settings # Set up figure #     fig = plt.figure(figsize=figsize) #     ax = fig.add_subplot() # if xlim is not None: #     ax.set_xlim(xlim) # if ylim is not None: #     ax.set_ylim(ylim) # Plot with the standard error # Plot based on colours Plots the average training or evaluation return over all runs for a     single data dictionary on an episodic environment. Plots shaded retions     as standard error.      Parameters     ----------     data : dict         The Python data dictionary generated from running main.py     type_ : str         Which type of data to plot, one of "eval" or "train"     ind : iter of int         The list of hyperparameter settings indices to plot     smooth_over : int         The number of previous data points to smooth over. Note that this         is *not* the number of timesteps to smooth over, but rather the number         of data points to smooth over. For example, if you save the return         every 1,000 timesteps, then setting this value to 15 will smooth         over the last 15 readings, or 15,000 timesteps.     fig : plt.figure         The figure to plot on, by default None. If None, creates a new figure     ax : plt.Axes         The axis to plot on, by default None, If None, creates a new axis     figsize : tuple(int, int)         The size of the figure to plot     name : str, optional         The name of the agent, used for the legend     last_ind : int, optional         The index of the last element to plot in the returns list,         by default -1. This is useful if you want to plot many things on the         same axis, but all of which have a different number of elements. This         way, we can plot the first last_ind elements of each returns for each         agent.     xlim : float, optional         The x limit for the plot, by default None     ylim : float, optional         The y limit for the plot, by default None     alpha : float, optional         The alpha channel for the plot, by default 0.1     colours : list of str         The colours to use for each plot of each hyperparameter setting     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'. By default         'continuing'      Returns     -------     plt.figure, plt.Axes         The figure and axes of the plot      Raises     ------     ValueError         When an axis is passed but no figure is passed         When an appropriate number of colours is not specified to cover all         hyperparameter settings # Set up figure # Plot with the standard error # data has consistent # of episodes, so treat as env_type="continuing" # Plot based on colours Plots the hyperparameter sensitivity for hp_name, where the combo argument     determines which hyperparameters are held constant.      Given hyperparameters a, b, and c, let's say we want to get all     hyperparameter settings indices where a varies, and b and c are constant.     if a, b, and c can each be 1 or 2, then there are four ways that a can     vary with b and c remaining constant:          [             ((a=1, b=1, c=1), (a=2, b=1, c=1)),         combo = 0             ((a=1, b=2, c=1), (a=2, b=2, c=1)),         combo = 1             ((a=1, b=1, c=2), (a=2, b=1, c=2)),         combo = 2             ((a=1, b=2, c=2), (a=2, b=2, c=2))          combo = 3         ]      The combo argument indexes into this list of hyperparameter settings      Parameters     ----------     dir : str         The directory containing all data dictionaries generated from running         main.py, separated into one file per HP index     type_ : str         Which type of data to plot, one of "eval" or "train"     combo : int         Determines the values of the constant hyperparameters. Given that         only one hyperparameter may vary, there are many different sets         having this hyperparameter varying with all others remaining constant         since each constant hyperparameter may take on many values. This         argument cycles through all sets of hyperparameter settings indices         that have only one hyperparameter varying and all others constant.     env_config : dict         The environment configuration file as a Python dictionary     agent_config : dict         The agent configuration file as a Python dictionary     figsize : tuple(int, int)         The size of the figure     fig : plt.figure, optional         The figure to plot on, by default None. If None, a new figure is         created     ax : plt.Axes, optional         The axis to plot on, by default None. If None, a new axis is created.     annotate : bool, optional         Whether or not to annoate each point drawn with its y-value     colour : str         The colour to use in the plot, by default None. If None, then default         values will be used.      Returns     -------     tuple of plt.figure, plt.Axes         The figure and axis that was plotted on      Raises     ------     ValueError         If only one of fig, ax are specified # Set up the figure Plot hyperparameter sensitivity by hyperparameter settings index.      Plots the hyperparameter sensitivity plot where all     constant hyperparameters are equal to those defined by the     argument hyperparameter settings index. The only hyperparameter that     is allowed to vary is the one corresponding to hp_name.      Parameters     ----------     dir : str         The directory containing all data dictionaries generated by running         main.py     type_ : str         Whether to plot the training or evaluation data, one of         'train', 'eval'     hp_name : str         The name of the hyperparameter to plot the sensitivity of     config : dict         A Python dictionary of the agent configuration file used to run         the experiment     ind : int         The hyperparameter settings index that should define the         values of all constant hyperparameters. Only the hyperparameter         defined by hp_name is allowed to vary, all others must be equal to         the values of those defined by the settings index ind.     env_config : dict         The environment configuration file as a Python dictionary     agent_config : dict         The agent configuration file as a Python dictionary     figsize : tuple(int, int)         The size of the figure     annotate : bool, optional         Whether or not to annoate each point drawn with its y-value      Returns     -------     tuple of plt.figure, plt.Axes         The figure and axis that was plotted on # Get all the combinations of hyperparameters with only hp_name varying # Get the index (into the list of tuples with only hp_name varying) of # the tuple which has hp_name varying only, but all other constant # hyperparameters equal to those defined by hyperparam setting number ind. # So hp_combo_constant_hps_equal_ind is the hyperparam combination with # all constant hyperparams equal to those in settings number ind. # Plot the hyperparameter sensitivity plot using all constant # hyperparameters equal to those defined in the settings # number ind We no longer store the state seen at each timestep             as this was becoming very memory intensive. We are storing the             states seen at each timestep as tile coded counts, similar to the             states used in updates Plots the state visitation distribution during training      Given a data dictionary, plots the state visitation distribution of a     single hyperparameter setting.      Parameters     ----------     data : dict         The data dictionary containing the runs of a single hyperparameter         setting     figsize : tuple, optional         The size of the figure to plot, by default (15, 10)     max_ : float, optional         The cutoff threshold for updates, by default 1.0. Any states that have         updates larger than this will be capped at this value. Note that         update counts are normalized between [0, 1].     xlabel : str, optional         The label for the x axis, by default ""     ylabel : str, optional         The label for the y axis, by default ""     title : str, optional         The title of the plot, by default ""     hp_ind : int, optional         The hyperparameter settings index in the data dictionary to use for         the plot, by default -1. If less than 0, then the first hyperparameter         setting in the dictionary is used.     rasterize : bool, optional         Whether or not to rasterize the graphics, by default True. Useful         to remove gridlines between cells when saving as a PDF or SVG; in         this case, set to False.     env_type : str, optional         The type of environment, one of 'continuing', 'episodic'; by default         "continuing"      Returns     -------     plt.figure, plt.Axes, array_like         The figure and axis plotted on and the array of state counts # Choose the hyperparameter setting if not st # Set up the TileCoder # If the type is continuing, then all data fits nicely into a single matrix # and we can easily calculate everything using array operations # If episodic, then since runs/episodes have different lengths, we instead # have an array, where each row /col is a different length (list of states) # since each run/episode has a different length. So, we have to deal with # this case differently, looping over all runs, then all states per # episode per run # For each episode, get the states and tile code, we go through # episodes separately since each has its own length Plot a heatmap of the states used in each update      This function will only plot a 2D heatmap of states. The first state     coordinate will refer to the rows of the heatmap, while the second will     refer to the columns. The minimum value of the first state coordinate     will be the top of the heatmap, while the max will be the bottom.     Similarly, the maximum of the second coordinate will be the right of     the plot, while the minimum will be the left.      This function assumes a vector that has as each element, the number of     updates performed to that corresponding state. If the state space is     continuous, simply use a high resolution tile coder and discretize the     state space.      Parameters     ----------     data : dict         The data dictionary containing the runs of a single hyperparameter         setting     figsize : tuple, optional         The size of the figure to plot, by default (15, 10)     max_ : float, optional         The cutoff threshold for updates, by default 1.0. Any states that have         updates larger than this will be capped at this value. Note that         update counts are normalized between [0, 1].     xlabel : str, optional         The label for the x axis, by default ""     ylabel : str, optional         The label for the y axis, by default ""     updates : bool, optional         Whether to plot the states used in updates or the states visited, by         default True.     use_next_state : bool, optional         Whether to plot the state or next state in the SARSA update,         by default False     hp_ind : int, optional         The hyperparameter settings index in the data dictionary to use for         the plot, by default -1. If less than 0, then the first hyperparameter         setting in the dictionary is used.     rasterize : bool, optional         Whether or not to rasterize the graphics, by default True. Useful         to remove gridlines between cells when saving as a PDF or SVG; in         this case, set to False.     time_ind : int, optional         In each experiment, the state update heatmap is stored at multiple         timesteps at certain intervals. This parameter controls which of the         update heatmaps to plot, with 0 being the first and -1 being the last         heatmap saved, by default -1.     use_interval_only : bool, optional         Whether to plot only the states visited between the time_ind reading         and the previous reading. That is, whether or not to plot only the         states used in updates or visited between the argument time_ind and         the previous, but none before. By default False.     title : str, optional         The title of the figure, by default ""      Returns     -------     plt.figure, plt.Axes, array_like         The figure and axis plotted on and the array of state counts      Raises     ------     ValueError         If attempting to plot the next states, but not plotting states used in         updates can only plot next states when plotting states                          used in updates, set 'updates=True' # Choose the hyperparameter setting if not st # Get the number of updates per each state per run # Get the mean states used in each update or visited at each timestep # over all runs # Calculate the states seen only for the current interval # Plot # Get the states used in each update or visited at each timestep # Get the average number of updates per state over all runs Plots a heatmap of states, where states are separated into bins      This function is designed to work with the TileCoder object. It takes in     a flattened array, where each index in the array refers to a sequential     bin of states. The element at each position is the number of times that     state was updated or encountered.      Parameters     ----------     visited_states : array_like         The flattened array of counts, where each index in the array         corresponds to the index of the bin     env_config : dict         The environment configuration file for the environment that the         experiment was run on     figsize : tuple, optional         The size of the figure to plot, by default (15, 10)     max_ : float, optional         The cutoff threshold for updates, by default 1.0. Any states that have         updates larger than this will be capped at this value. Note that         update counts are normalized between [0, 1].     xlabel : str, optional         The label for the x axis, by default ""     ylabel : str, optional         The label for the y axis, by default ""     title : str, optional         The title of the figure, by default ""     rasterize : bool, optional         Whether or not to rasterize the graphics, by default True. Useful         to remove gridlines between cells when saving as a PDF or SVG; in         this case, set to False.      Returns     -------     plt.figure, plt.Axes, array_like         The figure and axis plotted on and the array of state counts # Threshold the updates # Determine the bins per state dimension # Plot the heatmap # Set the axis tick locations # Label axis ticks Plot the action heatmap with a barchart above it      Parameters     ----------     data : dict         The data dictionary containing the runs of a single hyperparameter         setting     figsize : tuple, optional         The size of the figure to plot, by default (15, 10)     max_ : int, optional         The cutoff threshold for updates, by default 10. Any states that have         updates larger than this will be capped at this value.     xlabel : str, optional         The label for the x axis, by default ""     ylabel : str, optional         The label for the y axis, by default ""     hp_ind : int, optional         The hyperparameter settings index in the data dictionary to use for         the plot, by default -1. If less than 0, then the first hyperparameter         setting in the dictionary is used.     rasterize : bool, optional         Whether or not to rasterize the graphics, by default True. Useful         to remove gridlines between cells when saving as a PDF or SVG; in         this case, set to False.     ratio : tuple, optional         The ratio of bar char size to heatmap size, by default (2, 1)     time_ind : int, optional         In each experiment, the state update heatmap is stored at multiple         timesteps at certain intervals. This parameter controls which of the         update heatmaps to plot, with 0 being the first and -1 being the last         heatmap saved, by default -1.      Returns     -------     plt.figure, plt.Axes, plt.Axes         The figure, barchart axis, and heatmap axis plotte on # Set up the plot # Get the number updates per action bin # Get the average number of upates per state over all runs # Plot the barchart # Set the axis limits for the barchart # Threshold the updates # Determine the bins per action dimension # Plot the heatmap # Set the heatmap axis tick locations # Label heatmap axis ticks # Place the heatmaps ticks Factor to get the functions to label the x and y ticks      Parameters     ----------     num_ticks : int         The total number of ticks along the axis     env_config : dict         The environment configuration file as a dict     type_ : str, optional         The type of plot to get the ticks of the heatmap for, one of 'state',         'action', by default 'state'      Returns     -------     2-tuple of func         The function to generate the x tick labels and the function to         generate the y tick labels      Raises     -------     NotImplementedError         If trying to get tick labels for an unsupported environment Generates the tick labels for the heatmap on a 2D-state environment      Parameters     ----------     num_ticks : int         The total number of ticks along the axis     env_config : dict         The environment configuration file as a dict     type_ : str         The type of plot to get the ticks of the heatmap for, one of 'state',         'action'      Returns     -------     2-tuple of func         The function to generate the x tick labels and the function to         generate the y tick labels Plot the surface of states used in each update      Note that unlike the heatmap, the updates plotted for the surface are not     normalized betwee 0 and 1.      Parameters     ----------     data : dict         The data dictionary containing the runs of a single hyperparameter         setting     figsize : tuple, optional         The size of the figure to plot, by default (15, 10)     type_ : str, optional         The type of surface to plot, by default "surface". One of 'surface',         'wireframe', or 'bar'     max_ : int, optional         The cutoff threshold for updates, by default 100. Any states that have         updates larger than this will be capped at this value.     smooth : float, optional         The amount to smoothen the plot, by default 1.5. This value corresponds         to the standard deviation of a Gaussian convolution kernel applied to         the surface.     view : 2-tuple of int, optional         The view in 3D space, by default (30, -45)     linewidth : int, optional         The width of the lines separating cells, by default 0     colour : str         The colour to use for the plot, by default None. If None, uses the         default colours.     time_ind : int, optional         In each experiment, the state update heatmap is stored at multiple         timesteps at certain intervals. This parameter controls which of the         update heatmaps to plot, with 0 being the first and -1 being the last         heatmap saved, by default -1.      Returns     -------     plt.figure, plt.Axes3D         The figure and axis plotten on # Get the min/max x/y values # Determine the width, depth, and height of barx Plots the distribution of returns on either an episodic or continuing     environment      Parameters     ----------     data : dict         The data dictionary containing the runs of a single hyperparameter         setting     type_ : str, optional         The type of surface to plot, by default "surface". One of 'surface',         'wireframe', or 'bar'     hp_ind : int, optional         The hyperparameter settings index in the data dictionary to use for         the plot, by default -1. If less than 0, then the first hyperparameter         setting in the dictionary is used.     bins : Iterable, int         The bins to use for the plot. If an Iterable, then each value in the         Iterable is considered as a cutoff for bins. If an integer, separates         the returns into that many bins     figsize : tuple, optional         The size of the figure to plot, by default (15, 10)     xlim : 2-tuple of float, optional         The cutoff points for the x-axis to plot between, by default None     ylim : 2-tuple of float, optional         The cutoff points for the y-axis to plot between, by default None      Returns     -------     plt.figure, plt.Axes3D         The figure and axis plotten on # Get the episode returns for each run # Plot relative frequency on the y-axis Plots the mean returns of one hyperparameter vs another.      Given a list of hyperparameter indices and two hyperparameter names,     plots the mean returns generated by varying the second hyperparameter,     and leaving the first hyperparameter fixed as each hyperparameter index's     value. The plot is a heatmap. In this way, each row of the heatmap will     be a certain HP setting with a fixed HP, while the columns will be the     mean returns given the second HP varying for each HP index.      Parameters     ----------     dir_ : str         The directory which contains all data dictionaries     type_ : str         The type of returns to consider, one of 'train' or 'eval'     hp1 : str         The name of the hyperparameter setting to vary on the x-axis     hp2 : str         The name of the hyperparameter to plot on the y-axis for each         hyperparameter settings index plotted     inds : Iterable of int         An iterable of hyperparameter indices used to plot     env_config : dict         The environment configuration file used to run the experiment     agent_config : dict         The agent configuration file used to run the experiment     figsize : tuple, optional         The size of the figure, by default (12, 8)     annotate : bool, optional         Whether each box in the heatmap should be annotated with its value,         by default True     fmt : str, optional         If annotating, the format to use, by default ".2g"     rasterize : bool, optional         Whether or not the final heatmap should be rasterized, by default False     cmap : str, optional         The colourmap to use to plot, by default CMAP     linewidths : int, optional         The width of the lines used to separate boxes in the heatmap, by         default 0     env_type : str, optional         Whether the environment is continuing or episodic, by default         "continuing". One of 'episodic', 'continuing'     default_value : int, optional         The default value to use for heatmap boxes when no valid second         hyperparameter can be used for the first hyperparameter, by default         -1000. This happens, for example in hyperparameter settings that would         have batch size > replay capacity.     after : int, optional         Only consider returns after this episode, by default 0     before : int, optional         Only consider returns before this episode, by default None     fig : plt.Figure, optional         The figure to plot on, by default None. If None, a new Figure is         created.     ax : plt.Axes, optional         The axis to plot on, by default None. If None, a new Axes is created.     vmin : float, optional         Minimum value to anchor the colourmap to, otherwise it is inferred         from the data.     vmax : float, optional         Maximum value to anchor the colourmap to, otherwise it is inferred         from the data.      Returns     -------     plt.Figure, plt.Axes         The figure and axis plotted on      Raises     ------     ValueError         If the figure is not specified but the axis is         If the type of return considered is not one of "train" or "eval"         If an illegal environment type is given # Get all the combinations of hyperparameters with only hp_name varying # Get the index (into the list of tuples with only hp_name varying) of # the tuple which has hp_name varying only, but all other constant # hyperparameters equal to those defined by hyperparam setting ind. # So hp_combo_constant_hps_equal_ind is the hyperparam combination with # all constant hyperparams equal to those in settings number ind. # Get the mean returns for each hyperparameter setting # No hyperparameter setting for this combo of hyperparams # e.g. batch size > buffer size # Use the default value for the return # Open the data file and compute the mean returns # Calculate and save the mean return for this HP setting # Create the figure and axis # Plot the heatmap and set the labels # Define the tick generation function # Place the heatmaps ticks Plots the mean returns of each hyperparameter in a 2D matrix.      Given a 2D array of hyperparameter settings indices, plots the mean return     of each hyperparameter settings index as a heatmap. This function assumes     that each row in the matrix corresponds to a different value of one     hyperparameter, and each column corresponds to the values of a different     hyperparameter, and these hyperparameters are plotted as the x labels and     y labels.      Note that any hyperparameter setting less than 0 is ignored, so these     indices can be used to plot the default value in the heatmap.      Parameters     ----------     dir_ : str         The directory which contains all data dictionaries     type_ : str         The type of returns to consider, one of 'train' or 'eval'     hp1 : str         The name of the hyperparameter setting to vary on the x-axis     hp2 : str         The name of the hyperparameter to plot on the y-axis for each         hyperparameter settings index plotted     inds : 2D array_like of int         Matrix of hyperparameter indices to plot     env_config : dict         The environment configuration file used to run the experiment     agent_config : dict         The agent configuration file used to run the experiment     figsize : tuple, optional         The size of the figure, by default (12, 8)     annotate : bool, optional         Whether each box in the heatmap should be annotated with its value,         by default True     fmt : str, optional         If annotating, the format to use, by default ".2g"     rasterize : bool, optional         Whether or not the final heatmap should be rasterized, by default False     cmap : str, optional         The colourmap to use to plot, by default CMAP     linewidths : int, optional         The width of the lines used to separate boxes in the heatmap, by         default 0     env_type : str, optional         Whether the environment is continuing or episodic, by default         "continuing". One of 'episodic', 'continuing'     default_value : int, optional         The default value to use for heatmap boxes when no valid second         hyperparameter can be used for the first hyperparameter, by default         -1000. This happens, for example in hyperparameter settings that would         have batch size > replay capacity.     after : int, optional         Only consider returns after this episode, by default 0     before : int, optional         Only consider returns before this episode, by default None     fig : plt.Figure, optional         The figure to plot on, by default None. If None, a new Figure is         created.     ax : plt.Axes, optional         The axis to plot on, by default None. If None, a new Axes is created.     vmin : float, optional         Minimum value to anchor the colourmap to, otherwise it is inferred         from the data.     vmax : float, optional         Maximum value to anchor the colourmap to, otherwise it is inferred         from the data.      Returns     -------     plt.Figure, plt.Axes         The figure and axis plotted on      Raises     ------     ValueError         If the figure is not specified but the axis is         If the type of return considered is not one of "train" or "eval"         If an illegal environment type is given # Ensure the hyperparameter settings have consistent hp values in each # row and column # Open the data file and compute the mean returns # Calculate and save the mean return for this HP setting # Create the figure and axis # Plot the heatmap and set the labels # Define the tick generation function # Place the heatmaps ticks Check the columns of a matrix of hyperparameters indices for consistent     hyperparameter values.      Checks each row of a matrix of hyperparameter settings to ensure that each     hyperparameter setting in the row has the same value for a certain     hyperparameter. This is needed to ensure consistency when plotting a     matrix of hyperparameter settings as a heatmap.      Note that any hyperparameter setting less than 0 is ignored, so these     indices can be used to plot the default value in the heatmap.      Parameters     ----------     agent_config : dict         The agent configuration file used to run the experiment     env_config : dict         The environment configuration file used to run the experiment     inds : 2D array_like of int         Matrix of hyperparameter indices to check to ensure that each row         has hyperparameter settings that have the same hyperparameter value         for a specific hyperparameter     hp : str         The hyperparameter to ensure is consistent between hyperparameter         settings on each row.      Raises     ------     IndexError         If the hyperparameter indices do not refer to settings with the same         value for the argument hyperparameter.     ValueError         If the array is not 2D. # The value that all rows should have as HP value # No hyperparameter setting in this position # Get the value of the hyperparameter for the index in the current # position # This is the first HP value encountered, use it as a reference # value for all other HPs in the row # Illegal HP value Plot the average absolute value of a state feature per episode      Parameters     ----------     data : dict         The data dictionary containing the runs of a single hyperparameter         setting     after : int, optional         Consider the state features only after this timestep for each         episode, by default 0     figsize : tuple, optional         The size of the figure, by default (12, 6)     state_ind : int, optional         The index of the state feature to plot, by default 0     colour : str, optional         The colour for the plot, by default None     alpha : float, optional         The alpha value for the individual runs, by default 0.1     title : str, optional         The title for the plot, by default ""     xlabel : str, optional         The label for the x axis, by default ""     ylabel : str, optional         The label for the y axis, by default ""     xlim : 2-tuple of float, optional         The limits for the x axis, by default None     ylim : 2-tuple of float, optional         The limits for the y axis, by default None      Returns     -------     plt.figure, plt.Axes         The figure and Axes plotted on # Calculate the average performance, averaging over runs and then # timesteps to get the average angle per episode # Get the average angle per episode for each run # Get some prelimiary variables for the legend # Get the mean over all runs and timesteps for each episode # Set axis limits # Plot the average angle per episode for all runs # Plot average average absolute angle for each episode # (averaged over timesteps) for each separate run Recursively turns elements of an Iterable into strings representing     colours.      This function will turn each element of an Iterable into strings that     represent colours, recursively. If the elements of an Iterable are     also Iterable, then this function will recursively descend all the way     through every Iterable until it finds an Iterable with non-Iterable     elements. These elements will be replaced by strings that represent     colours. In effect, this function keeps the data structure, but replaces     non-Iterable elements by strings representing colours. Note that this     funcion assumes that all elements of an Iterable are of the same type,     and so it only checks if the first element of an Iterable object is     Iterable or not to stop the recursion.      Parameters     ----------     iter_ : collections.Iterable         The top-level Iterable object to turn into an Iterable of strings of         colours, recursively.      Returns     -------     list of list of ... of strings         A data structure that has the same architecture as the input Iterable         but with all non-Iterable elements replaced by strings. # Calculate the number of lists at the current level to go through # Return a list of colours if the elements of the list are not lists # For each list at the current level, get the colours corresponding to # this level Plots a curve with a shaded region.      Parameters     ----------     ax : plt.Axes         The axis to plot on     x : Iterable         The points on the x-axis     y : Iterable         The points on the y-axis     region : list or array_like         The region to shade about the y points. The shaded region will be         y +/- region. If region is a list or 1D np.ndarray, then the region         is used both for the + and - portions. If region is a 2D np.ndarray,         then the first row will be used as the lower bound (-) and the         second row will be used for the upper bound (+). That is, the region         between (lower_bound, upper_bound) will be shaded, and there will be         no subtraction/adding of the y-values.     colour : str         The colour to plot with     label : str         The label to use for the plot     alpha : float         The alpha value for the shaded region Resets the colours offset