plotting.py

# Plots for use in the App and beyond

import numpy as np
import matplotlib.pyplot as plt
import os
import math
import utilities
import model_core
from matplotlib.backends.backend_pdf import PdfPages

# some constants and indices
congr_upright_indices = [0, 3]
congr_updown_indices = [5, 6]
incongr_upright_indices = [1, 2]
incongr_updown_indices = [4, 7]


def _timesteps_transport_default():
    """Return the model transport duration, falling back to the current default."""
    return getattr(model_core, "timesteps_transport", 100)


def plot_focus(trial_data):
    """Plot event focus values over one trial."""
    # get relevant parts from the input data
    timesteps_reach = trial_data[0,5]
    # cut here, depending on timesteps_reach, etc.
    F = trial_data[:(int(timesteps_reach+_timesteps_transport_default())),19:24] #(?)

    fig, ax = plt.subplots()
    ax.plot(F[:,1], label="F[1] (reach)")
    ax.plot(F[:,2], label="F[2] (grasp)")
    ax.plot(F[:,3], label="F[3] (transp.)")
    ax.plot(F[:,4], label="F[4] (put-d.)")
    ax.plot(F[:,0], label="F[0] (idle)")
    ax.legend()
    ax.set_xlabel("Timestep")
    ax.set_ylabel("Focus value")
    ax.axvline(timesteps_reach, color='grey', alpha=0.5)
    return fig


def _focus_runs_by_condition(trials_data):
    """Reshape uncompressed trial data into run x condition x time x feature."""
    trials_data = np.asarray(trials_data)
    if trials_data.ndim != 3:
        raise ValueError(
            "Expected trials_data with shape (n_trials, n_timesteps, n_features)."
        )
    if trials_data.shape[0] < 8 or trials_data.shape[0] % 8 != 0:
        raise ValueError(
            "Expected uncompressed data ordered as 8 trial conditions per simulation run. "
            "Use the *_all.npz res_data output, not the 8-condition averaged file."
        )
    return trials_data.reshape(
        trials_data.shape[0] // 8,
        8,
        trials_data.shape[1],
        trials_data.shape[2],
    )


def _mean_and_band(values, band="sd"):
    """Return mean and symmetric uncertainty band across simulation runs."""
    values = np.asarray(values, dtype=float)
    mean = np.nanmean(values, axis=0)
    n_runs = values.shape[0]
    if n_runs < 2:
        spread = np.zeros_like(mean)
    else:
        sd = np.nanstd(values, axis=0, ddof=1)
        if band == "sd":
            spread = sd
        elif band == "sem":
            spread = sd / np.sqrt(n_runs)
        elif band in ("ci", "ci95", "95ci"):
            spread = 1.96 * sd / np.sqrt(n_runs)
        else:
            raise ValueError("band must be one of 'sd', 'sem', or 'ci'.")
    return mean, mean - spread, mean + spread


def plot_focus_uncertainty(
        trials_data,
        condition_groups=None,
        focus_indices=(1, 2),
        focus_labels=None,
        band="sd",
        colors=None,
        font_size=None,
):
    """Plot mean focus weights with across-run uncertainty bands.

    `trials_data` should be uncompressed `res_data`, where consecutive blocks
    of 8 trials are the 8 trial configurations for one simulation run. The
    default plot shows reach and grasp focus averaged over all trial conditions.
    """
    if condition_groups is None:
        condition_groups = [("All trials", list(range(8)))]
    if focus_labels is None:
        default_labels = {
            0: "F[0] (idle)",
            1: "F[1] (reach)",
            2: "F[2] (grasp)",
            3: "F[3] (transp.)",
            4: "F[4] (put-d.)",
        }
        focus_labels = [default_labels.get(i, f"F[{i}]") for i in focus_indices]
    if colors is None:
        colors = {
            0: "#9ecae1",
            1: "#1f77b4",
            2: "#ff7f0e",
            3: "#2ca02c",
            4: "#7f0000",
        }

    old_font_size = None
    if font_size is not None:
        old_font_size = plt.rcParams.get('font.size', None)
        plt.rcParams.update({'font.size': font_size})

    try:
        runs_by_condition = _focus_runs_by_condition(trials_data)
        n_groups = len(condition_groups)
        fig, axs = plt.subplots(1, n_groups, figsize=(6 * n_groups, 4.5), sharey=True)
        axs = np.atleast_1d(axs)

        for ax, (title, condition_indices) in zip(axs, condition_groups):
            condition_indices = list(condition_indices)
            grouped = runs_by_condition[:, condition_indices, :, :]
            timesteps_reach = int(np.round(np.nanmean(grouped[:, :, 0, 5])))
            n_timesteps = min(
                timesteps_reach + _timesteps_transport_default(),
                grouped.shape[2],
            )
            x = np.arange(n_timesteps)

            for focus_idx, label in zip(focus_indices, focus_labels):
                values = grouped[:, :, :n_timesteps, 19 + focus_idx]
                per_run = np.nanmean(values, axis=1)
                mean, low, high = _mean_and_band(per_run, band=band)
                color = colors.get(focus_idx, None)
                ax.plot(x, mean, label=label, color=color, linewidth=2)
                ax.fill_between(x, low, high, color=color, alpha=0.2, linewidth=0)

            ax.axvline(timesteps_reach, color='grey', alpha=0.5)
            ax.set_xlim(0, n_timesteps - 1)
            ax.set_ylim(0, 1)
            if n_groups > 1:
                ax.set_title(title)
            ax.set_xlabel("Timestep")
            ax.set_ylabel("Focus value")

        handles, labels = axs[0].get_legend_handles_labels()
        fig.legend(
            handles,
            labels,
            #loc="upper center",
            #bbox_to_anchor=(0.5, 0.95),
            loc = "right",
            bbox_to_anchor=(0.95,0.85),
            ncol=len(focus_indices),
            #frameon=False,
        )
        band_labels = {
            "sd": "SD",
            "sem": "SEM",
            "ci": "95% CI",
            "ci95": "95% CI",
            "95ci": "95% CI",
        }
        #fig.suptitle(
        #    f"Mean focus weights across simulation runs ({band_labels[band]} bands)",
        #    y=0.99,
        #)
        fig.tight_layout()#rect=(0, 0, 1, 0.9))
        return fig
    finally:
        if font_size is not None and old_font_size is not None and old_font_size != font_size:
            plt.rcParams.update({'font.size': old_font_size})

def plot_inf_event(trial_data):
    """Plot the inferred event sequence for one trial."""
    timesteps_reach = trial_data[0,5]
    fig, ax = plt.subplots()
    ax.plot(trial_data[:(int(timesteps_reach)+_timesteps_transport_default()),18], label="inferred event", color="black")
    ax.axhline(0, color="lightblue", linestyle="dotted")
    ax.axhline(1, color="blue", linestyle="dotted")
    ax.axhline(2, color="orange", linestyle="dotted")
    ax.axhline(3, color="green", linestyle="dotted")
    ax.axhline(4, color="purple", linestyle="dotted")
    ax.set_xlabel("Timestep")
    ax.set_yticks([0, 1, 2, 3, 4])
    ax.set_yticklabels(["idle", "reach", "grasp", "transport", "put-down"], rotation=45)
    ax.axvline(timesteps_reach, color='grey', alpha=0.5)
    return fig

