<reponame>colleenjg/OpenScope_CA_Analysis
"""
modif_analysis_plots.py

This script is used to modify plot designs on the fly for ROI and running 
analyses from dictionaries, e.g. for presentations or papers.

Authors: <NAME>

Date: October, 2018

Note: this code uses python 3.7.

"""

import copy
import logging
import os
import warnings

from joblib import Parallel, delayed
from matplotlib import pyplot as plt
import numpy as np

from util import file_util, gen_util, logger_util, plot_util
from sess_util import sess_plot_util, sess_str_util
from plot_fcts import roi_analysis_plots as roi_plots

logger = logging.getLogger(__name__)


TAB = "    "


#############################################
def plot_from_dict(dict_path, plt_bkend=None, fontdir=None, plot_tc=True, 
                   parallel=False, datetime=True):
    """
    plot_from_dict(dict_path)

    Plots data from dictionaries containing analysis parameters and results.

    Required args:
        - dict_path (str): path to dictionary to plot data from
    
    Optional_args:
        - plt_bkend (str): mpl backend to use for plotting (e.g., "agg")
                           default: None
        - fontdir (str)  : path to directory where additional fonts are stored
                           default: None
        - plot_tc (bool) : if True, tuning curves are plotted for each ROI 
                           default: True
        - parallel (bool): if True, some of the analysis is parallelized across 
                           CPU cores
                           default: False
        - datetime (bool): figpar["save"] datatime parameter (whether to 
                           place figures in a datetime folder)
                           default: True
    """

    logger.info(f"Plotting from dictionary: {dict_path}", 
        extra={"spacing": "\n"})
    
    figpar = sess_plot_util.init_figpar(
        plt_bkend=plt_bkend, fontdir=fontdir, datetime=datetime)
    plot_util.manage_mpl(cmap=False, **figpar["mng"])

    plt.rcParams["figure.titlesize"] = "xx-large"
    plt.rcParams["axes.titlesize"] = "xx-large"

    info = file_util.loadfile(dict_path)
    savedir = os.path.dirname(dict_path)

    analysis = info["extrapar"]["analysis"]

    # 0. Plots the full traces for each session
    if analysis == "f": # full traces
        plot_full_traces(figpar=figpar, savedir=savedir, **info)

    # 1. Plot average traces by quintile x surprise for each session 
    elif analysis == "t": # traces
        plot_traces_by_qu_surp_sess(figpar=figpar, savedir=savedir, **info)

    # 2. Plot average traces by quintile, locked to surprise for each session 
    elif analysis == "l": # surprise locked traces
        plot_traces_by_qu_lock_sess(figpar=figpar, savedir=savedir, **info)

    # 4. Plot autocorrelations
    elif analysis == "a": # autocorr
        plot_autocorr(figpar=figpar, savedir=savedir, **info)

    # 6. Plot colormaps and traces for orientations/directions
    elif analysis == "o": # colormaps
        plot_oridirs(figpar=figpar, savedir=savedir, parallel=parallel, **info)

    else:
        warnings.warn(f"No modified plotting function for analysis {analysis}")

    plt.close("all")

