//  [python]
def gen_figure(data, demographics, agreement, consensus, assignments, zrand):
    """
    Generates figure 3

    Parameters
    ----------
    data : list of pandas.DataFrame
    demographics : pandas.DataFrame
    agreement : (N, N) array_like
    consensus : (N,) array_like
    zrand : (C, K, M) array_like

    Returns
    -------
    fig : matplotlib.figure.Figure
        Plotted figure
    """

    # make figure
    fig = plt.figure(figsize=(16.5, 15))
    gs = gridspec.GridSpec(3, 15, figure=fig)
    gs.update(wspace=1.2, hspace=0.5)

    ax1, ax2, ax3 = (plt.subplot(gs[0, n:(n + 5)]) for n in [0, 5, 10])

    ###########################################################################
    # make hyperparameter similarity plot

    coll = ax1.imshow(zrand[1], cmap=defaults.similarity_cmap, alpha=0.7,
                      aspect='auto', vmin=50, vmax=150, rasterized=True)

    # make axis a bit prettier
    ax1.set(xticks=[], yticks=[], xticklabels=[], yticklabels=[])
    ax1.set_title('parameter stability', pad=20, fontsize=20)
    ax1.set_xlabel('scaling, μ', labelpad=7)
    ax1.set_ylabel('neighbors, K', labelpad=7)
    sns.despine(ax=ax1, left=True, bottom=True)

    # add colorbar to axis 1
    ax1_cbar = fig.colorbar(coll, ax=ax1, drawedges=False, ticks=[])
    ax1_cbar.outline.set(linewidth=0)
    ax1_cbar.ax.set_ylabel('local similarity', rotation=270, labelpad=25)
    ax1_cbar.ax.tick_params(axis='both', which='both', length=0)

    ###########################################################################
    # make cluster assignment map

    cbar_kws = {
        'ticks': [],
        'boundaries': np.arange(-0.5, 4.5)
    }
    # re-order community assignments for plotting purposes
    comms_plot, idxs = cluster.reorder_assignments(assignments, consensus,
                                                   seed=SEED)
    ax2 = sns.heatmap(comms_plot, cmap=defaults.four_cluster_cmap,
                      xticklabels=[], yticklabels=[],
                      ax=ax2, cbar_kws=cbar_kws, rasterized=True)
    hlines = np.where(np.diff(consensus[consensus.argsort()]))[0] + 1
    ax2.hlines(hlines, 0, comms_plot.shape[-1], color='white', linewidth=2)
    ax2.set_title('patient clusters', pad=20, fontsize=20)
    ax2.set_xlabel('cluster partitions', labelpad=7)
    ax2.set_ylabel('patients', labelpad=7)

    # modify colorbar
    cbar = ax2.collections[0].colorbar
    cbar.outline.set(linewidth=0)
    cbar.ax.set_ylabel('cluster label', rotation=270, labelpad=25)
    cbar.ax.tick_params(axis='both', which='both', length=0)

    ###########################################################################
    # make sorted cluster heatmap

    ax3 = plotting.plot_mod_heatmap(agreement, consensus, ax=ax3,
                                    inds=idxs[0].squeeze(), rasterized=True,
                                    cmap='viridis', edgecolor='white')
    ax3.set(xticks=[], yticks=[], xticklabels=[], yticklabels=[])
    ax3.set_title('patient co-assignment', pad=20, fontsize=20)
    ax3.set_xlabel('patients', labelpad=7)
    ax3.set_ylabel('patients', labelpad=7)

    # modify colorbar
    cbar = ax3.collections[0].colorbar
    cbar.set_ticks([])
    cbar.outline.set(linewidth=0)
    cbar.ax.set_ylabel('probability', rotation=270, labelpad=25)
    cbar.ax.tick_params(axis='both', which='both', length=0)

    # add figure labels
    for ax, text in zip([ax1, ax2, ax3], ['a', 'b', 'c']):
        ax.text(-0.2, 1.1, text, transform=ax.transAxes, fontdict=fontd)

    ###########################################################################
    # make rainplots

    axes = [plt.subplot(gs[1, r:(r + 3)]) for r in range(0, 15, 3)]
    titles = [
        'cortical thickness',
        'subcortical volume',
        'dopamine binding',
        'protein availability',
        'clinical score'
    ]
    xlabels = {
        'supramarginal_15_r': 'supramarginal gyrus',
        'substantia_nigra_pars_compacta': 'substantia nigra',
        'caudate_l': 'left caudate',
        'ttau': 'total tau',
        'tremor': 'tremor'
    }
    for dat, ax, title in zip(data, axes, titles):
        # z-score data for plotting
        zdat = sstats.zscore(dat, ddof=1)
        zdf = pd.DataFrame(zdat, index=dat.index, columns=dat.columns)
        currdata = pd.merge(demographics[['cluster']], zdf, on='participant')

        # find maximally discriminating feature in data
        grps = (zdat[consensus == cl] for cl in np.unique(consensus))
        idx = sstats.f_oneway(*grps).statistic.argmax()
        xlabel = xlabels.get(zdf.columns[idx])

        # make rainplot
        utils.rainplot(x=dat.columns[idx], y='cluster', data=currdata,
                       viol_kws={'linewidth': 0}, ax=ax,
                       palette=defaults.three_cluster_cmap)

        # make axis goodness
        ax.set(xticklabels=[-2.5, 2.5], yticklabels=[], ylabel='',
               xlabel=title, xlim=(-3.5, 3.5), xticks=[-2.5, 2.5])
        ax.set_title(xlabel, pad=15)
        sns.despine(ax=ax, left=True)
        utils.shift_axis(ax, lshift=0, rshift=0.02)

    axes[0].text(-0.3, 1.10, 'd', transform=axes[0].transAxes, fontdict=fontd)

    ###########################################################################
    # make lineplot

    behavior = load_longitudinal_behavior(demographics.index)
    axes_lp = [plt.subplot(gs[2, r:(r + 4)]) for r in [4, 9]]
    for n, (lp, ax) in enumerate(zip(['pigd', 'tremor'], axes_lp)):
        # get data for given test
        pd_test = pd.merge(clean_visits(behavior, lp), demographics,
                           on='participant')
        # normalize data based on first visit mean/stdev
        t = pd_test.query('visit == 0')[lp]
        pd_test.loc[:, lp] -= t.mean()
        pd_test.loc[:, lp] /= t.std(ddof=1)
        # plot it
        sns.lineplot(x='visit', y=lp, data=pd_test.dropna(subset=[lp]),
                     hue='cluster', palette=defaults.three_cluster_cmap,
                     legend=False, ci=68, ax=ax)
        ax.set(xticklabels=[], xlabel='time')
        if n == 0:
            ax.set_ylabel('clinical severity', labelpad=10)
            ax.text(-0.3, 1.10, 'e', transform=ax.transAxes, fontdict=fontd)
            ax.set(ylim=[-0.5, 2.5], yticks=[-0.5, 0.5, 1.5, 2.5],
                   yticklabels=[-0.5, 0.5, 1.5, 2.5])
        else:
            ax.set(ylim=[-0.75, 0.75], ylabel='',
                   yticks=[-0.75, 0, 0.75],
                   yticklabels=[-0.75, 0, 0.75])
        ax.set_title(lp, pad=15)
        sns.despine(ax=ax)

    ###########################################################################
    # shift axes to make room for everything

    utils.shift_axis(ax1, lshift=0.04)
    utils.shift_axis(ax1_cbar.ax, lshift=0.04)
    utils.shift_axis(ax2, lshift=0.06)
    utils.shift_axis(ax2.collections[0].colorbar.ax, lshift=0.06)
    utils.shift_axis(ax3, lshift=0.08)
    utils.shift_axis(ax3.collections[0].colorbar.ax, lshift=0.08)

    for a, lshift in zip(axes[1:], np.cumsum([0.02] * 4)):
        utils.shift_axis(a, lshift=lshift)

    utils.shift_axis(axes_lp[0], lshift=-0.05, rshift=0.01, tshift=-0.018)
    utils.shift_axis(axes_lp[1], rshift=0.05, lshift=-0.01, tshift=-0.018)

    return fig