def plot_base(trial_data, timestep_bottle_appearance=None, timestep_stim=None, timesteps_reach=None):
    """Plot the main trial progression."""
    if not timesteps_reach or timesteps_reach is None:
        timesteps_reach = int(trial_data[0,5])
    # possibly cut the data to the relevant timesteps
    trial_data = trial_data[:(int(timesteps_reach)+_timesteps_transport_default()),:]
    # other data points
    x = trial_data[:,12] # position/distance
    orient = trial_data[:,11] # orientation
    if not timestep_bottle_appearance or timestep_bottle_appearance is None:
        timestep_bottle_appearance = int(trial_data[0,3])
    if not timestep_stim or timestep_stim is None:
        timestep_stim = int(trial_data[0,4])
    if trial_data[timestep_stim,1] == utilities.THUMB:
        if trial_data[timestep_stim,2] == utilities.LEFT:
            color = 'yellow'
        else:
            color = 'orange'
    else:
        if trial_data[timestep_stim,2] == utilities.LEFT:
            color = 'red'
        else:
            color = 'purple'

    fig, ax = plt.subplots()
    ax.plot(x, color="black", label="Position/Distance")
    ax.plot(orient, color="lightblue", linestyle="solid", label="Orientation")
    ax.axvline(timestep_bottle_appearance, 0, 1, color='green',
               linestyle='dashed' if trial_data[timestep_bottle_appearance,0] == utilities.UPRIGHT else 'dotted',
               label="Bottle Appearance")
    ax.axvline(timestep_stim, color=color, linestyle='dashed', label="Stimulation/Distractor")
    try:
        plt.axvline(np.where(trial_data[:,10] != 0)[0][0], color='purple', label="Response") # response
    except:
        pass
    ax.plot(trial_data[:,6:8], color='grey', linestyle='dashed', label="Response density") # P_S
    ax.plot(trial_data[:,8:10], color='grey', linestyle='dotted', label="Response-side density") # P_R
    ax.set_xlabel("Timestep")
    ax.axvline(timesteps_reach, color='grey', alpha=0.5, label="Object contact")
    #ax.legend()
    return fig


def plot_movement(trial_data, timestep_bottle_appearance=None, timesteps_reach=None, figsize=(10,6)):
    """Plot only the movement: position/distance, orientation, bottle appearance and object contact."""
    if not timesteps_reach or timesteps_reach is None:
        timesteps_reach = int(trial_data[0,5])
    # possibly cut the data to the relevant timesteps
    trial_data = trial_data[:(int(timesteps_reach)+_timesteps_transport_default()),:]

    x = trial_data[:,12] # position/distance
    orient = trial_data[:,11] # orientation
    if timestep_bottle_appearance is None:
        timestep_bottle_appearance = int(trial_data[0,3])

    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(x, color="black", label="Position/Distance")
    ax.plot(orient, color="lightblue", linestyle="solid", label="Orientation")
    ax.axvline(timestep_bottle_appearance, 0, 1, color='green',
               linestyle='dashed' if trial_data[timestep_bottle_appearance,0] == utilities.UPRIGHT else 'dotted',
               label="Bottle Appearance")
    ax.axvline(timesteps_reach, color='grey', alpha=0.5, label="Object contact")
    ax.set_xlabel("Timestep")
    #ax.set_title("Movement: position & orientation")
    ax.legend()
    return fig


def plot_stim_response(trial_data, timestep_stim=None, timesteps_reach=None, figsize=(10,6)):
    """Plot stimulation and response (P_Q, P_R)."""
    if not timesteps_reach or timesteps_reach is None:
        timesteps_reach = int(trial_data[0,5])
    # cut to relevant timesteps
    trial_data = trial_data[:(int(timesteps_reach)+_timesteps_transport_default()),:]

    if timestep_stim is None:
        timestep_stim = int(trial_data[0,4])

    # determine stim color by finger and side
    try:
        if trial_data[timestep_stim,1] == utilities.THUMB:
            color = 'yellow' if trial_data[timestep_stim,2] == utilities.LEFT else 'orange'
        else:
            color = 'red' if trial_data[timestep_stim,2] == utilities.LEFT else 'purple'
    except Exception:
        color = 'orange'

    fig, ax = plt.subplots(figsize=figsize)
    ax.axvline(timestep_stim, color=color, linestyle='dashed', label="Stimulation/Distractor")
    try:
        resp_idx = np.where(trial_data[:,10] != 0)[0][0]
        ax.axvline(resp_idx, color='purple', label="Response")
    except Exception:
        resp_idx = None

    # response densities: columns 6,7 (P_Q) and 8,9 (P_R)
    try:
        #ax.plot(trial_data[:,6], color='grey', linestyle='dashed', label="Response-side density (P_S)")
        ax.plot(trial_data[:,6], color='grey', linestyle='dashed', label="Stimulus density (P_Q)")
        ax.plot(trial_data[:,7], color='grey', linestyle='dashed')
        ax.plot(trial_data[:,8], color='grey', linestyle='dotted', label="Response density (P_R)")
        ax.plot(trial_data[:,9], color='grey', linestyle='dotted')
    except Exception:
        pass

    ax.set_xlabel("Timestep")
    #ax.set_title("Stimulation & Response")
    ax.legend()
    return fig

def plot_trial_progression(trial_data, timestep_bottle_appearance=None, timestep_stim=None, timesteps_reach=None):
    """Plot movement, stimulation, response, and response densities for one trial."""
    if not timesteps_reach or timesteps_reach is None:
        timesteps_reach = int(trial_data[0,5])
    # possibly cut the data to the relevant timesteps
    trial_data = trial_data[:(int(timesteps_reach)+_timesteps_transport_default()),:]
    # other data points
    x = trial_data[:,12] # position/distance
    orient = trial_data[:,11] # orientation
    if not timestep_bottle_appearance or timestep_bottle_appearance is None:
        timestep_bottle_appearance = int(trial_data[0,3])
    if not timestep_stim or timestep_stim is None:
        timestep_stim = int(trial_data[0,4])
    if trial_data[timestep_stim,1] == utilities.THUMB:
        if trial_data[timestep_stim,2] == utilities.LEFT:
            color = 'yellow'
        else:
            color = 'orange'
    else:
        if trial_data[timestep_stim,2] == utilities.LEFT:
            color = 'red'
        else:
            color = 'purple'

    fig, ax = plt.subplots()
    ax.plot(x, color="black", label="Position/Distance")
    ax.plot(orient, color="lightblue", linestyle="solid", label="Orientation")
    ax.axvline(timestep_bottle_appearance, 0, 1, color='green',
               linestyle='dashed' if trial_data[timestep_bottle_appearance,0] == utilities.UPRIGHT else 'dotted',
               label="Bottle Appearance")
    ax.axvline(timestep_stim, color=color, linestyle='dashed', label="Stimulation/Distractor")
    try:
        plt.axvline(np.where(trial_data[:,10] != 0)[0][0], color='purple', label="Response") # response
    except:
        pass
    ax.plot(trial_data[:,6], color='grey', linestyle='dashed', label="Response density") # P_S
    ax.plot(trial_data[:,7], color='grey', linestyle='dashed') # P_S
    ax.plot(trial_data[:,8], color='grey', linestyle='dotted', label="Response-side density") # P_R
    ax.plot(trial_data[:,9], color='grey', linestyle='dotted') # P_R
    ax.set_xlabel("Timestep")
    ax.axvline(timesteps_reach, color='grey', alpha=0.5, label="Object contact")
    ax.legend()
    return fig

def plot_entropies(entropies, timesteps_reach=100):
    """Plot entropy traces by component for one trial."""
    # should be a 3D array of shape (timesteps, 5, 3)
    entropies = entropies[:(timesteps_reach+_timesteps_transport_default()),:,:] # cut to the right size
    # visualize
    fig, axs = plt.subplots(1, 3)#, figsize=(15, 5))
    for j in range(3):
        axs[j].plot(entropies[:,1,j], label="Reach")
        axs[j].plot(entropies[:,2,j], label="Grasp")
        axs[j].plot(entropies[:,3,j], label="Transport")
        axs[j].plot(entropies[:,4,j], label="Put-down")
        axs[j].plot(entropies[:,0,j], label="Idle")
        axs[j].set_xlabel("Timestep")
        axs[j].legend()
    axs[0].set_ylabel("Entropy value")
    axs[0].set_title("Initiator")
    axs[1].set_title("Dynamics")
    axs[2].set_title("Outcome")
    # tight layout
    plt.tight_layout()
    return fig