#############################################
def plot_full_traces(analyspar, sesspar, extrapar, sess_info, trace_info, 
                     figpar=None, savedir=None):
    """
    plot_full_traces(analyspar, sesspar, extrapar, sess_info, trace_info)

    From dictionaries, plots full traces for each session in a separate subplot.
    
    Returns figure name and save directory path.
    
    Required args:
        - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)    : dictionary with keys of SessPar namedtuple
        - extrapar (dict)   : dictionary containing additional analysis 
                              parameters
            ["analysis"] (str): analysis type (e.g., "t")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - sess_info (dict)  : dictionary containing information from each
                              session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - trace_info (dict): dictionary containing trace information
            ["all_tr"] (nested list): trace values structured as
                                          sess x 
                                          ([each_roi (scaled), averaged] 
                                            if datatype is "roi" x)
                                          (ROI if each_roi x)
                                          (me/err if averaged x)
                                          frames
            ["all_edges"] (list)    : edge values for each parameter, 
                                      structured as sess x block x 
                                                    edges ([start, end])
            ["all_pars"] (list)     : stimulus parameter strings structured as 
                                                    sess x block
                
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None    
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """
 
    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")
    
    sessstr = "sess{}_{}".format(sesspar["sess_n"], sesspar["plane"])
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])
    
    datatype = extrapar["datatype"]

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"))

    n_rows = 1
    if datatype == "roi":
        n_rows = 2
    
    if figpar is None:
        figpar = sess_plot_util.init_figpar()
    figpar = copy.deepcopy(figpar)
    figpar["init"]["subplot_wid"] *= 4
    figpar["init"]["subplot_hei"] *= 4
    figpar["init"]["sharex"] = True
    figpar["init"]["sharey"] = False
    
    fig, ax = plot_util.init_fig(n_sess*n_rows, n_sess, **figpar["init"])
    for i in range(n_sess):
        title = (f"Mouse {mouse_ns[i]} (sess {sess_ns[i]}, {lines[i]} "
            f"{planes[i]}{dendstr_pr}{nroi_strs[i]})")

        sub_axs = ax[:, i]
        sub_axs[0].set_title(title)
        if datatype == "roi":
            # average trace
            av_tr = np.asarray(trace_info["all_tr"][i][1])
            xran  = range(roi_tr_sep.shape[1])
            subtitle = "u{} across ROIs".format(statstr_pr)
            plot_util.plot_traces(
                sub_axs[1], xran, av_tr[0], av_tr[1:], title=subtitle, 
                xticks="auto")
            
            # each ROI trace
            roi_tr = np.asarray(trace_info["all_tr"][i][0])
            # values to add to each ROI to split them apart
            add = np.linspace(
                0, roi_tr.shape[0] * 1.5 + 1, roi_tr.shape[0])[:, np.newaxis]
            roi_tr_sep = roi_tr + add
            sub_axs[0].plot(roi_tr_sep.T)
        else:
            run_tr = np.asarray(trace_info["all_tr"][i])
            sub_axs[0].plot(run_tr)

        for b, block in enumerate(trace_info["all_edges"]):
            # all block labels to the lower plot
            plot_util.add_labels(
                sub_axs[-1], trace_info["all_pars"][b], np.mean(block), 0.85, 
                color="k")
            # add lines to both plots
            for r in range(n_rows):
                plot_util.add_bars(sub_axs[r], hbars=block)
                sess_plot_util.add_axislabels(
                    sub_axs[r], fluor=analyspar["fluor"], datatype=datatype, 
                    x_ax="")
                
    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype])

    savename = f"{datatype}_tr_{sessstr}{dendstr}"
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_traces_by_qu_surp_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats, figpar=None, 
                                savedir=None):
    """
    plot_traces_by_qu_surp_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats)

    From dictionaries, plots traces by quintile/surprise with each session in a 
    separate subplot.
    
    Returns figure name and save directory path.
    
    Required args:
        - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)    : dictionary with keys of SessPar namedtuple
        - stimpar (dict)    : dictionary with keys of StimPar namedtuple
        - extrapar (dict)   : dictionary containing additional analysis 
                              parameters
            ["analysis"] (str): analysis type (e.g., "t")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - quintpar (dict)   : dictionary with keys of QuintPar namedtuple
        - sess_info (dict)  : dictionary containing information from each
                              session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - trace_stats (dict): dictionary containing trace stats information
            ["xrans"] (list)           : time values for the 2p frames, for 
                                         each session
            ["all_stats"] (list)       : list of 4D arrays or lists of trace 
                                         data statistics across ROIs, 
                                         structured as:
                                            surp x quintiles x
                                            stats (me, err) x frames
            ["all_counts"] (array-like): number of sequences, structured as:
                                                sess x surp x quintiles
                
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None    
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """
 
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"],
        stimpar["gabk"], "print")
    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")
    
    sessstr = sess_str_util.sess_par_str(
        sesspar["sess_n"], stimpar["stimtype"], sesspar["plane"], 
        stimpar["bri_dir"], stimpar["bri_size"], stimpar["gabk"])
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])
    
    datatype = extrapar["datatype"]
    dimstr = sess_str_util.datatype_dim_str(datatype)

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"), style="par")

    xrans      = [np.asarray(xran) for xran in trace_stats["xrans"]]
    all_stats  = [np.asarray(sessst) for sessst in trace_stats["all_stats"]]
    all_counts = trace_stats["all_counts"]

    cols, lab_cols = sess_plot_util.get_quint_cols(quintpar["n_quints"])
    alpha = np.min([0.4, 0.8/quintpar["n_quints"]])

    surps = ["reg", "surp"]
    n = 6
    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    fig, ax = plot_util.init_fig(n_sess, **figpar["init"])
    for i in range(n_sess):
        sub_ax = plot_util.get_subax(ax, i)
        for s, [col, leg_ext] in enumerate(zip(cols, surps)):
            for q, qu_lab in enumerate(quintpar["qu_lab"]):
                if qu_lab != "":
                    qu_lab = f"{qu_lab.capitalize()} "
                line, plane = "5", "dendrites"
                if "23" in lines[i]:
                    line = "2/3"
                if "soma" in planes[i]:
                    plane = "somata"
                title = (f"M{mouse_ns[i]} - layer {line} {plane}{dendstr_pr}")
                # title=(f"Mouse {mouse_ns[i]} - {stimstr_pr} " + 
                #     u"{} ".format(statstr_pr) + f"across {dimstr}\n"
                #     f"(sess {sess_ns[i]}, {lines[i]} {planes[i]},"
                #     f"{nroi_strs[i]}")
                leg = None
                y_ax = ""
                if i == 0:
                    leg = f"{qu_lab}{leg_ext}"
                    y_ax = None
                # leg = f"{qu_lab}{leg_ext} ({all_counts[i][s][q]})"
                plot_util.plot_traces(
                    sub_ax, xrans[i], all_stats[i][s, q, 0], 
                    all_stats[i][s, q, 1:], title, color=col[q], alpha=alpha, 
                    label=leg, n_xticks=n, xticks="auto")
                sess_plot_util.add_axislabels(
                    sub_ax, fluor=analyspar["fluor"], datatype=datatype, 
                    y_ax=y_ax)

    plot_util.turn_off_extra(ax, n_sess)

    if stimpar["stimtype"] == "gabors": 
        sess_plot_util.plot_labels(
            ax, stimpar["gabfr"], "both", pre=stimpar["pre"], 
            post=stimpar["post"], cols=lab_cols, 
            sharey=figpar["init"]["sharey"])
    
    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["surp_qu"])

    qu_str = "_{}q".format(quintpar["n_quints"])
    if quintpar["n_quints"] == 1:
        qu_str = ""

    savename = f"{datatype}_av_{sessstr}{dendstr}{qu_str}"
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_traces_by_qu_lock_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats, 
                                figpar=None, savedir=None):
    """
    plot_traces_by_qu_lock_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats)

    From dictionaries, plots traces by quintile, locked to transitions from 
    surprise to regular or v.v. with each session in a separate subplot.
    
    Returns figure name and save directory path.
    
    Required args:
        - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)    : dictionary with keys of SessPar namedtuple
        - stimpar (dict)    : dictionary with keys of StimPar namedtuple
        - extrapar (dict)   : dictionary containing additional analysis 
                              parameters
            ["analysis"] (str): analysis type (e.g., "l")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - quintpar (dict)   : dictionary with keys of QuintPar namedtuple
        - sess_info (dict)  : dictionary containing information from each
                              session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - trace_stats (dict): dictionary containing trace stats information
            ["xrans"] (list)           : time values for the 2p frames, for 
                                         each session
            ["all_stats"] (list)       : list of 4D arrays or lists of trace 
                                         data statistics across ROIs, 
                                         structured as:
                                            (surp_len x) quintiles x
                                            stats (me, err) x frames
            ["all_counts"] (array-like): number of sequences, structured as:
                                                sess x (surp_len x) quintiles
            ["lock"] (str)             : value to which segments are locked:
                                         "surp", "reg" or "surp_split"
            ["baseline"] (num)         : number of seconds used for baseline
            ["reg_stats"] (list)       : list of 3D arrays or lists of trace 
                                         data statistics across ROIs for
                                         regular sampled sequences, 
                                         structured as:
                                            quintiles (1) x stats (me, err) 
                                            x frames
            ["reg_counts"] (array-like): number of sequences corresponding to
                                         reg_stats, structured as:
                                            sess x quintiles (1)
            
            if data is by surp_len:
            ["surp_lens"] (list)       : number of consecutive segments for
                                         each surp_len, structured by session
                
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None    
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """

    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"],
        stimpar["gabk"], "print")
    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")

    sessstr = sess_str_util.sess_par_str(
        sesspar["sess_n"], stimpar["stimtype"], sesspar["plane"], 
        stimpar["bri_dir"], stimpar["bri_size"], stimpar["gabk"])
    basestr = sess_str_util.base_par_str(trace_stats["baseline"])
    basestr_pr = sess_str_util.base_par_str(trace_stats["baseline"], "print")

    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])
    
    datatype = extrapar["datatype"]
    dimstr = sess_str_util.datatype_dim_str(datatype)

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"), style="par")

    xrans      = [np.asarray(xran) for xran in trace_stats["xrans"]]
    all_stats  = [np.asarray(sessst) for sessst in trace_stats["all_stats"]]
    reg_stats  = [np.asarray(regst) for regst in trace_stats["reg_stats"]]
    all_counts = trace_stats["all_counts"]
    reg_counts = trace_stats["reg_counts"]

    lock  = trace_stats["lock"]
    col_idx = 0
    if "surp" in lock:
        lock = "surp"
        col_idx = 1
    
    surp_lab, len_ext = "", ""
    surp_lens = [[None]] * n_sess
    offset = 0
    surp_len_default = True # plot default surprise lengths
    if (stimpar["stimtype"] == "gabors" and 
        stimpar["gabfr"] not in ["any", "all"]):
        offset = stimpar["gabfr"]
    if "surp_lens" in trace_stats.keys():
        surp_lens = trace_stats["surp_lens"]
        len_ext = "_bylen"
        if stimpar["stimtype"] == "gabors":
            offset = stimpar["gabfr"]
            surp_lens = [[sl * 1.5/4 - 0.3 * offset for sl in sls] 
                for sls in surp_lens]
        surp_len_default = False

    # plot surp_lens default values and RANGE TO PLOT
    if stimpar["stimtype"] == "gabors":
        DEFAULT_SURP_LENS = [3.0, 4.5, 6.0]
        end_val = 8.0
    else:
        DEFAULT_SURP_LENS = [2.0, 3.0, 4.0]
        end_val = 6.0
    if lock == "surp":
        st_val = -2.0
        inv = 1
    else:
        st_val = -end_val + 2
        end_val = end_val - 2
        inv = -1
    n_ticks = int((end_val - st_val)//2 + 1)
    
    if figpar is None:
        figpar = sess_plot_util.init_figpar()
    figpar = copy.deepcopy(figpar)
    figpar["init"]["subplot_wid"] = 6.5

    fig, ax = plot_util.init_fig(n_sess, **figpar["init"])
    for i, (stats, counts) in enumerate(zip(all_stats, all_counts)):
        xran = xrans[i]
        st, end = 0, len(xran)
        st_vals = list(filter(lambda i: xran[i] <= st_val, range(len(xran))))
        if len(st_vals) != 0:
            st = st_vals[-1]
        end_vals = list(filter(lambda i: xran[i] >= end_val, range(len(xran))))
        if len(end_vals) != 0:
            end = end_vals[0] + 1

        sub_ax = plot_util.get_subax(ax, i)

        line, plane = "5", "dendrites"
        if "23" in lines[i]:
            line = "2/3"
        if "soma" in planes[i]:
            plane = "somata"

        title = f"M{mouse_ns[i]} - layer {line} {plane}{dendstr_pr}"
        # title=(f"Mouse {mouse_ns[i]} - {stimstr_pr} " + 
        #     u"{statstr_pr} " + f"{lock} locked across {dimstr}{basestr_pr}"
        #     f"\n(sess {sess_ns[i]}, {lines[i]} {planes[i]}, {nroi_strs[i]}")
        y_ax = ""
        if i == 0:
            y_ax = None

        sess_plot_util.add_axislabels(
            sub_ax, fluor=analyspar["fluor"], datatype=datatype, y_ax=y_ax)
        plot_util.add_bars(sub_ax, hbars=0)
        n_lines = quintpar["n_quints"] * len(surp_lens[i])
        try: 
            cols = sess_plot_util.get_quint_cols(n_lines)[0][col_idx]
        except:
            cols = [None] * n_lines
        alpha = np.min([0.4, 0.8/n_lines])
        if stimpar["stimtype"] == "gabors":
            sess_plot_util.plot_gabfr_pattern(sub_ax, xran[st:end], 
                offset=offset, bars_omit=[0] + surp_lens[i])
        # plot regular data
        if reg_stats[i].shape[0] != 1:
            raise ValueError("Expected only one quintile for reg_stats.")
        
        leg = None
        if i == 0:
            leg = "reg"
        if surp_len_default:
            for leng in DEFAULT_SURP_LENS:
                edge = leng * inv
                if edge < 0:
                    edge = np.max([xran[st], edge])
                elif edge > 0:
                    edge = np.min([xran[end - 1], edge])
                plot_util.add_vshade(
                    sub_ax, 0, edge, color=cols[-1], alpha=0.1)
        # leg = f"reg (no lock) ({reg_counts[i][0]})"
        plot_util.plot_traces(
            sub_ax, xran[st:end], reg_stats[i][0][0, st:end], 
            reg_stats[i][0][1:, st:end], alpha=alpha, label=leg, 
            alpha_line=0.8, color="darkgray", xticks="auto")
        n = 0 # count lines plotted
        for s, surp_len in enumerate(surp_lens[i]):
            if surp_len is not None:
                edge = surp_len * inv
                if edge < 0:
                    edge = np.max([xran[st], edge])
                elif edge > 0:
                    edge = np.min([xran[end - 1], edge])
                plot_util.add_vshade(
                    sub_ax, 0, edge, color=cols[n], alpha=0.1)
                counts, stats = all_counts[i][s], all_stats[i][s]                
                surp_lab = f"surp len {surp_len}"
            else:
                # surp_lab = "surp lock"
                surp_lab = "surp"
            for q, qu_lab in enumerate(quintpar["qu_lab"]):
                if qu_lab != "":
                    qu_lab = f"{qu_lab.capitalize()} "
                lab = f"{qu_lab}{surp_lab}"
                if n == 2 and cols[n] is None:
                    sub_ax.plot([], []) # to advance the color cycle (past gray)
                #leg = f"{lab} ({counts[q]})"
                leg = None
                if i == 0:
                    leg = lab
                plot_util.plot_traces(
                    sub_ax, xran[st:end], stats[q][0, st:end], 
                    stats[q][1:, st:end], title, alpha=alpha, label=leg, 
                    n_xticks=n_ticks, alpha_line=0.8, color=cols[n], 
                    xticks="auto")
                n += 1
            if surp_len is not None:
                plot_util.add_bars(sub_ax, hbars=surp_len, 
                                   color=sub_ax.lines[-1].get_color(), alpha=1)
    
    plot_util.turn_off_extra(ax, n_sess)

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["surp_qu"], 
            f"{lock}_lock", basestr.replace("_", ""))

    qu_str = "_{}q".format(quintpar["n_quints"])
    if quintpar["n_quints"] == 1:
        qu_str = ""
 
    savename = (f"{datatype}_av_{lock}lock{len_ext}{basestr}_{sessstr}"
        f"{dendstr}{qu_str}")
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_autocorr(analyspar, sesspar, stimpar, extrapar, autocorrpar, 
                  sess_info, autocorr_data, figpar=None, savedir=None):
    """
    plot_autocorr(analyspar, sesspar, stimpar, extrapar, autocorrpar, 
                  sess_info, autocorr_data)

    From dictionaries, plots autocorrelation during stimulus blocks.

    Required args:
        - analyspar (dict)    : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)      : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)      : dictionary with keys of StimPar namedtuple
        - extrapar (dict)     : dictionary containing additional analysis 
                                parameters
            ["analysis"] (str): analysis type (e.g., "a")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - autocorrpar (dict)  : dictionary with keys of AutocorrPar namedtuple
        - sess_info (dict)    : dictionary containing information from each
                                session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - autocorr_data (dict): dictionary containing data to plot:
            ["xrans"] (list): list of lag values in seconds for each session
            ["stats"] (list): list of 3D arrays (or nested lists) of
                              autocorrelation statistics, structured as:
                                     sessions stats (me, err) 
                                     x ROI or 1x and 10x lag 
                                     x lag
    
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """


    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")

    sessstr = sess_str_util.sess_par_str(
        sesspar["sess_n"], stimpar["stimtype"], sesspar["plane"], 
        stimpar["bri_dir"],stimpar["bri_size"], stimpar["gabk"]) 
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])


    datatype = extrapar["datatype"]
    if datatype == "roi":
        fluorstr_pr = sess_str_util.fluor_par_str(
            analyspar["fluor"], str_type="print")
        title_str = u"{} autocorrelation".format(fluorstr_pr)
        if not autocorrpar["byitem"]:
            title_str = f"{title_str} across ROIs" 
    elif datatype == "run":
        datastr = sess_str_util.datatype_par_str(datatype)
        title_str = f"{datastr} autocorrelation"

    if stimpar["stimtype"] == "gabors":
        seq_bars = [-1.5, 1.5] # light lines
    else:
        seq_bars = [-1.0, 1.0] # light lines

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"), style="par")

    xrans = autocorr_data["xrans"]
    stats = [np.asarray(stat) for stat in autocorr_data["stats"]]

    lag_s = autocorrpar["lag_s"]
    xticks = np.linspace(-lag_s, lag_s, lag_s*2+1)
    yticks = np.linspace(0, 1, 6)

    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    byitemstr = ""
    if autocorrpar["byitem"]:
        byitemstr = "_byroi"

    fig, ax = plot_util.init_fig(n_sess, **figpar["init"])
    for i in range(n_sess):
        sub_ax = plot_util.get_subax(ax, i)
        line, plane = "5", "dendrites"
        if "23" in lines[i]:
            line = "2/3"
        if "soma" in planes[i]:
            plane = "somata"
        title=f"M{mouse_ns[i]} - layer {line} {plane}{dendstr_pr}"
        # title = (f"Mouse {mouse_ns[i]} - {stimstr_pr}" + 
        #     u"{}".format(statstr_pr) + 
        #     f"{title_str}\n(sess {sess_ns[i]}, {lines[i]} {planes[i]},"
        #     f"{nroi_strs[i]})")
        # transpose to ROI/lag x stats x series
        sess_stats = stats[i].transpose(1, 0, 2) 
        for s, sub_stats in enumerate(sess_stats):
            lab = None
            if i == 0:
                if not autocorrpar["byitem"]:
                    lab = ["actual lag", "10x lag"][s]
            plot_util.plot_traces(
                sub_ax, xrans[i], sub_stats[0], sub_stats[1:], xticks=xticks, 
                yticks=yticks, alpha=0.2, label=lab)
        plot_util.add_bars(sub_ax, hbars=seq_bars)
        sub_ax.set_ylim([0, 1])
        sub_ax.set_title(title)
        sub_ax.set_xlabel("Lag (s)")

    plot_util.turn_off_extra(ax, n_sess)

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["autocorr"])

    savename = f"{datatype}_autocorr{byitemstr}_{sessstr}{dendstr}"

    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_oridir_traces(analyspar, sesspar, stimpar, extrapar, quintpar, 
                        tr_data, sess_info, figpar=None, savedir=None):
    """
    plot_oridir_traces(analyspar, sesspar, stimpar, extrapar, quintpar, 
                       tr_data, sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI as colormaps for a single session and optionally
    a single quintile. 

    Required args:
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)  : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - extrapar (dict) : dictionary containing additional analysis 
                            parameters
            ["analysis"] (str): analysis type (e.g., "o")
            ["datatype"] (str): datatype (e.g., "roi")
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (only first session used)
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict): dictionary containing number of seqs for each
                               surprise x ori/dir combination under a 
                               separate key
            ["stats"] (dict) : dictionary containing trace mean/medians across
                               ROIs in 2D arrays or nested lists, 
                               structured as:
                                   stats (me, err) x frames
                               with each surprise x ori/dir combination under a 
                               separate key
                               (NaN arrays for combinations with 0 seqs.)
            ["xran"] (list)  : time values for the 2p frames

    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """

    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")

    stimstr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"])
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])

    datatype = extrapar["datatype"]
    if datatype != "roi":
        raise ValueError("Function only implemented for roi datatype.")
    dimstr = sess_str_util.datatype_dim_str(datatype)

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["oridir"])

    # extract some info from dictionaries
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_n, sess_n, line, plane] = [sess_info[key][0] for key in keys]

    nroi_str = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], style="par")[0]

    xran = tr_data["xran"]

    surps = ["reg", "surp"]
    if stimpar["stimtype"] == "gabors":
        surp_labs = surps
        deg = u"\u00B0"
        oridirs = stimpar["gab_ori"]
        n = 6
    elif stimpar["stimtype"] == "bricks":
        surp_labs = [f"{surps[i]} -> {surps[1-i]}" for i in range(len(surps))]
        deg  = ""
        oridirs = stimpar["bri_dir"]
        n = 7

    qu_str, qu_str_pr = quintpar["qu_lab"][0], quintpar["qu_lab_pr"][0]
    if qu_str != "":
        qu_str    = f"_{qu_str}"      
    if qu_str_pr != "":
        qu_str_pr = f" - {qu_str_pr.capitalize()}"

    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    figpar = copy.deepcopy(figpar)
    figpar["init"]["ncols"] = len(oridirs) 
    
    # suptitle = (f"Mouse {mouse_n} - {stimstr_pr}" + 
    #     u"{} ".format(statstr_pr) + f"across {dimstr}{qu_str_pr}\n"
    #     f"(sess {sess_n}, {line} {plane},{nroi_str}")
    line_str, plane_str = "5", "dendrites"
    if "23" in line:
        line_str = "2/3"
    if "soma" in plane:
        plane_str = "somata"
    suptitle = f"M{mouse_n} - layer {line_str} {plane_str}{dendstr_pr}"

    savename = (f"{datatype}_tr_m{mouse_n}_sess{sess_n}{qu_str}_{stimstr}_"
        f"{plane}{dendstr}")
    
    fig, ax = plot_util.init_fig(len(oridirs), **figpar["init"])
    for o, od in enumerate(oridirs):
        cols = []
        for surp, surp_lab in zip(surps, surp_labs): 
            sub_ax = plot_util.get_subax(ax, o)
            key = f"{surp}_{od}"
            stimtype_str_pr = stimpar["stimtype"][:-1].capitalize()
            title_tr = u"{} traces ({}{})".format(stimtype_str_pr, od, deg)
            lab = "{} (n={})".format(surp_lab, tr_data["n_seqs"][key])
            sess_plot_util.add_axislabels(sub_ax, datatype=datatype)
            me  = np.asarray(tr_data["stats"][key][0])
            err = np.asarray(tr_data["stats"][key][1:])
            plot_util.plot_traces(
                sub_ax, xran, me, err, title_tr, label=lab, n_xticks=n, 
                xticks="auto")
            cols.append(sub_ax.lines[-1].get_color())
    
    
    if stimpar["stimtype"] == "gabors":
        sess_plot_util.plot_labels(
            ax, stimpar["gabfr"], cols=cols, pre=stimpar["pre"], 
            post=stimpar["post"], sharey=figpar["init"]["sharey"])

    fig.suptitle(suptitle, y=1.08)
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir


#############################################
def plot_oridir_colormap(fig_type, analyspar, stimpar, quintpar, tr_data, 
                         sess_info, figpar=None, savedir=None, log_dir=True):
    """
    plot_oridir_colormap(fig_type, analyspar, stimpar, quintpar, tr_data, 
                         sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI for a single session and optionally a single 
    quintile. (Single figure type) 

    Required args:
        - fig_type (str)  : type of figure to plot, i.e., "byplot", "byreg", 
                            "byfir" or "by{}{}" (ori/dir, deg)
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (only first session used)
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict)    : dictionary containing number of seqs for 
                                   each surprise x ori/dir combination under a 
                                   separate key
            ["scale_vals"] (dict): dictionary containing 1D array or list of 
                                   scaling values for each surprise x ori/dir 
                                   combination under a separate key.
                                   (NaN arrays for combinations with 0 seqs.)
                ["{}_min"] (num): minimum value from corresponding tr_stats 
                                  mean/medians
                ["{}_max"] (num): maximum value from corresponding tr_stats 
                                  mean/medians
            ["roi_sort"] (dict)  : dictionary containing 1D arrays or list 
                                   of peak sorting order for each 
                                   surprise x ori/dir combination under a 
                                   separate key.
                                   (NaN arrays for combinations with 0 seqs.)
            ["roi_me"] (dict)    : dictionary containing trace mean/medians for 
                                   each ROI as 2D arrays or nested lists, 
                                   structured as:
                                       ROIs x frames, 
                                   with each surprise x ori/dir combination 
                                   under a separate key
                                   (NaN arrays for combinations with 0 seqs.)
            ["xran"] (list)      : time values for the 2p frames
    
    Optional args:
        - figpar (dict)  : dictionary containing the following figure parameter 
                           dictionaries
                           default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
            ["mng"]  (dict): dictionary with parameters to manage matplotlib
        - savedir (str)  : path of directory in which to save plots.
                           default: None
        - log_dir (bool) : if True, the figure saving directory is logged.
                           default: True
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
    """

    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"], "print")

    stimstr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"])

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"]["roi"], 
            figpar["dirs"]["oridir"])

    cmap = plot_util.manage_mpl(cmap=True, nbins=100, **figpar["mng"])

    # extract some info from sess_info (only one session)
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_n, sess_n, line, plane] = [sess_info[key][0] for key in keys]

    dendstr = sess_str_util.dend_par_str(analyspar["dend"], plane, "roi")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], plane, "roi", "print")

    surps = ["reg", "surp"]
    if stimpar["stimtype"] == "gabors":
        surp_labs = surps
        var_name = "orientation"
        deg  = "deg"
        deg_pr = u"\u00B0"
        oridirs = stimpar["gab_ori"]
        n = 6
    elif stimpar["stimtype"] == "bricks":
        surp_labs = [f"{surps[i]} -> {surps[1-i]}" for i in range(len(surps))]
        var_name = "direction"
        deg  = ""
        deg_pr = ""
        oridirs = stimpar["bri_dir"]
        n = 7
    
    qu_str, qu_str_pr = quintpar["qu_lab"][0], quintpar["qu_lab_pr"][0]
    if qu_str != "":
        qu_str    = f"_{qu_str}"      
    if qu_str_pr != "":
        qu_str_pr = f" - {qu_str_pr.capitalize()}"

    if figpar is None:
        figpar = sess_plot_util.init_figpar()
    figpar = copy.deepcopy(figpar)
    figpar["init"]["ncols"] = len(oridirs)
    figpar["init"]["sharex"] = True
    
    # plot colormaps
    # gentitle = (f"Mouse {mouse_n} - {stimstr_pr} " + 
    #     u"{} ".format(statstr_pr) + "across seqs colormaps"
    #     f"{qu_str_pr} \n(sess {sess_n}, {line} {plane})")
    line_str, plane_str = "5", "dendrites"
    if "23" in line:
        line_str = "2/3"
    if "soma" in plane:
        plane_str = "somata"
    gentitle = f"Mouse {mouse_n} - layer {line_str} {plane_str}{dendstr_pr}"
    
    gen_savename = (f"roi_cm_m{mouse_n}_sess{sess_n}{qu_str}_{stimstr}_"
        f"{plane}{dendstr}")

    gen_savename = f"colormap_m{mouse_n}s{sess_n}_{stimstr}_{plane}"

    if fig_type != "byfir":
        return ""

    if fig_type == "byplot":
        scale_type = "per plot"
        peak_sort  = ""
        figpar["init"]["sharey"] = False
    elif fig_type == "byreg":
        scale_type = f"within {var_name}"
        peak_sort  = f" of {surps[0]}"
        figpar["init"]["sharey"] = False
    elif fig_type == f"by{oridirs[0]}{deg}":
        scale_type = "within surp/reg"
        peak_sort  = f" of first {var_name}"
        figpar["init"]["sharey"] = True
    elif fig_type == "byfir":
        scale_type = "across plots"
        peak_sort  = " of first plot"
        figpar["init"]["sharey"] = True
    else:
        gen_util.accepted_values_error("fig_type", fig_type, 
            ["byplot", "byreg", f"by{oridirs[0]}{deg}", "byfir"])

    subtitle = (f"ROIs sorted by peak activity{peak_sort} and scaled "
                f"{scale_type}")
    logger.info(f"- {subtitle}", extra={"spacing": TAB})
    # suptitle = f"{gentitle}\n({subtitle})"
    suptitle = gentitle
    
    # get scaled and sorted ROI mean/medians (ROI x frame)
    scaled_sort_me = roi_plots.scale_sort_trace_data(
        tr_data, fig_type, surps, oridirs)
    fig, ax = plot_util.init_fig(len(oridirs) * len(surps), **figpar["init"])
    
    xran_edges = [np.min(tr_data["xran"]), np.max(tr_data["xran"])]

    nrois = scaled_sort_me[f"{surps[0]}_{oridirs[0]}"].shape[1]
    yticks_ev = int(10 * np.max([1, np.ceil(nrois/100)])) # avoid > 10 ticks
    for o, od in enumerate(oridirs):
        for s, (surp, surp_lab) in enumerate(zip(surps, surp_labs)):    
            sub_ax = ax[s][o]
            key = f"{surp}_{od}"
            title = u"{} seqs ({}{}) (n={})".format(
                surp_lab.capitalize(), od, deg_pr, tr_data["n_seqs"][key])
            if s == 0:
                od_pr = od
                if stimpar["stimtype"] == "bricks":
                    od_pr = od_pr.capitalize()
                title = u"{}{}".format(od_pr, deg_pr)
            else:
                title = None
            x_ax = None
            y_ax = "ROIs"
            if s != 1 or o != 0:
                y_ax = ""
            if stimpar["stimtype"] == "gabors":
                x_ax = ""
            sess_plot_util.add_axislabels(
                sub_ax, fluor=analyspar["fluor"], x_ax=x_ax, y_ax=y_ax, 
                datatype="roi")
            im = plot_util.plot_colormap(
                sub_ax, scaled_sort_me[key], title=title, cmap=cmap, n_xticks=n,
                yticks_ev=yticks_ev, xran=xran_edges, xticks="auto")
            
            if stimpar["stimtype"] == "bricks":
                plot_util.add_bars(sub_ax, 0)
            else:
                sub_ax.set_xticks([])

    for s, surp in enumerate(surps):
        sub_ax = ax[s:s+1]
        if stimpar["stimtype"] == "gabors":
            sess_plot_util.plot_labels(
                sub_ax, stimpar["gabfr"], surp, pre=stimpar["pre"], 
                post=stimpar["post"], sharey=figpar["init"]["sharey"], 
                t_heis=-0.05)
    
    plot_util.add_colorbar(fig, im, len(oridirs), cm_prop=0.06)
    fig.suptitle(suptitle, fontsize="xx-large", y=1.08)
    savename = f"{gen_savename}_{fig_type}"
    fulldir = plot_util.savefig(
        fig, savename, savedir, log_dir=log_dir, **figpar["save"])
    
    plt.close(fig)
    
    return fulldir


#############################################
def plot_oridir_colormaps(analyspar, sesspar, stimpar, extrapar, quintpar, 
                          tr_data, sess_info, figpar=None, savedir=None, 
                          parallel=False):
    """
    plot_oridir_colormaps(analyspar, sesspar, stimpar, extrapar, quintpar, 
                          tr_data, sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI as colormaps for a single session and optionally
    a single quintile. 

    Required args:
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)  : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - extrapar (dict) : dictionary containing additional analysis 
                            parameters
            ["analysis"] (str): analysis type (e.g., "o")
            ["datatype"] (str): datatype (e.g., "roi")
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (only first session used)
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict)    : dictionary containing number of seqs for 
                                   each surprise x ori/dir combination under a 
                                   separate key
            ["scale_vals"] (dict): dictionary containing 1D array or list of 
                                   scaling values for each surprise x ori/dir 
                                   combination under a separate key.
                                   (NaN arrays for combinations with 0 seqs.)
                ["{}_min"] (num): minimum value from corresponding tr_stats 
                                  mean/medians
                ["{}_max"] (num): maximum value from corresponding tr_stats 
                                  mean/medians
            ["roi_sort"] (dict)  : dictionary containing 1D arrays or list 
                                   of peak sorting order for each 
                                   surprise x ori/dir combination under a 
                                   separate key.
                                   (NaN arrays for combinations with 0 seqs.)
            ["roi_me"] (dict)    : dictionary containing trace mean/medians for 
                                   each ROI as 2D arrays or nested lists, 
                                   structured as:
                                       ROIs x frames, 
                                   with each surprise x ori/dir combination 
                                   under a separate key
                                   (NaN arrays for combinations with 0 seqs.)
    
    Optional args:
        - figpar (dict)  : dictionary containing the following figure parameter 
                           dictionaries
                           default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
            ["mng"]  (dict): dictionary with parameters to manage matplotlib
        - savedir (str)  : path of directory in which to save plots.
                           default: None
        - parallel (bool): if True, some of the analysis is parallelized across 
                           CPU cores
                           default: False
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
    """

    datatype = extrapar["datatype"]
    if datatype != "roi":
        raise ValueError("Function only implemented for roi datatype.")

    if stimpar["stimtype"] == "gabors":
        oridirs = stimpar["gab_ori"]
        deg  = "deg"
    elif stimpar["stimtype"] == "bricks":
        oridirs = stimpar["bri_dir"]
        deg = ""
    
    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    figpar = copy.deepcopy(figpar)
    if figpar["save"]["use_dt"] is None:
        figpar["save"]["use_dt"] = gen_util.create_time_str()
    figpar["save"]["fig_ext"] = "svg" # svg too big

    fig_types  = ["byplot", "byreg", f"by{oridirs[0]}{deg}", "byfir"]
    fig_types = ["byfir"]
    fig_last = len(fig_types) - 1
    
    # optionally runs in parallel
    if parallel:
        n_jobs = gen_util.get_n_jobs(len(fig_types))
        fulldirs = Parallel(n_jobs=n_jobs)(delayed(plot_oridir_colormap)
            (fig_type, analyspar, stimpar, quintpar, tr_data, sess_info, 
            figpar, savedir, (f == fig_last)) 
            for f, fig_type in enumerate(fig_types)) 
        fulldir = fulldirs[-1]
    else:
        for f, fig_type in enumerate(fig_types):
            log_dir = (f == fig_last)
            fulldir = plot_oridir_colormap(
                fig_type, analyspar, stimpar, quintpar, tr_data, sess_info, 
                figpar, savedir, log_dir)

    return fulldir


#############################################
def plot_oridirs(analyspar, sesspar, stimpar, extrapar, quintpar, 
                 tr_data, sess_info, figpar=None, savedir=None, 
                 parallel=False):
    """
    plot_oridirs(analyspar, sesspar, stimpar, extrapar, quintpar, 
                 tr_data, sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI as colormaps, as well as traces across ROIs for a 
    single session and optionally a single quintile. 

    Required args:
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)  : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - extrapar (dict) : dictionary containing additional analysis 
                            parameters
            ["analysis"] (str): analysis type (e.g., "o")
            ["datatype"] (str): datatype (e.g., "roi")
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (one first session used) 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict)    : dictionary containing number of seqs for each
                                   surprise x ori/dir combination under a 
                                   separate key
            ["scale_vals"] (dict): dictionary containing 1D array or list of 
                                   scaling values for each surprise x ori/dir 
                                   combination under a separate key.
                                   (NaN arrays for combinations with 0 seqs.)
                ["{}_min"] (num): minimum value from corresponding tr_stats 
                                  mean/medians
                ["{}_max"] (num): maximum value from corresponding tr_stats 
                                  mean/medians
            ["roi_sort"] (dict) : dictionary containing 1D arrays or list of 
                                  peak sorting order for each 
                                  surprise x ori/dir combination under a 
                                  separate key.
                                  (NaN arrays for combinations with 0 seqs.)
            ["roi_me"] (dict)   : dictionary containing trace mean/medians
                                  for each ROI as 2D arrays or nested lists, 
                                  structured as:
                                      ROIs x frames, 
                                  with each surprise x ori/dir combination 
                                  under a separate key.
                                  (NaN arrays for combinations with 0 seqs.)
            ["stats"] (dict)    : dictionary containing trace mean/medians 
                                  across ROIs in 2D arrays or nested lists, 
                                  structured as: 
                                      stats (me, err) x frames
                                  with each surprise x ori/dir combination 
                                  under a separate key
                                  (NaN arrays for combinations with 0 seqs.)

    Optional args:
        - figpar (dict)  : dictionary containing the following figure parameter 
                           dictionaries
                           default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
            ["mng"]  (dict): dictionary with parameters to manage matplotlib
        - savedir (str)  : path of directory in which to save plots.
                           default: None
        - parallel (bool): if True, some of the analysis is parallelized across 
                           CPU cores
                           default: False
    """

    datatype = extrapar["datatype"]
    if datatype != "roi":
        raise ValueError("Function only implemented for roi datatype.")

    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    figpar = copy.deepcopy(figpar)
    if figpar["save"]["use_dt"] is None:
        figpar["save"]["use_dt"] = gen_util.create_time_str()

    comm_info = {"analyspar": analyspar,
                 "sesspar"  : sesspar,
                 "stimpar"  : stimpar,
                 "extrapar" : extrapar,
                 "quintpar" : quintpar,
                 "sess_info": sess_info,
                 "tr_data" : tr_data,
                 "figpar"   : figpar,
                 }

    if "roi_me" in tr_data.keys():
        plot_oridir_colormaps(
            savedir=savedir, parallel=parallel, **comm_info)

    # if "stats" in tr_data.keys():
    #     plot_oridir_traces(savedir=savedir, **comm_info)


"""
modif_analysis_plots.py

This script is used to modify plot designs on the fly for ROI and running 
analyses from dictionaries, e.g. for presentations or papers.

Authors: <NAME>

Date: October, 2018

Note: this code uses python 3.7.

"""

import copy
import logging
import os
import warnings

from joblib import Parallel, delayed
from matplotlib import pyplot as plt
import numpy as np

from util import file_util, gen_util, logger_util, plot_util
from sess_util import sess_plot_util, sess_str_util
from plot_fcts import roi_analysis_plots as roi_plots

logger = logging.getLogger(__name__)


TAB = "    "


#############################################
def plot_from_dict(dict_path, plt_bkend=None, fontdir=None, plot_tc=True, 
                   parallel=False, datetime=True):
    """
    plot_from_dict(dict_path)

    Plots data from dictionaries containing analysis parameters and results.

    Required args:
        - dict_path (str): path to dictionary to plot data from
    
    Optional_args:
        - plt_bkend (str): mpl backend to use for plotting (e.g., "agg")
                           default: None
        - fontdir (str)  : path to directory where additional fonts are stored
                           default: None
        - plot_tc (bool) : if True, tuning curves are plotted for each ROI 
                           default: True
        - parallel (bool): if True, some of the analysis is parallelized across 
                           CPU cores
                           default: False
        - datetime (bool): figpar["save"] datatime parameter (whether to 
                           place figures in a datetime folder)
                           default: True
    """

    logger.info(f"Plotting from dictionary: {dict_path}", 
        extra={"spacing": "\n"})
    
    figpar = sess_plot_util.init_figpar(
        plt_bkend=plt_bkend, fontdir=fontdir, datetime=datetime)
    plot_util.manage_mpl(cmap=False, **figpar["mng"])

    plt.rcParams["figure.titlesize"] = "xx-large"
    plt.rcParams["axes.titlesize"] = "xx-large"

    info = file_util.loadfile(dict_path)
    savedir = os.path.dirname(dict_path)

    analysis = info["extrapar"]["analysis"]

    # 0. Plots the full traces for each session
    if analysis == "f": # full traces
        plot_full_traces(figpar=figpar, savedir=savedir, **info)

    # 1. Plot average traces by quintile x surprise for each session 
    elif analysis == "t": # traces
        plot_traces_by_qu_surp_sess(figpar=figpar, savedir=savedir, **info)

    # 2. Plot average traces by quintile, locked to surprise for each session 
    elif analysis == "l": # surprise locked traces
        plot_traces_by_qu_lock_sess(figpar=figpar, savedir=savedir, **info)

    # 4. Plot autocorrelations
    elif analysis == "a": # autocorr
        plot_autocorr(figpar=figpar, savedir=savedir, **info)

    # 6. Plot colormaps and traces for orientations/directions
    elif analysis == "o": # colormaps
        plot_oridirs(figpar=figpar, savedir=savedir, parallel=parallel, **info)

    else:
        warnings.warn(f"No modified plotting function for analysis {analysis}")

    plt.close("all")