def rt_plot2(rt_data):
    """Plot RT condition means with error bars."""
    data = utilities.get_rt_df_(rt_data)
    fig, ax = plt.subplots()
    plt.errorbar(
        [0.2, 0.8],
        [np.mean(data[data["Condition"] == "Congruent Upright"]["RT"]),
         np.mean(data[data["Condition"] == "Incongruent Upright"]["RT"])],
        yerr=[np.std(data[data["Condition"] == "Congruent Upright"]["RT"]),
              np.std(data[data["Condition"] == "Incongruent Upright"]["RT"])],
        fmt='d-', label="Upright"
    )
    plt.errorbar(
        [0.203, 0.803],
        [np.mean(data[data["Condition"] == "Congruent Upside-Down"]["RT"]),
         np.mean(data[data["Condition"] == "Incongruent Upside-Down"]["RT"])],
        yerr=[np.std(data[data["Condition"] == "Congruent Upside-Down"]["RT"]),
              np.std(data[data["Condition"] == "Incongruent Upside-Down"]["RT"])],
        fmt='d-', label="Upside-Down"
    )
    plt.legend(title="Bottle Orientation")
    plt.ylabel("RT")
    plt.xticks([0.2, 0.8], ["Congruent", "Incongruent"])
    return fig




def single_event_over_trial_plot(densities, event_name="Event", bins=None, save=None, filetype="pdf",
                                 max_timesteps=None, row_height=0.22, dim_names=None,
                                 vmax_scale=1.05, rows_per_page=50, id=None):
    """Plot one event density over time as paged rows of small histograms."""
    densities = np.asarray(densities)
    if densities.ndim != 3:
        raise ValueError("densities must be a 3D array (timesteps x 3 x bins) or (timesteps x bins x 3)")

    # Normalize shape to (T, 3, B)
    if densities.shape[1] == 3:
        T_full, D, B = densities.shape
    elif densities.shape[2] == 3:
        T_full, B, D = densities.shape
        densities = densities.transpose(0, 2, 1)
        D = 3
    elif densities.shape[0] == 3:
        D, T_full, B = densities.shape
        densities = densities.transpose(1, 0, 2)
        D = 3
    else:
        raise ValueError("Expected one dimension to be 3 (the three columns per timestep)")

    if D != 3:
        raise ValueError("Plotting expects exactly 3 columns/dimensions per timestep")

    # Respect max_timesteps as a hard cap
    if max_timesteps is not None:
        T_effective = min(T_full, int(max_timesteps))
        densities = densities[:T_effective]
    else:
        T_effective = T_full

    # Determine whether pagination/truncation is needed
    do_paginate = (rows_per_page is not None) and (rows_per_page > 0) and (T_effective > rows_per_page)

    # The figure we'll return should be the first page (possibly truncated)
    T_return = min(T_effective, rows_per_page) if do_paginate else T_effective

    if bins is None:
        # B is the original bins value
        bins = B

    if dim_names is None:
        dim_names = ["Dim 0", "Dim 1", "Dim 2"]

    def make_page_figure(dens_page, page_index=0):
        # dens_page has shape (T_page, 3, bins)
        T_page = dens_page.shape[0]
        fig_w = 9.0
        fig_h = max(2.0, row_height * T_page)
        fig = plt.figure(figsize=(fig_w, fig_h))
        # reduce GridSpec hspace (we control final spacing with subplots_adjust)
        gs = fig.add_gridspec(T_page, 3, hspace=0.02, wspace=0.22)
        x = np.arange(bins)
        global_ymax = dens_page.max() * vmax_scale if dens_page.size else 1.0
        for t in range(T_page):
            for j in range(3):
                ax = fig.add_subplot(gs[t, j])
                vals = dens_page[t, j]
                ax.bar(x, vals, color=f"C{j}", width=0.9)
                ax.set_xlim(-0.5, bins - 0.5)
                ax.set_ylim(0, global_ymax)
                ax.set_xticks([])
                ax.set_yticks([])
                # slightly smaller column title font to avoid collisions with suptitle
                if t == 0:
                    ax.set_title(dim_names[j], fontsize=7)
                if j == 0:
                    ax.text(-0.02, 0.5, f"t={t + page_index*rows_per_page}", transform=ax.transAxes,
                            ha="right", va="center", fontsize=7)
        # Annotate if truncated
        if do_paginate:
            # Put the suptitle close to the top but leave extra room for column titles
            fig.suptitle(f"{event_name} densities over time (showing {T_page}/{T_effective} timesteps; page {page_index+1})", fontsize=10, y=0.99)
        else:
            fig.suptitle(f"{event_name} densities over time (top: t=0)", fontsize=10, y=0.99)
        # Move subplots down a bit (lower `top`) so column titles have space below the suptitle,
        # while keeping hspace tight between rows. These values prevent suptitle/column-title overlap.
        fig.subplots_adjust(top=0.97, left=0.03, right=0.98, hspace=0.02, wspace=0.18)
        return fig

    # Build the returned (first page) figure
    dens_first_page = densities[:T_return]
    fig_first = make_page_figure(dens_first_page, page_index=0)

    # If saving, handle multipage PDF (or numbered parts for other filetypes)
    if save:
        outdir = save
        if outdir and not os.path.exists(outdir):
            os.makedirs(outdir, exist_ok=True)
        safe_name = "".join(c if (c.isalnum() or c in (" ", "_", "-")) else "_" for c in event_name).strip().replace(" ", "_")
        ext = str(filetype).lstrip('.') if filetype is not None else 'pdf'
        if id is not None:
            safe_name += f"_{id}"
        filename = f"{safe_name}.{ext}"
        fullpath = os.path.join(outdir, filename)

        if ext.lower() == 'pdf' and do_paginate:
            # write a multi-page PDF with PdfPages; iterate pages over the capped range (T_effective)
            with PdfPages(fullpath) as pdf:
                for i, start in enumerate(range(0, T_effective, rows_per_page)):
                    end = min(start + rows_per_page, T_effective)
                    dens_page = densities[start:end]
                    fig_page = make_page_figure(dens_page, page_index=i)
                    pdf.savefig(fig_page, bbox_inches='tight')
                    plt.close(fig_page)
        else:
            # For non-PDF or when no pagination needed: if pagination was desired but filetype != pdf,
            # save numbered part files; otherwise save single file
            if do_paginate and ext.lower() != 'pdf':
                for i, start in enumerate(range(0, T_effective, rows_per_page)):
                    end = min(start + rows_per_page, T_effective)
                    dens_page = densities[start:end]
                    fig_page = make_page_figure(dens_page, page_index=i)
                    part_name = f"{safe_name}_part{i+1}.{ext}"
                    part_path = os.path.join(outdir, part_name)
                    fig_page.savefig(part_path, bbox_inches='tight')
                    plt.close(fig_page)
            else:
                # Save the single (possibly truncated) returned figure
                fig_first.savefig(fullpath, bbox_inches='tight')
    return fig_first

def events_over_trial_plot_all(events_densities, save=None, filetype="png", id=None, event_name="All_Events", font_size=None):
    """Plot all event/component density heatmaps in one figure."""

    #plt.rc("font", 6)
    #plt.rcParams.update({'font.size': 6})
    if font_size is not None:
        old_font_size = plt.rcParams["font.size"]
        plt.rcParams["font.size"] = font_size

    fig, axs = plt.subplots(5, 9, figsize=(15, 10), constrained_layout=True)
    event_names = ["Idle", "Reach", "Grasp", "Transport", "Put-down"]
    dim_names = ["Init-Orient", "Init-Pos", "Init-Angle", "Dyn-Orient", "Dyn-Pos", "Dyn-Angle",
                 "Out-Orient", "Out-Pos", "Out-Angle"]

    for i, event in enumerate(event_names):
        for j, dim_all in enumerate(dim_names):
            dimension = dim_all.split("-")[1]
            component = dim_all.split("-")[0]
            comp = {"Init":0, "Dyn":1, "Out":2}[component]
            dim = {"Orient":0, "Pos":1, "Angle":2}[dimension]
            im = axs[i, j].imshow(events_densities[i, comp, dim, :, :], origin="lower", aspect="auto",
                                  extent=[0.0, 1.0, 0, events_densities.shape[-2]],
                                  cmap="viridis")
            if i == 0:
                axs[i, j].set_title(dim_all)
            if j == 0:
                axs[i, j].set_ylabel(event)
            #fig.colorbar(im, ax=axs[i, j])
            #fig.colorbar(im, cax=axs[i, j])

    if save:
        outdir = save
        if outdir and not os.path.exists(outdir):
            os.makedirs(outdir, exist_ok=True)
        safe_name = "".join(c if (c.isalnum() or c in (" ", "_", "-")) else "_" for c in event_name).strip().replace(" ", "_")
        ext = str(filetype).lstrip('.') if filetype is not None else 'pdf'
        if id is not None:
            safe_name += f"_{id}"
        filename = f"{safe_name}.{ext}"
        fullpath = os.path.join(outdir, filename)
        fig.savefig(fullpath, bbox_inches='tight', dpi=300)

    #return fig
    if font_size is not None:
        plt.rcParams["font.size"] = old_font_size

def plot_single_density(density, save=None, filetype="png", id=None, density_name="Density",
                        type="heatmap", rows_per_col=100, max_timesteps=None, show_title=True, font_size=None):
    """Plot one density over time as a heatmap or paged row plot."""
    # Set font size if necessary
    if font_size is not None:
        old_font_size = plt.rcParams["font.size"]
        plt.rcParams["font.size"] = font_size

    data = np.asarray(density)
    # cut timesteps if needed
    if max_timesteps is not None:
        data = data[:int(max_timesteps)]
    # If it's not a heatmap but instead a set of histograms/barplots, as above
    if type=="histograms":
        # Make several columns depending on the length
        n_cols = int(np.ceil(data.shape[0] / rows_per_col))
        fig, axs = plt.subplots(rows_per_col, n_cols, figsize=(3*n_cols, 0.22*rows_per_col), constrained_layout=True)
        for col in range(n_cols):
            for row in range(rows_per_col):
                t = col * rows_per_col + row
                if t >= data.shape[0]:
                    break
                ax = axs[row, col] if n_cols > 1 else axs[row]
                ax.bar(np.arange(data.shape[1]), data[t], color="C0", width=0.9)
                ax.set_xlim(-0.5, data.shape[1]-0.5)
                ax.set_ylim(0, data.max()*1.05)
                ax.set_xticks([])
                ax.set_yticks([])
                if row == 0:
                    ax.set_title(f"t={t}")
        if density_name is not None and show_title is not False:
            fig.suptitle(f"{density_name} over time", fontsize=10, y=0.99)
        fig.subplots_adjust(top=0.97, left=0.03, right=0.98, hspace=0.02, wspace=0.18)
    else:
        # heatmap
        fig, ax = plt.subplots(figsize=(9, 6))
        im = ax.imshow(data, origin="lower", aspect="auto",
                       extent=[0.0, 1.0, 0, data.shape[0]],
                       cmap="viridis")
        fig.colorbar(im, ax=ax)
        ax.set_xlabel("Density Support")
        ax.set_ylabel("Timestep")
        if density_name is not None and show_title is not False:
            ax.set_title(f"{density_name} over time")
    if save:
        outdir = save
        if outdir and not os.path.exists(outdir):
            os.makedirs(outdir, exist_ok=True)
        safe_name = "".join(c if (c.isalnum() or c in (" ", "_", "-")) else "_" for c in density_name).strip().replace(" ", "_")
        ext = str(filetype).lstrip('.') if filetype is not None else 'pdf'
        if id is not None:
            safe_name += f"_{id}"
        filename = f"{safe_name}.{ext}"
        fullpath = os.path.join(outdir, filename)
        fig.savefig(fullpath, bbox_inches='tight', dpi=300)

    if font_size is not None:
        plt.rcParams["font.size"] = old_font_size
    return fig

def plot_focus_fours(trials_data):
    """Plot focus values for congruent/incongruent and upright/upside-down trials."""
    # getting parts here depending on the indices
    congr_up_data = utilities.get_means_further_trials(trials_data[congr_upright_indices])
    incongr_up_data = utilities.get_means_further_trials(trials_data[incongr_upright_indices])
    congr_down_data = utilities.get_means_further_trials(trials_data[congr_updown_indices])
    incongr_down_data = utilities.get_means_further_trials(trials_data[incongr_updown_indices])

    # see if there are different timesteps_reach values or not, decide whether to share x row-wise or total
    t_reachs = trials_data[:,0,5]
    sharex = True if np.all(t_reachs == t_reachs[0]) else 'row'

    fig, axs = plt.subplots(2, 2, figsize=(12, 10), sharex=sharex, sharey=True)

    # helper function to plot the focus values
    def plot_focus_subplot(ax, data, title):
        timesteps_reach = int(data[0, 5])
        F = data[:(timesteps_reach + _timesteps_transport_default()), 19:24]
        ax.plot(F[:, 1], label="F[1] (reach)", alpha=0.7)
        ax.plot(F[:, 2], label="F[2] (grasp)", alpha=0.7)
        ax.plot(F[:, 3], label="F[3] (transp.)", alpha=0.7)
        ax.plot(F[:, 4], label="F[4] (put-d.)", alpha=0.7)
        ax.plot(F[:, 0], label="F[0] (idle)", alpha=0.7)
        ax.axvline(timesteps_reach, color='grey', alpha=0.5)
        ax.set_xlim(0, timesteps_reach + _timesteps_transport_default())
        ax.set_title(title)
        ax.set_xlabel("Timestep")
        ax.set_ylabel("Focus value")

    # Single subplots for each condition
    plot_focus_subplot(axs[0, 0], congr_up_data, "Congruent Upright")
    axs[0, 0].legend()

    plot_focus_subplot(axs[0, 1], incongr_up_data, "Incongruent Upright")

    plot_focus_subplot(axs[1, 0], congr_down_data, "Congruent Upside-Down")

    plot_focus_subplot(axs[1, 1], incongr_down_data, "Incongruent Upside-Down")

    # Layout
    plt.tight_layout()
    return fig

def plot_focus_subplot(ax, data, title):
    """Draw one focus subplot on an existing axis."""
    timesteps_reach = int(data[0, 5])
    F = data[:(timesteps_reach + _timesteps_transport_default()), 19:24]
    ax.plot(F[:, 1], label="F[1] (reach)", alpha=0.7)
    ax.plot(F[:, 2], label="F[2] (grasp)", alpha=0.7)
    ax.plot(F[:, 3], label="F[3] (transp.)", alpha=0.7)
    ax.plot(F[:, 4], label="F[4] (put-d.)", alpha=0.7)
    ax.plot(F[:, 0], label="F[0] (idle)", alpha=0.7)
    ax.axvline(timesteps_reach, color='grey', alpha=0.5)
    ax.set_xlim(0, timesteps_reach + _timesteps_transport_default())
    ax.set_title(title)
    ax.set_xlabel("Timestep")
    ax.set_ylabel("Focus value")
    #return ax

def plot_focus_twos(trials_data, font_size=None):
    """Plot focus averaged by upright and upside-down orientation."""
    if font_size is not None:
        old_font_size = plt.rcParams.get('font.size', None)
        plt.rcParams.update({'font.size': font_size})
    try:
        # getting parts here depending on the indices
        upright_indices = congr_upright_indices + incongr_upright_indices
        updown_indices = congr_updown_indices + incongr_updown_indices
        upright_data = utilities.get_means_further_trials(trials_data[upright_indices])
        updown_data = utilities.get_means_further_trials(trials_data[updown_indices])

        # # see if there are different timesteps_reach values or not, decide whether to share x row-wise or total
        # t_reachs = trials_data[:,0,5]
        # sharex = True if np.all(t_reachs == t_reachs[0]) else 'row'
        # (for now, make them horizontally, not vertically)

        fig, axs = plt.subplots(1, 2, figsize=(12, 5), sharey=True)

        # Single subplots for each condition
        plot_focus_subplot(axs[0], upright_data, "Upright")
        axs[0].legend()
        plot_focus_subplot(axs[1], updown_data, "Upside-Down")

        # Layout
        plt.tight_layout()
        return fig
    finally:
        # restore old font size if necessary
        if font_size is not None and old_font_size is not None and old_font_size != font_size:
            plt.rcParams.update({'font.size': old_font_size})