#############################################
def plot_full_traces(analyspar, sesspar, extrapar, sess_info, trace_info, 
                     figpar=None, savedir=None):
    """
    plot_full_traces(analyspar, sesspar, extrapar, sess_info, trace_info)

    From dictionaries, plots full traces for each session in a separate subplot.
    
    Returns figure name and save directory path.
    
    Required args:
        - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)    : dictionary with keys of SessPar namedtuple
        - extrapar (dict)   : dictionary containing additional analysis 
                              parameters
            ["analysis"] (str): analysis type (e.g., "t")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - sess_info (dict)  : dictionary containing information from each
                              session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - trace_info (dict): dictionary containing trace information
            ["all_tr"] (nested list): trace values structured as
                                          sess x 
                                          ([each_roi (scaled), averaged] 
                                            if datatype is "roi" x)
                                          (ROI if each_roi x)
                                          (me/err if averaged x)
                                          frames
            ["all_edges"] (list)    : edge values for each parameter, 
                                      structured as sess x block x 
                                                    edges ([start, end])
            ["all_pars"] (list)     : stimulus parameter strings structured as 
                                                    sess x block
                
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None    
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """
 
    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")
    
    sessstr = "sess{}_{}".format(sesspar["sess_n"], sesspar["plane"])
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])
    
    datatype = extrapar["datatype"]

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"))

    n_rows = 1
    if datatype == "roi":
        n_rows = 2
    
    if figpar is None:
        figpar = sess_plot_util.init_figpar()
    figpar = copy.deepcopy(figpar)
    figpar["init"]["subplot_wid"] *= 4
    figpar["init"]["subplot_hei"] *= 4
    figpar["init"]["sharex"] = True
    figpar["init"]["sharey"] = False
    
    fig, ax = plot_util.init_fig(n_sess*n_rows, n_sess, **figpar["init"])
    for i in range(n_sess):
        title = (f"Mouse {mouse_ns[i]} (sess {sess_ns[i]}, {lines[i]} "
            f"{planes[i]}{dendstr_pr}{nroi_strs[i]})")

        sub_axs = ax[:, i]
        sub_axs[0].set_title(title)
        if datatype == "roi":
            # average trace
            av_tr = np.asarray(trace_info["all_tr"][i][1])
            xran  = range(roi_tr_sep.shape[1])
            subtitle = "u{} across ROIs".format(statstr_pr)
            plot_util.plot_traces(
                sub_axs[1], xran, av_tr[0], av_tr[1:], title=subtitle, 
                xticks="auto")
            
            # each ROI trace
            roi_tr = np.asarray(trace_info["all_tr"][i][0])
            # values to add to each ROI to split them apart
            add = np.linspace(
                0, roi_tr.shape[0] * 1.5 + 1, roi_tr.shape[0])[:, np.newaxis]
            roi_tr_sep = roi_tr + add
            sub_axs[0].plot(roi_tr_sep.T)
        else:
            run_tr = np.asarray(trace_info["all_tr"][i])
            sub_axs[0].plot(run_tr)

        for b, block in enumerate(trace_info["all_edges"]):
            # all block labels to the lower plot
            plot_util.add_labels(
                sub_axs[-1], trace_info["all_pars"][b], np.mean(block), 0.85, 
                color="k")
            # add lines to both plots
            for r in range(n_rows):
                plot_util.add_bars(sub_axs[r], hbars=block)
                sess_plot_util.add_axislabels(
                    sub_axs[r], fluor=analyspar["fluor"], datatype=datatype, 
                    x_ax="")
                
    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype])

    savename = f"{datatype}_tr_{sessstr}{dendstr}"
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_traces_by_qu_surp_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats, figpar=None, 
                                savedir=None):
    """
    plot_traces_by_qu_surp_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats)

    From dictionaries, plots traces by quintile/surprise with each session in a 
    separate subplot.
    
    Returns figure name and save directory path.
    
    Required args:
        - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)    : dictionary with keys of SessPar namedtuple
        - stimpar (dict)    : dictionary with keys of StimPar namedtuple
        - extrapar (dict)   : dictionary containing additional analysis 
                              parameters
            ["analysis"] (str): analysis type (e.g., "t")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - quintpar (dict)   : dictionary with keys of QuintPar namedtuple
        - sess_info (dict)  : dictionary containing information from each
                              session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - trace_stats (dict): dictionary containing trace stats information
            ["xrans"] (list)           : time values for the 2p frames, for 
                                         each session
            ["all_stats"] (list)       : list of 4D arrays or lists of trace 
                                         data statistics across ROIs, 
                                         structured as:
                                            surp x quintiles x
                                            stats (me, err) x frames
            ["all_counts"] (array-like): number of sequences, structured as:
                                                sess x surp x quintiles
                
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None    
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """
 
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"],
        stimpar["gabk"], "print")
    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")
    
    sessstr = sess_str_util.sess_par_str(
        sesspar["sess_n"], stimpar["stimtype"], sesspar["plane"], 
        stimpar["bri_dir"], stimpar["bri_size"], stimpar["gabk"])
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])
    
    datatype = extrapar["datatype"]
    dimstr = sess_str_util.datatype_dim_str(datatype)

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"), style="par")

    xrans      = [np.asarray(xran) for xran in trace_stats["xrans"]]
    all_stats  = [np.asarray(sessst) for sessst in trace_stats["all_stats"]]
    all_counts = trace_stats["all_counts"]

    cols, lab_cols = sess_plot_util.get_quint_cols(quintpar["n_quints"])
    alpha = np.min([0.4, 0.8/quintpar["n_quints"]])

    surps = ["reg", "surp"]
    n = 6
    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    fig, ax = plot_util.init_fig(n_sess, **figpar["init"])
    for i in range(n_sess):
        sub_ax = plot_util.get_subax(ax, i)
        for s, [col, leg_ext] in enumerate(zip(cols, surps)):
            for q, qu_lab in enumerate(quintpar["qu_lab"]):
                if qu_lab != "":
                    qu_lab = f"{qu_lab.capitalize()} "
                line, plane = "5", "dendrites"
                if "23" in lines[i]:
                    line = "2/3"
                if "soma" in planes[i]:
                    plane = "somata"
                title = (f"M{mouse_ns[i]} - layer {line} {plane}{dendstr_pr}")
                # title=(f"Mouse {mouse_ns[i]} - {stimstr_pr} " + 
                #     u"{} ".format(statstr_pr) + f"across {dimstr}\n"
                #     f"(sess {sess_ns[i]}, {lines[i]} {planes[i]},"
                #     f"{nroi_strs[i]}")
                leg = None
                y_ax = ""
                if i == 0:
                    leg = f"{qu_lab}{leg_ext}"
                    y_ax = None
                # leg = f"{qu_lab}{leg_ext} ({all_counts[i][s][q]})"
                plot_util.plot_traces(
                    sub_ax, xrans[i], all_stats[i][s, q, 0], 
                    all_stats[i][s, q, 1:], title, color=col[q], alpha=alpha, 
                    label=leg, n_xticks=n, xticks="auto")
                sess_plot_util.add_axislabels(
                    sub_ax, fluor=analyspar["fluor"], datatype=datatype, 
                    y_ax=y_ax)

    plot_util.turn_off_extra(ax, n_sess)

    if stimpar["stimtype"] == "gabors": 
        sess_plot_util.plot_labels(
            ax, stimpar["gabfr"], "both", pre=stimpar["pre"], 
            post=stimpar["post"], cols=lab_cols, 
            sharey=figpar["init"]["sharey"])
    
    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["surp_qu"])

    qu_str = "_{}q".format(quintpar["n_quints"])
    if quintpar["n_quints"] == 1:
        qu_str = ""

    savename = f"{datatype}_av_{sessstr}{dendstr}{qu_str}"
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_traces_by_qu_lock_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats, 
                                figpar=None, savedir=None):
    """
    plot_traces_by_qu_lock_sess(analyspar, sesspar, stimpar, extrapar, 
                                quintpar, sess_info, trace_stats)

    From dictionaries, plots traces by quintile, locked to transitions from 
    surprise to regular or v.v. with each session in a separate subplot.
    
    Returns figure name and save directory path.
    
    Required args:
        - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)    : dictionary with keys of SessPar namedtuple
        - stimpar (dict)    : dictionary with keys of StimPar namedtuple
        - extrapar (dict)   : dictionary containing additional analysis 
                              parameters
            ["analysis"] (str): analysis type (e.g., "l")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - quintpar (dict)   : dictionary with keys of QuintPar namedtuple
        - sess_info (dict)  : dictionary containing information from each
                              session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - trace_stats (dict): dictionary containing trace stats information
            ["xrans"] (list)           : time values for the 2p frames, for 
                                         each session
            ["all_stats"] (list)       : list of 4D arrays or lists of trace 
                                         data statistics across ROIs, 
                                         structured as:
                                            (surp_len x) quintiles x
                                            stats (me, err) x frames
            ["all_counts"] (array-like): number of sequences, structured as:
                                                sess x (surp_len x) quintiles
            ["lock"] (str)             : value to which segments are locked:
                                         "surp", "reg" or "surp_split"
            ["baseline"] (num)         : number of seconds used for baseline
            ["reg_stats"] (list)       : list of 3D arrays or lists of trace 
                                         data statistics across ROIs for
                                         regular sampled sequences, 
                                         structured as:
                                            quintiles (1) x stats (me, err) 
                                            x frames
            ["reg_counts"] (array-like): number of sequences corresponding to
                                         reg_stats, structured as:
                                            sess x quintiles (1)
            
            if data is by surp_len:
            ["surp_lens"] (list)       : number of consecutive segments for
                                         each surp_len, structured by session
                
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None    
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """

    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"],
        stimpar["gabk"], "print")
    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")

    sessstr = sess_str_util.sess_par_str(
        sesspar["sess_n"], stimpar["stimtype"], sesspar["plane"], 
        stimpar["bri_dir"], stimpar["bri_size"], stimpar["gabk"])
    basestr = sess_str_util.base_par_str(trace_stats["baseline"])
    basestr_pr = sess_str_util.base_par_str(trace_stats["baseline"], "print")

    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])
    
    datatype = extrapar["datatype"]
    dimstr = sess_str_util.datatype_dim_str(datatype)

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"), style="par")

    xrans      = [np.asarray(xran) for xran in trace_stats["xrans"]]
    all_stats  = [np.asarray(sessst) for sessst in trace_stats["all_stats"]]
    reg_stats  = [np.asarray(regst) for regst in trace_stats["reg_stats"]]
    all_counts = trace_stats["all_counts"]
    reg_counts = trace_stats["reg_counts"]

    lock  = trace_stats["lock"]
    col_idx = 0
    if "surp" in lock:
        lock = "surp"
        col_idx = 1
    
    surp_lab, len_ext = "", ""
    surp_lens = [[None]] * n_sess
    offset = 0
    surp_len_default = True # plot default surprise lengths
    if (stimpar["stimtype"] == "gabors" and 
        stimpar["gabfr"] not in ["any", "all"]):
        offset = stimpar["gabfr"]
    if "surp_lens" in trace_stats.keys():
        surp_lens = trace_stats["surp_lens"]
        len_ext = "_bylen"
        if stimpar["stimtype"] == "gabors":
            offset = stimpar["gabfr"]
            surp_lens = [[sl * 1.5/4 - 0.3 * offset for sl in sls] 
                for sls in surp_lens]
        surp_len_default = False

    # plot surp_lens default values and RANGE TO PLOT
    if stimpar["stimtype"] == "gabors":
        DEFAULT_SURP_LENS = [3.0, 4.5, 6.0]
        end_val = 8.0
    else:
        DEFAULT_SURP_LENS = [2.0, 3.0, 4.0]
        end_val = 6.0
    if lock == "surp":
        st_val = -2.0
        inv = 1
    else:
        st_val = -end_val + 2
        end_val = end_val - 2
        inv = -1
    n_ticks = int((end_val - st_val)//2 + 1)
    
    if figpar is None:
        figpar = sess_plot_util.init_figpar()
    figpar = copy.deepcopy(figpar)
    figpar["init"]["subplot_wid"] = 6.5

    fig, ax = plot_util.init_fig(n_sess, **figpar["init"])
    for i, (stats, counts) in enumerate(zip(all_stats, all_counts)):
        xran = xrans[i]
        st, end = 0, len(xran)
        st_vals = list(filter(lambda i: xran[i] <= st_val, range(len(xran))))
        if len(st_vals) != 0:
            st = st_vals[-1]
        end_vals = list(filter(lambda i: xran[i] >= end_val, range(len(xran))))
        if len(end_vals) != 0:
            end = end_vals[0] + 1

        sub_ax = plot_util.get_subax(ax, i)

        line, plane = "5", "dendrites"
        if "23" in lines[i]:
            line = "2/3"
        if "soma" in planes[i]:
            plane = "somata"

        title = f"M{mouse_ns[i]} - layer {line} {plane}{dendstr_pr}"
        # title=(f"Mouse {mouse_ns[i]} - {stimstr_pr} " + 
        #     u"{statstr_pr} " + f"{lock} locked across {dimstr}{basestr_pr}"
        #     f"\n(sess {sess_ns[i]}, {lines[i]} {planes[i]}, {nroi_strs[i]}")
        y_ax = ""
        if i == 0:
            y_ax = None

        sess_plot_util.add_axislabels(
            sub_ax, fluor=analyspar["fluor"], datatype=datatype, y_ax=y_ax)
        plot_util.add_bars(sub_ax, hbars=0)
        n_lines = quintpar["n_quints"] * len(surp_lens[i])
        try: 
            cols = sess_plot_util.get_quint_cols(n_lines)[0][col_idx]
        except:
            cols = [None] * n_lines
        alpha = np.min([0.4, 0.8/n_lines])
        if stimpar["stimtype"] == "gabors":
            sess_plot_util.plot_gabfr_pattern(sub_ax, xran[st:end], 
                offset=offset, bars_omit=[0] + surp_lens[i])
        # plot regular data
        if reg_stats[i].shape[0] != 1:
            raise ValueError("Expected only one quintile for reg_stats.")
        
        leg = None
        if i == 0:
            leg = "reg"
        if surp_len_default:
            for leng in DEFAULT_SURP_LENS:
                edge = leng * inv
                if edge < 0:
                    edge = np.max([xran[st], edge])
                elif edge > 0:
                    edge = np.min([xran[end - 1], edge])
                plot_util.add_vshade(
                    sub_ax, 0, edge, color=cols[-1], alpha=0.1)
        # leg = f"reg (no lock) ({reg_counts[i][0]})"
        plot_util.plot_traces(
            sub_ax, xran[st:end], reg_stats[i][0][0, st:end], 
            reg_stats[i][0][1:, st:end], alpha=alpha, label=leg, 
            alpha_line=0.8, color="darkgray", xticks="auto")
        n = 0 # count lines plotted
        for s, surp_len in enumerate(surp_lens[i]):
            if surp_len is not None:
                edge = surp_len * inv
                if edge < 0:
                    edge = np.max([xran[st], edge])
                elif edge > 0:
                    edge = np.min([xran[end - 1], edge])
                plot_util.add_vshade(
                    sub_ax, 0, edge, color=cols[n], alpha=0.1)
                counts, stats = all_counts[i][s], all_stats[i][s]                
                surp_lab = f"surp len {surp_len}"
            else:
                # surp_lab = "surp lock"
                surp_lab = "surp"
            for q, qu_lab in enumerate(quintpar["qu_lab"]):
                if qu_lab != "":
                    qu_lab = f"{qu_lab.capitalize()} "
                lab = f"{qu_lab}{surp_lab}"
                if n == 2 and cols[n] is None:
                    sub_ax.plot([], []) # to advance the color cycle (past gray)
                #leg = f"{lab} ({counts[q]})"
                leg = None
                if i == 0:
                    leg = lab
                plot_util.plot_traces(
                    sub_ax, xran[st:end], stats[q][0, st:end], 
                    stats[q][1:, st:end], title, alpha=alpha, label=leg, 
                    n_xticks=n_ticks, alpha_line=0.8, color=cols[n], 
                    xticks="auto")
                n += 1
            if surp_len is not None:
                plot_util.add_bars(sub_ax, hbars=surp_len, 
                                   color=sub_ax.lines[-1].get_color(), alpha=1)
    
    plot_util.turn_off_extra(ax, n_sess)

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["surp_qu"], 
            f"{lock}_lock", basestr.replace("_", ""))

    qu_str = "_{}q".format(quintpar["n_quints"])
    if quintpar["n_quints"] == 1:
        qu_str = ""
 
    savename = (f"{datatype}_av_{lock}lock{len_ext}{basestr}_{sessstr}"
        f"{dendstr}{qu_str}")
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_autocorr(analyspar, sesspar, stimpar, extrapar, autocorrpar, 
                  sess_info, autocorr_data, figpar=None, savedir=None):
    """
    plot_autocorr(analyspar, sesspar, stimpar, extrapar, autocorrpar, 
                  sess_info, autocorr_data)

    From dictionaries, plots autocorrelation during stimulus blocks.

    Required args:
        - analyspar (dict)    : dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)      : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)      : dictionary with keys of StimPar namedtuple
        - extrapar (dict)     : dictionary containing additional analysis 
                                parameters
            ["analysis"] (str): analysis type (e.g., "a")
            ["datatype"] (str): datatype (e.g., "run", "roi")
        - autocorrpar (dict)  : dictionary with keys of AutocorrPar namedtuple
        - sess_info (dict)    : dictionary containing information from each
                                session 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - autocorr_data (dict): dictionary containing data to plot:
            ["xrans"] (list): list of lag values in seconds for each session
            ["stats"] (list): list of 3D arrays (or nested lists) of
                              autocorrelation statistics, structured as:
                                     sessions stats (me, err) 
                                     x ROI or 1x and 10x lag 
                                     x lag
    
    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """


    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")

    sessstr = sess_str_util.sess_par_str(
        sesspar["sess_n"], stimpar["stimtype"], sesspar["plane"], 
        stimpar["bri_dir"],stimpar["bri_size"], stimpar["gabk"]) 
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])


    datatype = extrapar["datatype"]
    if datatype == "roi":
        fluorstr_pr = sess_str_util.fluor_par_str(
            analyspar["fluor"], str_type="print")
        title_str = u"{} autocorrelation".format(fluorstr_pr)
        if not autocorrpar["byitem"]:
            title_str = f"{title_str} across ROIs" 
    elif datatype == "run":
        datastr = sess_str_util.datatype_par_str(datatype)
        title_str = f"{datastr} autocorrelation"

    if stimpar["stimtype"] == "gabors":
        seq_bars = [-1.5, 1.5] # light lines
    else:
        seq_bars = [-1.0, 1.0] # light lines

    # extract some info from sess_info
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_ns, sess_ns, lines, planes] = [sess_info[key] for key in keys]
    
    n_sess = len(mouse_ns)
    nroi_strs = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], 
        empty=(datatype!="roi"), style="par")

    xrans = autocorr_data["xrans"]
    stats = [np.asarray(stat) for stat in autocorr_data["stats"]]

    lag_s = autocorrpar["lag_s"]
    xticks = np.linspace(-lag_s, lag_s, lag_s*2+1)
    yticks = np.linspace(0, 1, 6)

    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    byitemstr = ""
    if autocorrpar["byitem"]:
        byitemstr = "_byroi"

    fig, ax = plot_util.init_fig(n_sess, **figpar["init"])
    for i in range(n_sess):
        sub_ax = plot_util.get_subax(ax, i)
        line, plane = "5", "dendrites"
        if "23" in lines[i]:
            line = "2/3"
        if "soma" in planes[i]:
            plane = "somata"
        title=f"M{mouse_ns[i]} - layer {line} {plane}{dendstr_pr}"
        # title = (f"Mouse {mouse_ns[i]} - {stimstr_pr}" + 
        #     u"{}".format(statstr_pr) + 
        #     f"{title_str}\n(sess {sess_ns[i]}, {lines[i]} {planes[i]},"
        #     f"{nroi_strs[i]})")
        # transpose to ROI/lag x stats x series
        sess_stats = stats[i].transpose(1, 0, 2) 
        for s, sub_stats in enumerate(sess_stats):
            lab = None
            if i == 0:
                if not autocorrpar["byitem"]:
                    lab = ["actual lag", "10x lag"][s]
            plot_util.plot_traces(
                sub_ax, xrans[i], sub_stats[0], sub_stats[1:], xticks=xticks, 
                yticks=yticks, alpha=0.2, label=lab)
        plot_util.add_bars(sub_ax, hbars=seq_bars)
        sub_ax.set_ylim([0, 1])
        sub_ax.set_title(title)
        sub_ax.set_xlabel("Lag (s)")

    plot_util.turn_off_extra(ax, n_sess)

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["autocorr"])

    savename = f"{datatype}_autocorr{byitemstr}_{sessstr}{dendstr}"

    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir, savename