def plot_inf_event_fours(trials_data):
    """Plot inferred event traces for the four trial-condition groups."""
    # getting parts here depending on the indices
    congr_up_data = utilities.get_means_further_trials(trials_data[congr_upright_indices])
    incongr_up_data = utilities.get_means_further_trials(trials_data[incongr_upright_indices])
    congr_down_data = utilities.get_means_further_trials(trials_data[congr_updown_indices])
    incongr_down_data = utilities.get_means_further_trials(trials_data[incongr_updown_indices])

    # see if there are different timesteps_reach values or not, decide whether to share x row-wise or total
    t_reachs = trials_data[:,0,5]
    sharex = True if np.all(t_reachs == t_reachs[0]) else 'row'

    fig, axs = plt.subplots(2, 2, figsize=(12, 10), sharex=sharex, sharey=True)

    # helper function to plot the inferred event values
    def plot_inf_event_subplot(ax, data, title):
        timesteps_reach = int(data[0, 5])
        ax.plot(data[:(timesteps_reach + _timesteps_transport_default()), 18], label="inferred event", color="black")
        ax.axhline(0, color="lightblue", linestyle="dotted")
        ax.axhline(1, color="blue", linestyle="dotted")
        ax.axhline(2, color="orange", linestyle="dotted")
        ax.axhline(3, color="green", linestyle="dotted")
        ax.axhline(4, color="purple", linestyle="dotted")
        ax.set_xlabel("Timestep")
        ax.set_yticks([0, 1, 2, 3, 4])
        ax.set_yticklabels(["idle", "reach", "grasp", "transport", "put-down"], rotation=45)
        ax.axvline(timesteps_reach, color='grey', alpha=0.5)
        ax.set_title(title)

    # Single subplots for each condition
    plot_inf_event_subplot(axs[0, 0], congr_up_data, "Congruent Upright")

    plot_inf_event_subplot(axs[0, 1], incongr_up_data, "Incongruent Upright")

    plot_inf_event_subplot(axs[1, 0], congr_down_data, "Congruent Upside-Down")

    plot_inf_event_subplot(axs[1, 1], incongr_down_data, "Incongruent Upside-Down")

    # Layout
    plt.tight_layout()
    return fig


def plot_multiple_conditions_focus(datasets,which_params, which_params_values, mod_name): # todo: for now, saving back where this is called? ...
    """Plot focus comparisons for a two-parameter condition grid."""

    # Normalize inputs
    if which_params is None:
        which_params = []
    if which_params_values is None:
        which_params_values = []

    # Determine grid dimensions
    if len(which_params) >= 2:
        param_row_name = which_params[0]
        param_col_name = which_params[1]
        row_vals = list(which_params_values[0])
        col_vals = list(which_params_values[1])
        rows = len(row_vals)
        cols = len(col_vals)
    elif len(which_params) == 1:
        param_row_name = None
        param_col_name = which_params[0]
        row_vals = [None]
        col_vals = list(which_params_values[0])
        rows = 1
        cols = len(col_vals)
    else:
        # fallback: place all datasets in a single row
        param_row_name = None
        param_col_name = None
        row_vals = [None]
        cols = len(datasets)
        col_vals = [None] * cols
        rows = 1

    total_cells = rows * cols

    # Ensure datasets is a list and has correct length
    datasets_list = list(datasets)
    if len(datasets_list) < total_cells:
        datasets_list = datasets_list + [None] * (total_cells - len(datasets_list))
    else:
        datasets_list = datasets_list[:total_cells]

    # Create figure and axes grid
    fig, axs = plt.subplots(rows, cols, figsize=(max(15, cols*5), max(9, rows*3)), squeeze=False)

    # For each cell, create two small axes and plot upright/upside-down using plot_focus_subplot
    upright_indices = congr_upright_indices + incongr_upright_indices
    updown_indices = congr_updown_indices + incongr_updown_indices

    for r in range(rows):
        for c in range(cols):
            ax = axs[r][c]
            idx = r * cols + c
            cell_data = datasets_list[idx]

            # create two inset axes side-by-side inside the cell axis
            left_ax = ax.inset_axes([0.0, 0.0, 0.48, 1.0])
            right_ax = ax.inset_axes([0.52, 0.0, 0.48, 1.0])
            # hide the parent cell axis frame so only the inset axes are visible
            ax.set_frame_on(False)
            # ensure the parent axis has no ticks or background patch, but remains present
            try:
                ax.patch.set_alpha(0)
            except Exception:
                pass
            ax.set_xticks([])
            ax.set_yticks([])

            if cell_data is None:
                # no data: annotate blank cell
                left_ax.text(0.5, 0.5, "no data", ha='center', va='center', color='red')
                right_ax.axis('off')
                left_ax.set_xticks([])
                left_ax.set_yticks([])
            else:
                try:
                    # Expect cell_data to be an array-like with shape (N_trials, timesteps, features)
                    td = np.asarray(cell_data)
                    # select upright/upside-down trials using the standard indices
                    u_idxs = [i for i in upright_indices if i < td.shape[0]]
                    d_idxs = [i for i in updown_indices if i < td.shape[0]]

                    if len(u_idxs) > 0:
                        upright_data = utilities.get_means_further_trials(td[u_idxs])
                    else:
                        upright_data = None
                    if len(d_idxs) > 0:
                        updown_data = utilities.get_means_further_trials(td[d_idxs])
                    else:
                        updown_data = None

                except Exception as e:
                    upright_data = None
                    updown_data = None

                # plot into the small axes using existing helper
                # We'll plot manually into the inset axes but avoid titles/axis labels
                def _plot_focus_on_ax(ax_inner, data_inner):
                    if data_inner is None:
                        ax_inner.text(0.5, 0.5, "no data", ha='center', va='center', color='gray')
                        ax_inner.set_xticks([])
                        return
                    timesteps_reach = int(data_inner[0, 5])
                    F = data_inner[:(timesteps_reach + _timesteps_transport_default()), 19:24]
                    x = np.arange(F.shape[0])
                    # color mapping: 1->blue, 2->orange, 3->green, 4->purple, 0->lightblue
                    colmap = {1: 'tab:blue', 2: 'tab:orange', 3: 'tab:green', 4: 'tab:purple', 0: 'lightblue'}
                    ax_inner.plot(x, F[:, 1], color=colmap[1], alpha=0.8)
                    ax_inner.plot(x, F[:, 2], color=colmap[2], alpha=0.8)
                    ax_inner.plot(x, F[:, 3], color=colmap[3], alpha=0.8)
                    ax_inner.plot(x, F[:, 4], color=colmap[4], alpha=0.8)
                    ax_inner.plot(x, F[:, 0], color=colmap[0], alpha=0.8)
                    ax_inner.axvline(timesteps_reach, color='grey', alpha=0.5)
                    ax_inner.set_xlim(0, timesteps_reach + _timesteps_transport_default())
                    # do not set titles or y-labels here (outer labels are used)
                    #  try and set them to none here ...
                    #ax_inner.set_yticks([])

                _plot_focus_on_ax(left_ax, upright_data)
                _plot_focus_on_ax(right_ax, updown_data)
                # only show a legend once (top-left cell)
                if r == 0 and c == 0:
                    # create legend in the left inset using labels
                    # build proxy lines for legend
                    lab_lines = []
                    lab_labels = ["F[1] (reach)", "F[2] (grasp)", "F[3] (transp.)", "F[4] (put-d.)", "F[0] (idle)"]
                    # match the same colors as the plots
                    colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:purple', 'lightblue']
                    for color, lab in zip(colors, lab_labels):
                        line, = left_ax.plot([], [], color=color, label=lab)
                        lab_lines.append(line)
                    left_ax.legend(handles=lab_lines, loc='upper right', fontsize='small')

                # By default disable ticks and labels on both inset axes to avoid duplicates
                left_ax.tick_params(axis='both', which='both', labelbottom=False, bottom=False)
                right_ax.tick_params(axis='both', which='both', labelbottom=False, bottom=False)

                # Enable x ticks/label only for the bottom row on the right inset
                if r == rows - 1:
                    right_ax.tick_params(axis='x', which='both', labelbottom=True, bottom=True)
                    right_ax.set_xlabel('Timestep')

                # Enable y ticks/label only for the leftmost column on the left inset
                if c == 0:
                    left_ax.tick_params(axis='y', which='both', labelleft=True, left=True)
                    left_ax.set_ylabel('Focus value')

    # Layout
    plt.tight_layout()
    return fig


def plot_multiple_conditions_focus_subfigures(datasets, which_params, which_params_values, mod_name=None, figsize=None):
    """Plot focus comparisons for a two-parameter grid using subfigures."""
    # Normalize inputs
    if which_params is None:
        which_params = []
    if which_params_values is None:
        which_params_values = []

    # determine grid shape
    if len(which_params) >= 2:
        param_row_name = which_params[0]
        param_col_name = which_params[1]
        row_vals = list(which_params_values[0])
        col_vals = list(which_params_values[1])
        rows = len(row_vals)
        cols = len(col_vals)
    elif len(which_params) == 1:
        param_row_name = None
        param_col_name = which_params[0]
        row_vals = [None]
        col_vals = list(which_params_values[0])
        rows = 1
        cols = len(col_vals)
    else:
        param_row_name = None
        param_col_name = None
        row_vals = [None]
        cols = len(datasets)
        col_vals = [None] * cols
        rows = 1

    total_cells = rows * cols
    datasets_list = list(datasets)
    if len(datasets_list) < total_cells:
        datasets_list = datasets_list + [None] * (total_cells - len(datasets_list))
    else:
        datasets_list = datasets_list[:total_cells]

    # figure size heuristic
    if figsize is None:
        fig_w = int(max(12, cols * 5.5))
        fig_h = int(max(4, rows * 3.5))
    else:
        fig_w, fig_h = figsize

    # create a gridspec with a top row and left column reserved for labels
    # layout: (rows+1) x (cols+1) grid; [0,0] unused, [0,1:] column labels, [1:,0] row labels, [1:,1:] plots
    # reduce the size reserved for labels (top/left) and tighten spacing
    fig = plt.figure(figsize=(fig_w, fig_h))
    outer_gs = fig.add_gridspec(rows + 1, cols + 1,
                                width_ratios=[0.08] + [1.0]*cols,
                                height_ratios=[0.02] + [1.0]*rows,
                                hspace=0.08, wspace=0.10)

    upright_indices = congr_upright_indices + incongr_upright_indices
    updown_indices = congr_updown_indices + incongr_updown_indices
    colmap = {1: 'tab:blue', 2: 'tab:orange', 3: 'tab:green', 4: 'tab:red', 0: 'lightblue'}

    legend_drawn = False

    # Top: parameter name (center) and per-column values
    if param_col_name is not None:
        ax_colname = fig.add_subplot(outer_gs[0, 1:])
        ax_colname.axis('off')
        # place the parameter name a bit lower so it's closer to the column values and plots
        ax_colname.text(0.5, 0.95, str(param_col_name), ha='center', va='bottom', fontsize=12, weight='bold')
        for c, col_val in enumerate(col_vals):
            # Make a small axis for each column value for reliable placement
            ax_c = fig.add_subplot(outer_gs[0, c + 1])
            ax_c.axis('off')
            # raise the value text so it sits closer to the plots
            ax_c.text(0.5, 0.5, str(col_val), ha='center', va='center', fontsize=10)
    else:
        # consume top row but leave it empty
        for c in range(cols):
            ax_c = fig.add_subplot(outer_gs[0, c + 1])
            ax_c.axis('off')

    # Left: parameter name (vertical) and per-row values
    if param_row_name is not None:
        ax_rowname = fig.add_subplot(outer_gs[1:, 0])
        ax_rowname.axis('off')
        # place the row parameter name slightly to the right so it visually aligns with row values
        ax_rowname.text(0.5, 0.5, str(param_row_name), ha='center', va='center', rotation='vertical', fontsize=12, weight='bold')
        for r, row_val in enumerate(row_vals):
            ax_r = fig.add_subplot(outer_gs[r + 1, 0])
            ax_r.axis('off')
            # shift row values a bit toward the plots for tighter grouping
            ax_r.text(0.3, 0.7, str(row_val), ha='center', va='center', rotation='vertical', fontsize=10)
    else:
        for r in range(rows):
            ax_r = fig.add_subplot(outer_gs[r + 1, 0])
            ax_r.axis('off')

    # Inner grid: the plots
    for r in range(rows):
        for c in range(cols):
            cell_gs = outer_gs[r + 1, c + 1].subgridspec(1, 2, wspace=0.05)
            left_ax = fig.add_subplot(cell_gs[0, 0])
            right_ax = fig.add_subplot(cell_gs[0, 1], sharey=left_ax)

            idx = r * cols + c
            cell_data = datasets_list[idx]

            if cell_data is None:
                left_ax.text(0.5, 0.5, 'no data', ha='center', va='center', color='red')
                right_ax.axis('off')
            else:
                td = np.asarray(cell_data)
                u_idxs = [i for i in upright_indices if i < td.shape[0]]
                d_idxs = [i for i in updown_indices if i < td.shape[0]]
                upright_data = utilities.get_means_further_trials(td[u_idxs]) if len(u_idxs) > 0 else None
                updown_data = utilities.get_means_further_trials(td[d_idxs]) if len(d_idxs) > 0 else None

                def _plot_focus(ax_plot, data_plot):
                    if data_plot is None:
                        ax_plot.text(0.5, 0.5, 'no data', ha='center', va='center', color='gray')
                        ax_plot.set_xticks([])
                        return
                    treach = int(data_plot[0, 5])
                    F = data_plot[:(treach + _timesteps_transport_default()), 19:24]
                    x = np.arange(F.shape[0])
                    ax_plot.plot(x, F[:, 1], color=colmap[1], alpha=0.9)
                    ax_plot.plot(x, F[:, 2], color=colmap[2], alpha=0.9)
                    ax_plot.plot(x, F[:, 3], color=colmap[3], alpha=0.9)
                    ax_plot.plot(x, F[:, 4], color=colmap[4], alpha=0.9)
                    ax_plot.plot(x, F[:, 0], color=colmap[0], alpha=0.9)
                    ax_plot.axvline(treach, color='grey', alpha=0.5)
                    ax_plot.set_xlim(0, treach + _timesteps_transport_default())

                _plot_focus(left_ax, upright_data)
                _plot_focus(right_ax, updown_data)

                if (not legend_drawn) and (upright_data is not None or updown_data is not None):
                    proxy_lines = []
                    labels = ['F[1] (reach)', 'F[2] (grasp)', 'F[3] (transp.)', 'F[4] (put-d.)', 'F[0] (idle)']
                    colors = [colmap[1], colmap[2], colmap[3], colmap[4], colmap[0]]
                    for color, lab in zip(colors, labels):
                        line, = left_ax.plot([], [], color=color, label=lab)
                        proxy_lines.append(line)
                    left_ax.legend(handles=proxy_lines, loc='upper right', fontsize='small')
                    legend_drawn = True

            # hide inner spines and ticks on shared boundary
            try:
                left_ax.spines['right'].set_visible(False)
            except Exception:
                pass
            try:
                right_ax.spines['left'].set_visible(False)
            except Exception:
                pass
            try:
                left_ax.yaxis.set_ticks_position('left')
                left_ax.tick_params(axis='y', which='both', right=False)
            except Exception:
                pass
            try:
                right_ax.yaxis.set_ticks_position('right')
                right_ax.tick_params(axis='y', which='both', left=False)
            except Exception:
                pass

            # tick control: only leftmost column shows y ticks; only bottom row shows x ticks
            if c == 0:
                left_ax.set_ylabel('Focus value')
            else:
                left_ax.set_yticklabels([])
                left_ax.set_yticks([])
                right_ax.set_yticklabels([])
                right_ax.set_yticks([])

            if r == rows - 1:
                right_ax.set_xlabel('Timestep')
            else:
                left_ax.set_xticklabels([])
                left_ax.set_xticks([])
                right_ax.set_xticklabels([])
                right_ax.set_xticks([])

    # small final layout adjustment
    fig.subplots_adjust(top=0.96, left=0.03, right=0.98, bottom=0.03)
    return fig