#############################################
def plot_oridir_traces(analyspar, sesspar, stimpar, extrapar, quintpar, 
                        tr_data, sess_info, figpar=None, savedir=None):
    """
    plot_oridir_traces(analyspar, sesspar, stimpar, extrapar, quintpar, 
                       tr_data, sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI as colormaps for a single session and optionally
    a single quintile. 

    Required args:
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)  : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - extrapar (dict) : dictionary containing additional analysis 
                            parameters
            ["analysis"] (str): analysis type (e.g., "o")
            ["datatype"] (str): datatype (e.g., "roi")
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (only first session used)
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict): dictionary containing number of seqs for each
                               surprise x ori/dir combination under a 
                               separate key
            ["stats"] (dict) : dictionary containing trace mean/medians across
                               ROIs in 2D arrays or nested lists, 
                               structured as:
                                   stats (me, err) x frames
                               with each surprise x ori/dir combination under a 
                               separate key
                               (NaN arrays for combinations with 0 seqs.)
            ["xran"] (list)  : time values for the 2p frames

    Optional args:
        - figpar (dict): dictionary containing the following figure parameter 
                         dictionaries
                         default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
        - savedir (str): path of directory in which to save plots.
                         default: None
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
        - savename (str): name under which the figure is saved
    """

    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"], "print")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"], "print")

    stimstr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"])
    dendstr = sess_str_util.dend_par_str(
        analyspar["dend"], sesspar["plane"], extrapar["datatype"])

    datatype = extrapar["datatype"]
    if datatype != "roi":
        raise ValueError("Function only implemented for roi datatype.")
    dimstr = sess_str_util.datatype_dim_str(datatype)

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"][datatype], 
            figpar["dirs"]["oridir"])

    # extract some info from dictionaries
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_n, sess_n, line, plane] = [sess_info[key][0] for key in keys]

    nroi_str = sess_str_util.get_nroi_strs(
        sess_info, analyspar["remnans"], analyspar["fluor"], style="par")[0]

    xran = tr_data["xran"]

    surps = ["reg", "surp"]
    if stimpar["stimtype"] == "gabors":
        surp_labs = surps
        deg = u"\u00B0"
        oridirs = stimpar["gab_ori"]
        n = 6
    elif stimpar["stimtype"] == "bricks":
        surp_labs = [f"{surps[i]} -> {surps[1-i]}" for i in range(len(surps))]
        deg  = ""
        oridirs = stimpar["bri_dir"]
        n = 7

    qu_str, qu_str_pr = quintpar["qu_lab"][0], quintpar["qu_lab_pr"][0]
    if qu_str != "":
        qu_str    = f"_{qu_str}"      
    if qu_str_pr != "":
        qu_str_pr = f" - {qu_str_pr.capitalize()}"

    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    figpar = copy.deepcopy(figpar)
    figpar["init"]["ncols"] = len(oridirs) 
    
    # suptitle = (f"Mouse {mouse_n} - {stimstr_pr}" + 
    #     u"{} ".format(statstr_pr) + f"across {dimstr}{qu_str_pr}\n"
    #     f"(sess {sess_n}, {line} {plane},{nroi_str}")
    line_str, plane_str = "5", "dendrites"
    if "23" in line:
        line_str = "2/3"
    if "soma" in plane:
        plane_str = "somata"
    suptitle = f"M{mouse_n} - layer {line_str} {plane_str}{dendstr_pr}"

    savename = (f"{datatype}_tr_m{mouse_n}_sess{sess_n}{qu_str}_{stimstr}_"
        f"{plane}{dendstr}")
    
    fig, ax = plot_util.init_fig(len(oridirs), **figpar["init"])
    for o, od in enumerate(oridirs):
        cols = []
        for surp, surp_lab in zip(surps, surp_labs): 
            sub_ax = plot_util.get_subax(ax, o)
            key = f"{surp}_{od}"
            stimtype_str_pr = stimpar["stimtype"][:-1].capitalize()
            title_tr = u"{} traces ({}{})".format(stimtype_str_pr, od, deg)
            lab = "{} (n={})".format(surp_lab, tr_data["n_seqs"][key])
            sess_plot_util.add_axislabels(sub_ax, datatype=datatype)
            me  = np.asarray(tr_data["stats"][key][0])
            err = np.asarray(tr_data["stats"][key][1:])
            plot_util.plot_traces(
                sub_ax, xran, me, err, title_tr, label=lab, n_xticks=n, 
                xticks="auto")
            cols.append(sub_ax.lines[-1].get_color())
    
    
    if stimpar["stimtype"] == "gabors":
        sess_plot_util.plot_labels(
            ax, stimpar["gabfr"], cols=cols, pre=stimpar["pre"], 
            post=stimpar["post"], sharey=figpar["init"]["sharey"])

    fig.suptitle(suptitle, y=1.08)
    fulldir = plot_util.savefig(fig, savename, savedir, **figpar["save"])

    return fulldir


#############################################
def plot_oridir_colormap(fig_type, analyspar, stimpar, quintpar, tr_data, 
                         sess_info, figpar=None, savedir=None, log_dir=True):
    """
    plot_oridir_colormap(fig_type, analyspar, stimpar, quintpar, tr_data, 
                         sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI for a single session and optionally a single 
    quintile. (Single figure type) 

    Required args:
        - fig_type (str)  : type of figure to plot, i.e., "byplot", "byreg", 
                            "byfir" or "by{}{}" (ori/dir, deg)
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (only first session used)
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict)    : dictionary containing number of seqs for 
                                   each surprise x ori/dir combination under a 
                                   separate key
            ["scale_vals"] (dict): dictionary containing 1D array or list of 
                                   scaling values for each surprise x ori/dir 
                                   combination under a separate key.
                                   (NaN arrays for combinations with 0 seqs.)
                ["{}_min"] (num): minimum value from corresponding tr_stats 
                                  mean/medians
                ["{}_max"] (num): maximum value from corresponding tr_stats 
                                  mean/medians
            ["roi_sort"] (dict)  : dictionary containing 1D arrays or list 
                                   of peak sorting order for each 
                                   surprise x ori/dir combination under a 
                                   separate key.
                                   (NaN arrays for combinations with 0 seqs.)
            ["roi_me"] (dict)    : dictionary containing trace mean/medians for 
                                   each ROI as 2D arrays or nested lists, 
                                   structured as:
                                       ROIs x frames, 
                                   with each surprise x ori/dir combination 
                                   under a separate key
                                   (NaN arrays for combinations with 0 seqs.)
            ["xran"] (list)      : time values for the 2p frames
    
    Optional args:
        - figpar (dict)  : dictionary containing the following figure parameter 
                           dictionaries
                           default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
            ["mng"]  (dict): dictionary with parameters to manage matplotlib
        - savedir (str)  : path of directory in which to save plots.
                           default: None
        - log_dir (bool) : if True, the figure saving directory is logged.
                           default: True
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
    """

    statstr_pr = sess_str_util.stat_par_str(
        analyspar["stats"], analyspar["error"], "print")
    stimstr_pr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"], "print")

    stimstr = sess_str_util.stim_par_str(
        stimpar["stimtype"], stimpar["bri_dir"], stimpar["bri_size"], 
        stimpar["gabk"])

    if savedir is None:
        savedir = os.path.join(
            figpar["dirs"]["roi"], 
            figpar["dirs"]["oridir"])

    cmap = plot_util.manage_mpl(cmap=True, nbins=100, **figpar["mng"])

    # extract some info from sess_info (only one session)
    keys = ["mouse_ns", "sess_ns", "lines", "planes"]
    [mouse_n, sess_n, line, plane] = [sess_info[key][0] for key in keys]

    dendstr = sess_str_util.dend_par_str(analyspar["dend"], plane, "roi")
    dendstr_pr = sess_str_util.dend_par_str(
        analyspar["dend"], plane, "roi", "print")

    surps = ["reg", "surp"]
    if stimpar["stimtype"] == "gabors":
        surp_labs = surps
        var_name = "orientation"
        deg  = "deg"
        deg_pr = u"\u00B0"
        oridirs = stimpar["gab_ori"]
        n = 6
    elif stimpar["stimtype"] == "bricks":
        surp_labs = [f"{surps[i]} -> {surps[1-i]}" for i in range(len(surps))]
        var_name = "direction"
        deg  = ""
        deg_pr = ""
        oridirs = stimpar["bri_dir"]
        n = 7
    
    qu_str, qu_str_pr = quintpar["qu_lab"][0], quintpar["qu_lab_pr"][0]
    if qu_str != "":
        qu_str    = f"_{qu_str}"      
    if qu_str_pr != "":
        qu_str_pr = f" - {qu_str_pr.capitalize()}"

    if figpar is None:
        figpar = sess_plot_util.init_figpar()
    figpar = copy.deepcopy(figpar)
    figpar["init"]["ncols"] = len(oridirs)
    figpar["init"]["sharex"] = True
    
    # plot colormaps
    # gentitle = (f"Mouse {mouse_n} - {stimstr_pr} " + 
    #     u"{} ".format(statstr_pr) + "across seqs colormaps"
    #     f"{qu_str_pr} \n(sess {sess_n}, {line} {plane})")
    line_str, plane_str = "5", "dendrites"
    if "23" in line:
        line_str = "2/3"
    if "soma" in plane:
        plane_str = "somata"
    gentitle = f"Mouse {mouse_n} - layer {line_str} {plane_str}{dendstr_pr}"
    
    gen_savename = (f"roi_cm_m{mouse_n}_sess{sess_n}{qu_str}_{stimstr}_"
        f"{plane}{dendstr}")

    gen_savename = f"colormap_m{mouse_n}s{sess_n}_{stimstr}_{plane}"

    if fig_type != "byfir":
        return ""

    if fig_type == "byplot":
        scale_type = "per plot"
        peak_sort  = ""
        figpar["init"]["sharey"] = False
    elif fig_type == "byreg":
        scale_type = f"within {var_name}"
        peak_sort  = f" of {surps[0]}"
        figpar["init"]["sharey"] = False
    elif fig_type == f"by{oridirs[0]}{deg}":
        scale_type = "within surp/reg"
        peak_sort  = f" of first {var_name}"
        figpar["init"]["sharey"] = True
    elif fig_type == "byfir":
        scale_type = "across plots"
        peak_sort  = " of first plot"
        figpar["init"]["sharey"] = True
    else:
        gen_util.accepted_values_error("fig_type", fig_type, 
            ["byplot", "byreg", f"by{oridirs[0]}{deg}", "byfir"])

    subtitle = (f"ROIs sorted by peak activity{peak_sort} and scaled "
                f"{scale_type}")
    logger.info(f"- {subtitle}", extra={"spacing": TAB})
    # suptitle = f"{gentitle}\n({subtitle})"
    suptitle = gentitle
    
    # get scaled and sorted ROI mean/medians (ROI x frame)
    scaled_sort_me = roi_plots.scale_sort_trace_data(
        tr_data, fig_type, surps, oridirs)
    fig, ax = plot_util.init_fig(len(oridirs) * len(surps), **figpar["init"])
    
    xran_edges = [np.min(tr_data["xran"]), np.max(tr_data["xran"])]

    nrois = scaled_sort_me[f"{surps[0]}_{oridirs[0]}"].shape[1]
    yticks_ev = int(10 * np.max([1, np.ceil(nrois/100)])) # avoid > 10 ticks
    for o, od in enumerate(oridirs):
        for s, (surp, surp_lab) in enumerate(zip(surps, surp_labs)):    
            sub_ax = ax[s][o]
            key = f"{surp}_{od}"
            title = u"{} seqs ({}{}) (n={})".format(
                surp_lab.capitalize(), od, deg_pr, tr_data["n_seqs"][key])
            if s == 0:
                od_pr = od
                if stimpar["stimtype"] == "bricks":
                    od_pr = od_pr.capitalize()
                title = u"{}{}".format(od_pr, deg_pr)
            else:
                title = None
            x_ax = None
            y_ax = "ROIs"
            if s != 1 or o != 0:
                y_ax = ""
            if stimpar["stimtype"] == "gabors":
                x_ax = ""
            sess_plot_util.add_axislabels(
                sub_ax, fluor=analyspar["fluor"], x_ax=x_ax, y_ax=y_ax, 
                datatype="roi")
            im = plot_util.plot_colormap(
                sub_ax, scaled_sort_me[key], title=title, cmap=cmap, n_xticks=n,
                yticks_ev=yticks_ev, xran=xran_edges, xticks="auto")
            
            if stimpar["stimtype"] == "bricks":
                plot_util.add_bars(sub_ax, 0)
            else:
                sub_ax.set_xticks([])

    for s, surp in enumerate(surps):
        sub_ax = ax[s:s+1]
        if stimpar["stimtype"] == "gabors":
            sess_plot_util.plot_labels(
                sub_ax, stimpar["gabfr"], surp, pre=stimpar["pre"], 
                post=stimpar["post"], sharey=figpar["init"]["sharey"], 
                t_heis=-0.05)
    
    plot_util.add_colorbar(fig, im, len(oridirs), cm_prop=0.06)
    fig.suptitle(suptitle, fontsize="xx-large", y=1.08)
    savename = f"{gen_savename}_{fig_type}"
    fulldir = plot_util.savefig(
        fig, savename, savedir, log_dir=log_dir, **figpar["save"])
    
    plt.close(fig)
    
    return fulldir