def plot_condition_values_focus(datasets, which_param, which_param_values, mod_name=None, figsize=None):
    """Plot focus comparisons for multiple values of one parameter."""
    # Normalize inputs
    if which_param is None:
        which_param = "Condition"
    if which_param_values is None:
        which_param_values = [f"Value {i}" for i in range(len(datasets))]

    rows = max(1, len(which_param_values))
    cols = 2

    # Figure size heuristic
    if figsize is None:
        fig_w = 10
        fig_h = max(3.0, rows * 2.8)
    else:
        fig_w, fig_h = figsize

    fig = plt.figure(figsize=(fig_w, fig_h))

    upright_indices = congr_upright_indices + incongr_upright_indices
    updown_indices = congr_updown_indices + incongr_updown_indices

    use_subfigs = hasattr(fig, 'subfigures')

    axes_rows = []

    if use_subfigs:
        subfigs = fig.subfigures(nrows=rows, ncols=1, squeeze=False)
        for r in range(rows):
            subfig = subfigs[r, 0]
            axs = subfig.subplots(1, 2, squeeze=False, sharey=True)
            # tighten internal spacing
            try:
                subfig.subplots_adjust(left=0.03, right=0.99, wspace=0.025)
            except Exception:
                pass
            left_ax = axs[0][0]
            right_ax = axs[0][1]
            axes_rows.append((left_ax, right_ax, subfig))
    else:
        axs_all = fig.subplots(rows, 2, squeeze=False, sharey=True, gridspec_kw={'hspace': 0.18, 'wspace': 0.025})
        for r in range(rows):
            left_ax = axs_all[r][0]
            right_ax = axs_all[r][1]
            axes_rows.append((left_ax, right_ax, None))

    # plotting helper
    def _plot_focus(ax_plot, data_plot):
        if data_plot is None:
            ax_plot.text(0.5, 0.5, 'no data', ha='center', va='center', color='gray')
            ax_plot.set_xticks([])
            return
        treach = int(data_plot[0, 5])
        F = data_plot[:(treach + _timesteps_transport_default()), 19:24]
        x = np.arange(F.shape[0])
        colmap = {1: 'tab:blue', 2: 'tab:orange', 3: 'tab:green', 4: 'tab:red', 0: 'lightblue'}
        ax_plot.plot(x, F[:, 1], color=colmap[1], alpha=0.9)
        ax_plot.plot(x, F[:, 2], color=colmap[2], alpha=0.9)
        ax_plot.plot(x, F[:, 3], color=colmap[3], alpha=0.9)
        ax_plot.plot(x, F[:, 4], color=colmap[4], alpha=0.9)
        ax_plot.plot(x, F[:, 0], color=colmap[0], alpha=0.9)
        ax_plot.axvline(treach, color='grey', alpha=0.5)
        ax_plot.set_xlim(0, treach + _timesteps_transport_default())

    # Fill rows
    datasets_list = list(datasets)
    if len(datasets_list) < rows:
        datasets_list += [None] * (rows - len(datasets_list))

    legend_drawn = False
    for r in range(rows):
        cell_data = datasets_list[r]
        left_ax, right_ax, subfig = axes_rows[r]

        if cell_data is None:
            left_ax.text(0.5, 0.5, 'no data', ha='center', va='center', color='red')
            right_ax.text(0.5, 0.5, 'no data', ha='center', va='center', color='red')
            left_ax.set_xticks([])
            right_ax.set_xticks([])
            left_ax.set_yticks([])
            right_ax.set_yticks([])
            # ensure x-axis spine visible
            try:
                left_ax.spines['bottom'].set_visible(True)
                right_ax.spines['bottom'].set_visible(True)
            except Exception:
                pass
            continue

        td = np.asarray(cell_data)
        u_idxs = [i for i in upright_indices if i < td.shape[0]]
        d_idxs = [i for i in updown_indices if i < td.shape[0]]
        upright_data = utilities.get_means_further_trials(td[u_idxs]) if len(u_idxs) > 0 else None
        updown_data = utilities.get_means_further_trials(td[d_idxs]) if len(d_idxs) > 0 else None

        _plot_focus(left_ax, upright_data)
        _plot_focus(right_ax, updown_data)

        # show y ticks only on left column
        if r == 0 and not legend_drawn and (upright_data is not None or updown_data is not None):
            # draw legend once in the top-left
            proxy_lines = []
            labels = ['F[1] (reach)', 'F[2] (grasp)', 'F[3] (transp.)', 'F[4] (put-d.)', 'F[0] (idle)']
            colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'lightblue']
            for color, lab in zip(colors, labels):
                line, = left_ax.plot([], [], color=color, label=lab)
                proxy_lines.append(line)
            left_ax.legend(handles=proxy_lines, loc='upper right', fontsize='small')
            legend_drawn = True

        if r == rows - 1:
            # per-axis xlabel on bottom row
            right_ax.set_xlabel('Timestep')

        # Ensure xticks/ticklabels visible on all rows
        try:
            left_ax.tick_params(axis='x', which='both', labelbottom=True)
            right_ax.tick_params(axis='x', which='both', labelbottom=True)
        except Exception:
            pass

        # y label only for left column
        if True:
            left_ax.set_ylabel('Focus value')
        # hide y ticks for right column
        right_ax.set_yticklabels([])
        right_ax.set_yticks([])

        # make bottom spines visible for clarity
        try:
            left_ax.spines['bottom'].set_visible(True)
            right_ax.spines['bottom'].set_visible(True)
        except Exception:
            pass

    # Force a draw so bboxes are available
    fig.canvas.draw()

    # Compute required bottom margin from xtick/xlabel bboxes
    try:
        renderer = fig.canvas.get_renderer()
        max_y1 = 0.0
        for left_ax, right_ax, _ in axes_rows:
            # consider xtick labels for both axes
            for lab in left_ax.get_xticklabels() + right_ax.get_xticklabels():
                try:
                    bb = lab.get_window_extent(renderer)
                    bbf = bb.transformed(fig.transFigure.inverted())
                    max_y1 = max(max_y1, bbf.y1)
                except Exception:
                    pass
            # also consider axis xlabels
            try:
                xl = left_ax.xaxis.get_label()
                bb = xl.get_window_extent(renderer)
                bbf = bb.transformed(fig.transFigure.inverted())
                max_y1 = max(max_y1, bbf.y1)
            except Exception:
                pass
            try:
                xr = right_ax.xaxis.get_label()
                bb = xr.get_window_extent(renderer)
                bbf = bb.transformed(fig.transFigure.inverted())
                max_y1 = max(max_y1, bbf.y1)
            except Exception:
                pass

        # conservative padding
        required_bottom = min(0.18, max(0.06, max_y1 + 0.03))
        # apply margin
        fig.subplots_adjust(bottom=required_bottom)

        # create a tiny bottom axis for the global xlabel
        try:
            label_ax = fig.add_axes([0.03, 0.0, 0.94, required_bottom])
            label_ax.axis('off')
            label_ax.text(0.5, 0.5, 'Timestep', ha='center', va='center', fontsize=10)
        except Exception:
            # fallback to fig.text
            fig.text(0.5, required_bottom * 0.45, 'Timestep', ha='center', va='center', fontsize=10)
    except Exception:
        # if something goes wrong, ensure a decent bottom
        try:
            fig.subplots_adjust(bottom=0.12)
            fig.text(0.5, 0.06, 'Timestep', ha='center', va='center', fontsize=10)
        except Exception:
            pass

    # Column headers: centered over left/right columns using the first row's axes positions
    try:
        left0, right0, _ = axes_rows[0]
        left_pos = left0.get_position()
        right_pos = right0.get_position()
        center_left = 0.5 * (left_pos.x0 + left_pos.x1)
        center_right = 0.5 * (right_pos.x0 + right_pos.x1)
        col_y = min(0.98, 0.96)
        fig.text(center_left, col_y, 'Upright', ha='center', va='bottom', fontsize=11, weight='bold')
        fig.text(center_right, col_y, 'Upside-down', ha='center', va='bottom', fontsize=11, weight='bold')
    except Exception:
        pass

    # Row labels (centered titles)
    try:
        for r in range(rows):
            left_ax, right_ax, subfig = axes_rows[r]
            label = f"{which_param} = {which_param_values[r]}"
            # Prefer subfigure suptitle when available so Matplotlib manages spacing.
            if subfig is not None:
                try:
                    # place slightly below the top of the subfigure
                    subfig.suptitle(label, fontsize=10, weight='bold', y=0.98)
                    continue
                except Exception:
                    pass
            # Fallback: place centered above the row using figure coordinates
            try:
                lp = left_ax.get_position()
                rp = right_ax.get_position()
                center_x = 0.5 * (lp.x0 + rp.x1)
                row_top = max(lp.y1, rp.y1)
                y = min(row_top + 0.02, 0.92)
                fig.text(center_x, y, label, ha='center', va='bottom', fontsize=10, weight='bold')
            except Exception:
                # last-resort: place near top center
                try:
                    fig.text(0.5, 0.92 - (r * 0.02), label, ha='center', va='bottom', fontsize=10, weight='bold')
                except Exception:
                    pass
    except Exception:
        pass

    return fig


def plot_event_mappings(mappings=None, cmap='Blues', figsize=(10, 4), save_folder=None, n_columns=3):
    """Plot event mapping matrices, optionally saving the figure."""
    # default try to load common names from module globals (works when placed in same file)
    if mappings is None:
        candidates = ["M_none","M_even","M_default","M_reach","M_upright","M_upsidedown",
                      "M_grasp","M_idle","M_transport","M_putdown"]
        mappings = {}
        g = globals()
        for name in candidates:
            if name in g:
                mappings[name] = np.asarray(g[name], dtype=float)
    if not mappings:
        raise ValueError("No mappings found.")

    n = len(mappings)
    cols = min(n_columns, n)
    rows = math.ceil(n / cols)
    fig, axs = plt.subplots(rows, cols, figsize=figsize, squeeze=False)
    tick_x = ["Left", "Right"]     # Distractor axis
    tick_y = ["Thumb", "Index"]    # Stimulation axis

    for idx, (name, mat) in enumerate(mappings.items()):
        r = idx // cols
        c = idx % cols
        ax = axs[r][c]
        mat = np.asarray(mat)
        if mat.shape != (2,2):
            # try to reshape or warn
            try:
                mat = mat.reshape((2,2))
            except:
                raise ValueError(f"Mapping {name} has shape {mat.shape}, expected (2,2).")
        im = ax.imshow(mat, vmin=0, vmax=max(1.0, mat.max()), cmap=cmap, aspect='auto')
        # annotate values
        for i in range(2):
            for j in range(2):
                ax.text(j, i, f"{mat[i,j]:.2f}", ha='center', va='center', color='black', fontsize=10)
        ax.set_yticks([0,1])
        ax.set_xticks([0,1])
        # show x tick labels only on the bottom row, and y tick labels only on the leftmost column
        if r == rows - 1:
            ax.set_xticklabels(tick_x)
            ax.set_xlabel("Distractor (Side)")
        else:
            ax.set_xticklabels(["", ""])  # hide x tick labels for non-bottom rows

        if c == 0:
            ax.set_yticklabels(tick_y)
            ax.set_ylabel("Stimulation (Finger)")
        else:
            ax.set_yticklabels(["", ""])  # hide y tick labels for non-left columns

        ax.set_title(name)
        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    # hide any unused subplots
    total_plots = rows * cols
    for idx2 in range(n, total_plots):
        r = idx2 // cols
        c = idx2 % cols
        axs[r][c].axis('off')

    fig.tight_layout()
    if save_folder is not None:
        # create output dir if needed
        os.makedirs(save_folder, exist_ok=True)
        # base filename and extension
        base_name = "event_mappings"
        ext = "png"
        # if the plain name is available, use it; otherwise append increasing suffixes
        filename = f"{base_name}.{ext}"
        fullpath = os.path.join(save_folder, filename)
        if os.path.exists(fullpath):
            i = 1
            while True:
                filename = f"{base_name}_{i}.{ext}"
                fullpath = os.path.join(save_folder, filename)
                if not os.path.exists(fullpath):
                    break
                i += 1
        fig.savefig(fullpath, dpi=300, bbox_inches='tight')
    plt.show()


def plot_focus_one(trials_data, t_calc=None, font_size=None):
    """Plot average focus values across all supplied trials."""
    old_font_size = None
    if font_size is not None:
        old_font_size = plt.rcParams.get('font.size', None)
        plt.rcParams.update({'font.size': font_size})

    try:
        # Mean over all passed trial groups (same helper used by split focus plots).
        data = utilities.get_means_further_trials(trials_data)
        timesteps_reach = int(data[0, 5])
        max_t = min(int(timesteps_reach + _timesteps_transport_default()), data.shape[0])
        F = data[:max_t, 19:24]

        fig, ax = plt.subplots(figsize=(8, 5))
        plot_order = [1, 2, 3, 4, 0]
        labels = ["F[1] (reach)", "F[2] (grasp)", "F[3] (transp.)", "F[4] (put-d.)", "F[0] (idle)"]

        lines = []
        for idx, label in zip(plot_order, labels):
            line, = ax.plot(F[:, idx], label=label, alpha=0.8)
            lines.append((idx, line))

        ax.axvline(timesteps_reach, color='grey', alpha=0.6)
        ax.set_xlim(0, max_t)
        ax.set_xlabel("Timestep")
        ax.set_ylabel("Focus value")
        ax.set_title("Focus (overall mean)")
        ax.legend()

        if t_calc is not None:
            t_points = sorted({int(t) for t in t_calc if 0 <= int(t) < max_t})
            if t_points:
                # Mark requested timesteps and add an in-plot value summary for each one.
                summary_lines = []
                for t in t_points:
                    ax.axvline(t, color='lightgrey', linestyle=':', linewidth=0.8, alpha=0.8)
                    for idx, line in lines:
                        ax.scatter(t, F[t, idx], color=line.get_color(), s=16, zorder=4)
                    summary_lines.append(
                        f"t={t}: i={F[t,0]:.2f}, r={F[t,1]:.2f}, g={F[t,2]:.2f}, tr={F[t,3]:.2f}, pd={F[t,4]:.2f}"
                    )

                ax.text(
                    0.01,
                    0.99,
                    "\n".join(summary_lines),
                    transform=ax.transAxes,
                    va='top',
                    ha='left',
                    fontsize=max(7, plt.rcParams.get('font.size', 10) - 1),
                    bbox=dict(facecolor='white', alpha=0.75, edgecolor='lightgrey'),
                )

        plt.tight_layout()
        return fig
    finally:
        if font_size is not None and old_font_size is not None and old_font_size != font_size:
            plt.rcParams.update({'font.size': old_font_size})