#############################################
def plot_oridir_colormaps(analyspar, sesspar, stimpar, extrapar, quintpar, 
                          tr_data, sess_info, figpar=None, savedir=None, 
                          parallel=False):
    """
    plot_oridir_colormaps(analyspar, sesspar, stimpar, extrapar, quintpar, 
                          tr_data, sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI as colormaps for a single session and optionally
    a single quintile. 

    Required args:
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)  : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - extrapar (dict) : dictionary containing additional analysis 
                            parameters
            ["analysis"] (str): analysis type (e.g., "o")
            ["datatype"] (str): datatype (e.g., "roi")
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (only first session used)
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict)    : dictionary containing number of seqs for 
                                   each surprise x ori/dir combination under a 
                                   separate key
            ["scale_vals"] (dict): dictionary containing 1D array or list of 
                                   scaling values for each surprise x ori/dir 
                                   combination under a separate key.
                                   (NaN arrays for combinations with 0 seqs.)
                ["{}_min"] (num): minimum value from corresponding tr_stats 
                                  mean/medians
                ["{}_max"] (num): maximum value from corresponding tr_stats 
                                  mean/medians
            ["roi_sort"] (dict)  : dictionary containing 1D arrays or list 
                                   of peak sorting order for each 
                                   surprise x ori/dir combination under a 
                                   separate key.
                                   (NaN arrays for combinations with 0 seqs.)
            ["roi_me"] (dict)    : dictionary containing trace mean/medians for 
                                   each ROI as 2D arrays or nested lists, 
                                   structured as:
                                       ROIs x frames, 
                                   with each surprise x ori/dir combination 
                                   under a separate key
                                   (NaN arrays for combinations with 0 seqs.)
    
    Optional args:
        - figpar (dict)  : dictionary containing the following figure parameter 
                           dictionaries
                           default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
            ["mng"]  (dict): dictionary with parameters to manage matplotlib
        - savedir (str)  : path of directory in which to save plots.
                           default: None
        - parallel (bool): if True, some of the analysis is parallelized across 
                           CPU cores
                           default: False
    
    Returns:
        - fulldir (str) : final name of the directory in which the figure is 
                          saved (may differ from input savedir, if datetime 
                          subfolder is added.)
    """

    datatype = extrapar["datatype"]
    if datatype != "roi":
        raise ValueError("Function only implemented for roi datatype.")

    if stimpar["stimtype"] == "gabors":
        oridirs = stimpar["gab_ori"]
        deg  = "deg"
    elif stimpar["stimtype"] == "bricks":
        oridirs = stimpar["bri_dir"]
        deg = ""
    
    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    figpar = copy.deepcopy(figpar)
    if figpar["save"]["use_dt"] is None:
        figpar["save"]["use_dt"] = gen_util.create_time_str()
    figpar["save"]["fig_ext"] = "svg" # svg too big

    fig_types  = ["byplot", "byreg", f"by{oridirs[0]}{deg}", "byfir"]
    fig_types = ["byfir"]
    fig_last = len(fig_types) - 1
    
    # optionally runs in parallel
    if parallel:
        n_jobs = gen_util.get_n_jobs(len(fig_types))
        fulldirs = Parallel(n_jobs=n_jobs)(delayed(plot_oridir_colormap)
            (fig_type, analyspar, stimpar, quintpar, tr_data, sess_info, 
            figpar, savedir, (f == fig_last)) 
            for f, fig_type in enumerate(fig_types)) 
        fulldir = fulldirs[-1]
    else:
        for f, fig_type in enumerate(fig_types):
            log_dir = (f == fig_last)
            fulldir = plot_oridir_colormap(
                fig_type, analyspar, stimpar, quintpar, tr_data, sess_info, 
                figpar, savedir, log_dir)

    return fulldir


#############################################
def plot_oridirs(analyspar, sesspar, stimpar, extrapar, quintpar, 
                 tr_data, sess_info, figpar=None, savedir=None, 
                 parallel=False):
    """
    plot_oridirs(analyspar, sesspar, stimpar, extrapar, quintpar, 
                 tr_data, sess_info)

    From dictionaries, plots average activity across gabor orientations or 
    brick directions per ROI as colormaps, as well as traces across ROIs for a 
    single session and optionally a single quintile. 

    Required args:
        - analyspar (dict): dictionary with keys of AnalysPar namedtuple
        - sesspar (dict)  : dictionary with keys of SessPar namedtuple 
        - stimpar (dict)  : dictionary with keys of StimPar namedtuple
        - extrapar (dict) : dictionary containing additional analysis 
                            parameters
            ["analysis"] (str): analysis type (e.g., "o")
            ["datatype"] (str): datatype (e.g., "roi")
        - quintpar (dict) : dictionary with keys of QuintPar namedtuple
        - sess_info (dict): dictionary containing information from each
                            session (one first session used) 
            ["mouse_ns"] (list)   : mouse numbers
            ["sess_ns"] (list)    : session numbers  
            ["lines"] (list)      : mouse lines
            ["planes"] (list)     : imaging planes
            ["nrois"] (list)      : number of ROIs in session
            ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F 
                                    traces ("raw", "dff")
        - tr_data (dict)   : dictionary containing information to plot colormap.
                             Surprise x ori/dir keys are formatted as 
                             [{s}_{od}] for surp in ["reg", "surp"]
                                        and od in [0, 45, 90, 135] or 
                                                  ["right", "left"]
            ["n_seqs"] (dict)    : dictionary containing number of seqs for each
                                   surprise x ori/dir combination under a 
                                   separate key
            ["scale_vals"] (dict): dictionary containing 1D array or list of 
                                   scaling values for each surprise x ori/dir 
                                   combination under a separate key.
                                   (NaN arrays for combinations with 0 seqs.)
                ["{}_min"] (num): minimum value from corresponding tr_stats 
                                  mean/medians
                ["{}_max"] (num): maximum value from corresponding tr_stats 
                                  mean/medians
            ["roi_sort"] (dict) : dictionary containing 1D arrays or list of 
                                  peak sorting order for each 
                                  surprise x ori/dir combination under a 
                                  separate key.
                                  (NaN arrays for combinations with 0 seqs.)
            ["roi_me"] (dict)   : dictionary containing trace mean/medians
                                  for each ROI as 2D arrays or nested lists, 
                                  structured as:
                                      ROIs x frames, 
                                  with each surprise x ori/dir combination 
                                  under a separate key.
                                  (NaN arrays for combinations with 0 seqs.)
            ["stats"] (dict)    : dictionary containing trace mean/medians 
                                  across ROIs in 2D arrays or nested lists, 
                                  structured as: 
                                      stats (me, err) x frames
                                  with each surprise x ori/dir combination 
                                  under a separate key
                                  (NaN arrays for combinations with 0 seqs.)

    Optional args:
        - figpar (dict)  : dictionary containing the following figure parameter 
                           dictionaries
                           default: None
            ["init"] (dict): dictionary with figure initialization parameters
            ["save"] (dict): dictionary with figure saving parameters
            ["dirs"] (dict): dictionary with additional figure parameters
            ["mng"]  (dict): dictionary with parameters to manage matplotlib
        - savedir (str)  : path of directory in which to save plots.
                           default: None
        - parallel (bool): if True, some of the analysis is parallelized across 
                           CPU cores
                           default: False
    """

    datatype = extrapar["datatype"]
    if datatype != "roi":
        raise ValueError("Function only implemented for roi datatype.")

    if figpar is None:
        figpar = sess_plot_util.init_figpar()

    figpar = copy.deepcopy(figpar)
    if figpar["save"]["use_dt"] is None:
        figpar["save"]["use_dt"] = gen_util.create_time_str()

    comm_info = {"analyspar": analyspar,
                 "sesspar"  : sesspar,
                 "stimpar"  : stimpar,
                 "extrapar" : extrapar,
                 "quintpar" : quintpar,
                 "sess_info": sess_info,
                 "tr_data" : tr_data,
                 "figpar"   : figpar,
                 }

    if "roi_me" in tr_data.keys():
        plot_oridir_colormaps(
            savedir=savedir, parallel=parallel, **comm_info)

    # if "stats" in tr_data.keys():
    #     plot_oridir_traces(savedir=savedir, **comm_info)
modif_analysis_plots.py  This script is used to modify plot designs on the fly for ROI and running  analyses from dictionaries, e.g. for presentations or papers.  Authors: <NAME>  Date: October, 2018  Note: this code uses python 3.7. ############################################# plot_from_dict(dict_path)      Plots data from dictionaries containing analysis parameters and results.      Required args:         - dict_path (str): path to dictionary to plot data from          Optional_args:         - plt_bkend (str): mpl backend to use for plotting (e.g., "agg")                            default: None         - fontdir (str)  : path to directory where additional fonts are stored                            default: None         - plot_tc (bool) : if True, tuning curves are plotted for each ROI                             default: True         - parallel (bool): if True, some of the analysis is parallelized across                             CPU cores                            default: False         - datetime (bool): figpar["save"] datatime parameter (whether to                             place figures in a datetime folder)                            default: True # 0. Plots the full traces for each session # full traces # 1. Plot average traces by quintile x surprise for each session # traces # 2. Plot average traces by quintile, locked to surprise for each session # surprise locked traces # 4. Plot autocorrelations # autocorr # 6. Plot colormaps and traces for orientations/directions # colormaps ############################################# plot_full_traces(analyspar, sesspar, extrapar, sess_info, trace_info)      From dictionaries, plots full traces for each session in a separate subplot.          Returns figure name and save directory path.          Required args:         - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple         - sesspar (dict)    : dictionary with keys of SessPar namedtuple         - extrapar (dict)   : dictionary containing additional analysis                                parameters             ["analysis"] (str): analysis type (e.g., "t")             ["datatype"] (str): datatype (e.g., "run", "roi")         - sess_info (dict)  : dictionary containing information from each                               session              ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - trace_info (dict): dictionary containing trace information             ["all_tr"] (nested list): trace values structured as                                           sess x                                            ([each_roi (scaled), averaged]                                              if datatype is "roi" x)                                           (ROI if each_roi x)                                           (me/err if averaged x)                                           frames             ["all_edges"] (list)    : edge values for each parameter,                                        structured as sess x block x                                                      edges ([start, end])             ["all_pars"] (list)     : stimulus parameter strings structured as                                                      sess x block                      Optional args:         - figpar (dict): dictionary containing the following figure parameter                           dictionaries                          default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters         - savedir (str): path of directory in which to save plots.                          default: None              Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.)         - savename (str): name under which the figure is saved # extract some info from sess_info # average trace # each ROI trace # values to add to each ROI to split them apart # all block labels to the lower plot # add lines to both plots ############################################# plot_traces_by_qu_surp_sess(analyspar, sesspar, stimpar, extrapar,                                  quintpar, sess_info, trace_stats)      From dictionaries, plots traces by quintile/surprise with each session in a      separate subplot.          Returns figure name and save directory path.          Required args:         - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple         - sesspar (dict)    : dictionary with keys of SessPar namedtuple         - stimpar (dict)    : dictionary with keys of StimPar namedtuple         - extrapar (dict)   : dictionary containing additional analysis                                parameters             ["analysis"] (str): analysis type (e.g., "t")             ["datatype"] (str): datatype (e.g., "run", "roi")         - quintpar (dict)   : dictionary with keys of QuintPar namedtuple         - sess_info (dict)  : dictionary containing information from each                               session              ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - trace_stats (dict): dictionary containing trace stats information             ["xrans"] (list)           : time values for the 2p frames, for                                           each session             ["all_stats"] (list)       : list of 4D arrays or lists of trace                                           data statistics across ROIs,                                           structured as:                                             surp x quintiles x                                             stats (me, err) x frames             ["all_counts"] (array-like): number of sequences, structured as:                                                 sess x surp x quintiles                      Optional args:         - figpar (dict): dictionary containing the following figure parameter                           dictionaries                          default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters         - savedir (str): path of directory in which to save plots.                          default: None              Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.)         - savename (str): name under which the figure is saved # extract some info from sess_info # title=(f"Mouse {mouse_ns[i]} - {stimstr_pr} " + #     u"{} ".format(statstr_pr) + f"across {dimstr}\n" #     f"(sess {sess_ns[i]}, {lines[i]} {planes[i]}," #     f"{nroi_strs[i]}") # leg = f"{qu_lab}{leg_ext} ({all_counts[i][s][q]})" ############################################# plot_traces_by_qu_lock_sess(analyspar, sesspar, stimpar, extrapar,                                  quintpar, sess_info, trace_stats)      From dictionaries, plots traces by quintile, locked to transitions from      surprise to regular or v.v. with each session in a separate subplot.          Returns figure name and save directory path.          Required args:         - analyspar (dict)  : dictionary with keys of AnalysPar namedtuple         - sesspar (dict)    : dictionary with keys of SessPar namedtuple         - stimpar (dict)    : dictionary with keys of StimPar namedtuple         - extrapar (dict)   : dictionary containing additional analysis                                parameters             ["analysis"] (str): analysis type (e.g., "l")             ["datatype"] (str): datatype (e.g., "run", "roi")         - quintpar (dict)   : dictionary with keys of QuintPar namedtuple         - sess_info (dict)  : dictionary containing information from each                               session              ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - trace_stats (dict): dictionary containing trace stats information             ["xrans"] (list)           : time values for the 2p frames, for                                           each session             ["all_stats"] (list)       : list of 4D arrays or lists of trace                                           data statistics across ROIs,                                           structured as:                                             (surp_len x) quintiles x                                             stats (me, err) x frames             ["all_counts"] (array-like): number of sequences, structured as:                                                 sess x (surp_len x) quintiles             ["lock"] (str)             : value to which segments are locked:                                          "surp", "reg" or "surp_split"             ["baseline"] (num)         : number of seconds used for baseline             ["reg_stats"] (list)       : list of 3D arrays or lists of trace                                           data statistics across ROIs for                                          regular sampled sequences,                                           structured as:                                             quintiles (1) x stats (me, err)                                              x frames             ["reg_counts"] (array-like): number of sequences corresponding to                                          reg_stats, structured as:                                             sess x quintiles (1)                          if data is by surp_len:             ["surp_lens"] (list)       : number of consecutive segments for                                          each surp_len, structured by session                      Optional args:         - figpar (dict): dictionary containing the following figure parameter                           dictionaries                          default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters         - savedir (str): path of directory in which to save plots.                          default: None              Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.)         - savename (str): name under which the figure is saved # extract some info from sess_info # plot default surprise lengths # plot surp_lens default values and RANGE TO PLOT # title=(f"Mouse {mouse_ns[i]} - {stimstr_pr} " + #     u"{statstr_pr} " + f"{lock} locked across {dimstr}{basestr_pr}" #     f"\n(sess {sess_ns[i]}, {lines[i]} {planes[i]}, {nroi_strs[i]}") # plot regular data # leg = f"reg (no lock) ({reg_counts[i][0]})" # count lines plotted # surp_lab = "surp lock" # to advance the color cycle (past gray) #leg = f"{lab} ({counts[q]})" ############################################# plot_autocorr(analyspar, sesspar, stimpar, extrapar, autocorrpar,                    sess_info, autocorr_data)      From dictionaries, plots autocorrelation during stimulus blocks.      Required args:         - analyspar (dict)    : dictionary with keys of AnalysPar namedtuple         - sesspar (dict)      : dictionary with keys of SessPar namedtuple          - stimpar (dict)      : dictionary with keys of StimPar namedtuple         - extrapar (dict)     : dictionary containing additional analysis                                  parameters             ["analysis"] (str): analysis type (e.g., "a")             ["datatype"] (str): datatype (e.g., "run", "roi")         - autocorrpar (dict)  : dictionary with keys of AutocorrPar namedtuple         - sess_info (dict)    : dictionary containing information from each                                 session              ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - autocorr_data (dict): dictionary containing data to plot:             ["xrans"] (list): list of lag values in seconds for each session             ["stats"] (list): list of 3D arrays (or nested lists) of                               autocorrelation statistics, structured as:                                      sessions stats (me, err)                                       x ROI or 1x and 10x lag                                       x lag          Optional args:         - figpar (dict): dictionary containing the following figure parameter                           dictionaries                          default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters         - savedir (str): path of directory in which to save plots.                          default: None          Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.)         - savename (str): name under which the figure is saved # light lines # light lines # extract some info from sess_info # title = (f"Mouse {mouse_ns[i]} - {stimstr_pr}" + #     u"{}".format(statstr_pr) + #     f"{title_str}\n(sess {sess_ns[i]}, {lines[i]} {planes[i]}," #     f"{nroi_strs[i]})") # transpose to ROI/lag x stats x series ############################################# plot_oridir_traces(analyspar, sesspar, stimpar, extrapar, quintpar,                         tr_data, sess_info)      From dictionaries, plots average activity across gabor orientations or      brick directions per ROI as colormaps for a single session and optionally     a single quintile.       Required args:         - analyspar (dict): dictionary with keys of AnalysPar namedtuple         - sesspar (dict)  : dictionary with keys of SessPar namedtuple          - stimpar (dict)  : dictionary with keys of StimPar namedtuple         - extrapar (dict) : dictionary containing additional analysis                              parameters             ["analysis"] (str): analysis type (e.g., "o")             ["datatype"] (str): datatype (e.g., "roi")         - quintpar (dict) : dictionary with keys of QuintPar namedtuple         - sess_info (dict): dictionary containing information from each                             session (only first session used)             ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - tr_data (dict)   : dictionary containing information to plot colormap.                              Surprise x ori/dir keys are formatted as                               [{s}_{od}] for surp in ["reg", "surp"]                                         and od in [0, 45, 90, 135] or                                                    ["right", "left"]             ["n_seqs"] (dict): dictionary containing number of seqs for each                                surprise x ori/dir combination under a                                 separate key             ["stats"] (dict) : dictionary containing trace mean/medians across                                ROIs in 2D arrays or nested lists,                                 structured as:                                    stats (me, err) x frames                                with each surprise x ori/dir combination under a                                 separate key                                (NaN arrays for combinations with 0 seqs.)             ["xran"] (list)  : time values for the 2p frames      Optional args:         - figpar (dict): dictionary containing the following figure parameter                           dictionaries                          default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters         - savedir (str): path of directory in which to save plots.                          default: None          Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.)         - savename (str): name under which the figure is saved # extract some info from dictionaries # suptitle = (f"Mouse {mouse_n} - {stimstr_pr}" + #     u"{} ".format(statstr_pr) + f"across {dimstr}{qu_str_pr}\n" #     f"(sess {sess_n}, {line} {plane},{nroi_str}") ############################################# plot_oridir_colormap(fig_type, analyspar, stimpar, quintpar, tr_data,                           sess_info)      From dictionaries, plots average activity across gabor orientations or      brick directions per ROI for a single session and optionally a single      quintile. (Single figure type)       Required args:         - fig_type (str)  : type of figure to plot, i.e., "byplot", "byreg",                              "byfir" or "by{}{}" (ori/dir, deg)         - analyspar (dict): dictionary with keys of AnalysPar namedtuple         - stimpar (dict)  : dictionary with keys of StimPar namedtuple         - quintpar (dict) : dictionary with keys of QuintPar namedtuple         - sess_info (dict): dictionary containing information from each                             session (only first session used)             ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - tr_data (dict)   : dictionary containing information to plot colormap.                              Surprise x ori/dir keys are formatted as                               [{s}_{od}] for surp in ["reg", "surp"]                                         and od in [0, 45, 90, 135] or                                                    ["right", "left"]             ["n_seqs"] (dict)    : dictionary containing number of seqs for                                     each surprise x ori/dir combination under a                                     separate key             ["scale_vals"] (dict): dictionary containing 1D array or list of                                     scaling values for each surprise x ori/dir                                     combination under a separate key.                                    (NaN arrays for combinations with 0 seqs.)                 ["{}_min"] (num): minimum value from corresponding tr_stats                                    mean/medians                 ["{}_max"] (num): maximum value from corresponding tr_stats                                    mean/medians             ["roi_sort"] (dict)  : dictionary containing 1D arrays or list                                     of peak sorting order for each                                     surprise x ori/dir combination under a                                     separate key.                                    (NaN arrays for combinations with 0 seqs.)             ["roi_me"] (dict)    : dictionary containing trace mean/medians for                                     each ROI as 2D arrays or nested lists,                                     structured as:                                        ROIs x frames,                                     with each surprise x ori/dir combination                                     under a separate key                                    (NaN arrays for combinations with 0 seqs.)             ["xran"] (list)      : time values for the 2p frames          Optional args:         - figpar (dict)  : dictionary containing the following figure parameter                             dictionaries                            default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters             ["mng"]  (dict): dictionary with parameters to manage matplotlib         - savedir (str)  : path of directory in which to save plots.                            default: None         - log_dir (bool) : if True, the figure saving directory is logged.                            default: True          Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.) # extract some info from sess_info (only one session) # plot colormaps # gentitle = (f"Mouse {mouse_n} - {stimstr_pr} " + #     u"{} ".format(statstr_pr) + "across seqs colormaps" #     f"{qu_str_pr} \n(sess {sess_n}, {line} {plane})") # suptitle = f"{gentitle}\n({subtitle})" # get scaled and sorted ROI mean/medians (ROI x frame) # avoid > 10 ticks ############################################# plot_oridir_colormaps(analyspar, sesspar, stimpar, extrapar, quintpar,                            tr_data, sess_info)      From dictionaries, plots average activity across gabor orientations or      brick directions per ROI as colormaps for a single session and optionally     a single quintile.       Required args:         - analyspar (dict): dictionary with keys of AnalysPar namedtuple         - sesspar (dict)  : dictionary with keys of SessPar namedtuple          - stimpar (dict)  : dictionary with keys of StimPar namedtuple         - extrapar (dict) : dictionary containing additional analysis                              parameters             ["analysis"] (str): analysis type (e.g., "o")             ["datatype"] (str): datatype (e.g., "roi")         - quintpar (dict) : dictionary with keys of QuintPar namedtuple         - sess_info (dict): dictionary containing information from each                             session (only first session used)             ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - tr_data (dict)   : dictionary containing information to plot colormap.                              Surprise x ori/dir keys are formatted as                               [{s}_{od}] for surp in ["reg", "surp"]                                         and od in [0, 45, 90, 135] or                                                    ["right", "left"]             ["n_seqs"] (dict)    : dictionary containing number of seqs for                                     each surprise x ori/dir combination under a                                     separate key             ["scale_vals"] (dict): dictionary containing 1D array or list of                                     scaling values for each surprise x ori/dir                                     combination under a separate key.                                    (NaN arrays for combinations with 0 seqs.)                 ["{}_min"] (num): minimum value from corresponding tr_stats                                    mean/medians                 ["{}_max"] (num): maximum value from corresponding tr_stats                                    mean/medians             ["roi_sort"] (dict)  : dictionary containing 1D arrays or list                                     of peak sorting order for each                                     surprise x ori/dir combination under a                                     separate key.                                    (NaN arrays for combinations with 0 seqs.)             ["roi_me"] (dict)    : dictionary containing trace mean/medians for                                     each ROI as 2D arrays or nested lists,                                     structured as:                                        ROIs x frames,                                     with each surprise x ori/dir combination                                     under a separate key                                    (NaN arrays for combinations with 0 seqs.)          Optional args:         - figpar (dict)  : dictionary containing the following figure parameter                             dictionaries                            default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters             ["mng"]  (dict): dictionary with parameters to manage matplotlib         - savedir (str)  : path of directory in which to save plots.                            default: None         - parallel (bool): if True, some of the analysis is parallelized across                             CPU cores                            default: False          Returns:         - fulldir (str) : final name of the directory in which the figure is                            saved (may differ from input savedir, if datetime                            subfolder is added.) # svg too big # optionally runs in parallel ############################################# plot_oridirs(analyspar, sesspar, stimpar, extrapar, quintpar,                   tr_data, sess_info)      From dictionaries, plots average activity across gabor orientations or      brick directions per ROI as colormaps, as well as traces across ROIs for a      single session and optionally a single quintile.       Required args:         - analyspar (dict): dictionary with keys of AnalysPar namedtuple         - sesspar (dict)  : dictionary with keys of SessPar namedtuple          - stimpar (dict)  : dictionary with keys of StimPar namedtuple         - extrapar (dict) : dictionary containing additional analysis                              parameters             ["analysis"] (str): analysis type (e.g., "o")             ["datatype"] (str): datatype (e.g., "roi")         - quintpar (dict) : dictionary with keys of QuintPar namedtuple         - sess_info (dict): dictionary containing information from each                             session (one first session used)              ["mouse_ns"] (list)   : mouse numbers             ["sess_ns"] (list)    : session numbers               ["lines"] (list)      : mouse lines             ["planes"] (list)     : imaging planes             ["nrois"] (list)      : number of ROIs in session             ["nanrois_{}"] (list) : list of ROIs with NaNs/Infs in raw or dF/F                                      traces ("raw", "dff")         - tr_data (dict)   : dictionary containing information to plot colormap.                              Surprise x ori/dir keys are formatted as                               [{s}_{od}] for surp in ["reg", "surp"]                                         and od in [0, 45, 90, 135] or                                                    ["right", "left"]             ["n_seqs"] (dict)    : dictionary containing number of seqs for each                                    surprise x ori/dir combination under a                                     separate key             ["scale_vals"] (dict): dictionary containing 1D array or list of                                     scaling values for each surprise x ori/dir                                     combination under a separate key.                                    (NaN arrays for combinations with 0 seqs.)                 ["{}_min"] (num): minimum value from corresponding tr_stats                                    mean/medians                 ["{}_max"] (num): maximum value from corresponding tr_stats                                    mean/medians             ["roi_sort"] (dict) : dictionary containing 1D arrays or list of                                    peak sorting order for each                                    surprise x ori/dir combination under a                                    separate key.                                   (NaN arrays for combinations with 0 seqs.)             ["roi_me"] (dict)   : dictionary containing trace mean/medians                                   for each ROI as 2D arrays or nested lists,                                    structured as:                                       ROIs x frames,                                    with each surprise x ori/dir combination                                    under a separate key.                                   (NaN arrays for combinations with 0 seqs.)             ["stats"] (dict)    : dictionary containing trace mean/medians                                    across ROIs in 2D arrays or nested lists,                                    structured as:                                        stats (me, err) x frames                                   with each surprise x ori/dir combination                                    under a separate key                                   (NaN arrays for combinations with 0 seqs.)      Optional args:         - figpar (dict)  : dictionary containing the following figure parameter                             dictionaries                            default: None             ["init"] (dict): dictionary with figure initialization parameters             ["save"] (dict): dictionary with figure saving parameters             ["dirs"] (dict): dictionary with additional figure parameters             ["mng"]  (dict): dictionary with parameters to manage matplotlib         - savedir (str)  : path of directory in which to save plots.                            default: None         - parallel (bool): if True, some of the analysis is parallelized across                             CPU cores                            default: False # if "stats" in tr_data.keys(): #     plot_oridir_traces(savedir=savedir, **comm_info)