""" Utility module with various helper functions and subplot generaters"""
import pandas as pd
import numpy as np
import os
import sys
import re 
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import pyplot as plt
from matplotlib.patches import Patch
from matplotlib.colors import is_color_like
from scipy.cluster.hierarchy import distance
from scipy.cluster.hierarchy import linkage 
from scipy.cluster.hierarchy import dendrogram 
from scipy.cluster.hierarchy import optimal_leaf_ordering
from scipy.spatial.distance import pdist

import DPre.main.config as config
from DPre.main._logger import logger, spacer


def _add_mg_types(data, down):
    """Add markergene type index (up and down) to columns at level 0"""
    orig_order = data.columns.unique(0)
    updown_idx = ['up']*data.shape[1]
    if down:
        updown_idx.extend(['down']*data.shape[1])
        data = pd.concat((data, data), axis=1)
    data.columns = _add_level(data.columns, updown_idx)
    return data.reindex(orig_order, axis=1, level=1)

def _bool_to_int_genes(_diff, trans_updown=True, return_merged=False):
    """Take _diff input and convert up-genes to +1, down-genes to -1, optionally 
       transfer up- and down values to each other makeing up- and down subframes
       equal"""
    int_diff = _diff.astype(int)
    if 'down' in _diff.columns.unique(0):
        int_diff['down'] *= -1
    if trans_updown and 'down' in _diff.columns.unique(0):
        int_diff['up'] = int_diff['up'].mask(_diff['down'], -1)
        int_diff['down'] = int_diff['down'].mask(_diff['up'], 1)
    if not return_merged:
        return int_diff
    else:
        return int_diff.xs('up', 1, 0)

def _add_mgtmean(agg):
    """Prodcue the mean between aggregated up- and down mg similarity values"""
    agg_mean = agg.groupby(axis=1, level=1, sort=False).mean()
    agg_mean.columns = _add_level(agg_mean.columns, 'mean')
    return pd.concat([agg, agg_mean], axis=1)

def _add_log2_z(expr, rowwise_sd=False):
    """Compute log2 and z-transformed expression data. Substitute read count in
       expression. Optionally, compute the standad deviation row(gene)-wise. 
       Is used for large datasets like reference transcriptome libraries. 
    """
    expr = np.log2(expr +1)
    expr.columns = _add_level(expr.columns, 'log2', at=1)

    m = expr.values.mean()
    s = expr.values.std() if not rowwise_sd else expr.std(1)
        
    z_expr = expr.apply(lambda c: (c-m) /s)
    z_expr.columns = _add_level(z_expr.columns, 'z', 1)
    return pd.concat((expr, z_expr), axis=1).reindex(expr.columns.unique(0), 
                                                     axis=1, level=0)

def _add_level(index, label, at=0, replace=False, name=''):
    """Add a level with labels 'label' to a pd.MultiIndex"""
    index = pd.DataFrame(index=index)
    if replace:
        index.reset_index(level=at, drop=True, inplace=True)
    index[name] = label
    order = list(range(index.index.nlevels))
    order.insert(at, -1)
    return index.set_index(name, append=True).reorder_levels(order).index

def _get_gene_ann(species):
    """Open the gene annotation reference file (mouse/ human) and return it"""
    path = os.path.dirname(__file__)
    if species == 'mouse':
        return pd.read_pickle(path + '/../gene_ann/mg_ensembl96_GRCm38.p6.gzip')
    elif species == 'human':
        return pd.read_pickle(path + '/../gene_ann/hg_GRCh38.p12.gzip')
    else:
        logger.info('')
        logger.error('Invalid input for species: `{}`. Valid are `mouse` and '
                     '`human`'.format(species))
        sys.exit(1)

def annotate(ensgs, species):
    """ Annotate mouse or human ensg keys. Return the gene names.

    DPre references the ensembl gene annotation v.96 located at 
    DPre/gene_ann. 

    Args:
        ensgs (list, pandas.Index): The collection of ensg keys to annotate
        species (str): The origin species of the genes, 'mouse' or 'human'.
    
    Returns:
        annotated pandas.Index
    """
    ref = _get_gene_ann(species)
    try:
        return pd.Index(ref.reindex(ensgs).name.values)
    except Exception as e:
        logger.error('{}\nDPre references the ensembl gene annotaiton v.96. '
                     'Differently annotated datasets may cause problems.'
                     .format(e))
        sys.exit(1)

def get_ensgs(names, species):
    """ Return the ensg keys for a list of gene names.

    DPre references the ensembl gene annotation v.96 located at 
    DPre/gene_ann. If a gene name has multiple ensg keys, this gene will appear 
    last in the DataFrame regardless of the input order.
    
    Args:
        names (list pandas.Index): The collection of names to return ensg keys 
            for
        species (str): The origin species of the genes, 'mouse' or 'human'.
    
    Returns:
        pandas.Index of ensg keys 
    """
    ref = _get_gene_ann(species)
    try:
        ann = ref.reindex(ref.index[ref.name.isin(names)]).reset_index()
        if ann.name.duplicated().any():
            dupl = pd.Index(ann.name).duplicated()
            ann_dr = ann[~dupl]
            ann_du = ann[dupl]
            ann_dr = ann_dr.set_index('name').reindex(names).reset_index()
            ann_dr.rename({'index': 'name'}, axis=1, inplace=1)
            ann = ann_dr.append(ann_du, sort=False)
            ann.index = np.arange(ann.shape[0])
        else:
            ann = ann.set_index('name').reindex(names).reset_index()
            ann.rename({'index': 'name'}, axis=1, inplace=1)
        return ann
    except Exception as e:
        logger.error('{}\nDPre references the ensembl gene annotaiton v.96. '
                     'Differently annotated datasets may cause problems.'
                     .format(e))
        sys.exit(1)

def _align_indices(data, order, axis=1):
    """Align the indices/ columns in a collection of pandas objects to order"""
    for i in range(len(data)):
        if data[i] is not None:
            data[i] = data[i].reindex(order, axis=axis)
    return data


def _init_figure(fig_widths, fig_heights, nplts, spacers):
    """Calculate the size proportion of each plot element, create figure"""
    width, height = sum(fig_widths), sum(fig_heights)
    ratio = {'width_ratios': list(map(lambda w: w/width, 
                                        fig_widths[1:-2])),
             'height_ratios': list(map(lambda h: h/height, 
                                        fig_heights[1:-2]))}
    # init figure
    fig, axes = plt.subplots(*nplts, figsize=(width, height), 
                                gridspec_kw=ratio)
    if not isinstance(axes, np.ndarray):
        axes = np.array([axes])
    axes = _clean_axes(axes)
    
    wspace_prop = spacers[0] /np.array(fig_widths[1:-2]).mean()
    hspace_prop = spacers[1] /np.array(fig_heights[1:-2]).mean()
    adj_args = {'left': fig_widths[0] /width, 
                'wspace': wspace_prop,
                'right': 1 - fig_widths[-1] /width,
                'top': 1 - fig_heights[0] /height, 
                'hspace': hspace_prop,
                'bottom': fig_heights[-1] /height}
    fig.subplots_adjust(**adj_args)
    return fig, axes

def _open_file(filename):
    """Open a file based on the filename ending or if not present
       on config.SAVE_FORMAT. Must be supporte by matplotlib."""
    valid = plt.figure().canvas.get_supported_filetypes()
    if not any([filename.endswith(val_format) for val_format in valid]):
        if config.SAVE_FORMAT in valid:
            filename += '.' + config.SAVE_FORMAT
        else:
            logger.error('The value for config.SAVE_FORMAT `{}` is not '
                         'supported by matplotlib. Valid formats are:\n{}'
                         .format(config.SAVE_FORMAT, ', '.join(list(valid.keys()))))
            sys.exit(1)
        
    if filename.endswith('.pdf'):
        return filename, PdfPages(filename)
    else:
        return filename, None
def _save_file(fig, filename=None, pp=None, close_pp=False):
    """Save pdf if pp is passed, otherwise use filename to save as 
       config.SAVE_FORMAT"""
    if pp:
        fig.savefig(pp, format='pdf')
        if close_pp:
            pp.close()
            plt.close(fig)
    elif filename:
        replace = ['$\\mathit{', '}$']
        for repl in replace:
            filename = filename.replace(repl, '')
        fig.savefig(filename)
        plt.close(fig)

def _clean_axes(axes):
    """Remove all spines, ticks and tickalabels"""
    np.array([axes])
    for ax in axes.flatten():
        [s.set_visible(False) for s in ax.spines.values()]
        ax.tick_params(bottom=False, left=False, labelbottom=False, 
                       labelleft=False)
    return axes

def _make_title(differential, metric, el1, el2, pref='', postf=''):
    """Produce the plot title based on plot parmaters, pref and posf are used
       for plot specific adjustments; return the title string"""

    metric_title = 'metric: '
    if metric == 'euclid':
        metric_title += 'L1 Euclidean distance'
    elif metric == 'pearson':
        metric_title += 'Pearson correlation' 
    elif metric == 'cosine':
        metric_title += 'cosine similarity' 
    elif metric == 'intersect':
         metric_title += 'marker gene intersect'

    if differential:
        dtype = 'Change in '
    else:
        dtype = 'Absolute '
    if pref:
        dtype = dtype.lower()
    title = ('{}{}{}transcriptional similarity \nof {} & {}\n{}'
              .format(pref, dtype, postf, el1, el2, metric_title))
    return title[0].upper() + title[1:]

def _heatmap_cluster(dat, where, ax, metric):
    """Cluster the columns or index with scipy; return the new order"""
    ax.set_visible(True)
    d = dat.T if metric == 'columns' else dat 
    Y = pdist(d, metric='euclidean')
    Z = linkage(Y, method='complete', metric='euclidean')
    order = dendrogram(Z,
                        count_sort = True,
                        no_labels = True,
                        orientation = where, 
                        labels = d.index, 
                        above_threshold_color = config.dendrogram_colors[0],
                        ax = ax)['ivl']
    if metric == 'rows':
        # for some reason reversed?
        order = order[::-1]
    return order

def _plot_distance_bar(axes, data, ctrl_lbl, bar_args, draw_colorbar=False, 
                       cb_lbl=None, fig=None, pivot=None, w=None, h=None):
    """Draw the distance bar on top of the heatmap"""
    # set ylabel on the left 
    axes[0].tick_params(labelleft=True)
    axes[0].set_ylim(0, 1)
    axes[0].set_yticks((.5,))
    axes[0].set_yticklabels((ctrl_lbl,), x=.5)
    # draw the heatmap
    ax = axes[1]
    [s.set_visible(True) for s in ax.spines.values()]
    im = ax.imshow(data.values, aspect='auto', **bar_args)

    # setup the colorbar legend    
    if draw_colorbar:
        at = (config.CB_LEFT_SEC/w, 1- config.CB_TOP/h, config.CB_WIDTH/w, 
              config.CB_HEIGHT/h)
        cb = ax.figure.colorbar(im, cax=fig.add_axes(at), alpha =.3,
                                orientation='horizontal')
        bar_ticks = (bar_args['vmin'], bar_args['vmax'])
        cb.set_ticks(bar_ticks)
        cb.ax.set_xticklabels(bar_ticks)
        if pivot:
            cb.ax.tick_params(labelrotation=90)
        cb.ax.set_xlabel(cb_lbl)
        cb.ax.get_xaxis().set_label_position('top')

def _setup_heatmap_xy(x_y, ax, lbls, pivot, hide_lbls, lbl_size, colors):
    """Setting all paramters for the x- and y axis of the two heatmap plots"""
    dim = len(lbls)
    if x_y == 'x':
        # X-axis setup, colorbar bottom
        ax.set_xlim(0, dim)
        ticks = np.arange(.5, dim)
        ax.set_xticks(ticks)
        if not hide_lbls:
            ax.tick_params(labelbottom=True)
            fs = lbl_size*config.FONTS if lbl_size else config.FONTS
            if not pivot:
                ax.set_xticklabels(lbls, rotation=45, ha='right', fontsize=fs, 
                                   rotation_mode='anchor', y=-.5)
            else:
                ax.set_xticklabels(lbls, rotation=90, ha='right', va='center', 
                                   fontsize=fs, rotation_mode='anchor', y=-.5)
        if colors:
            ax.bar(ticks, 1, 1, color=colors)
    
    elif x_y == 'y':
        ax.set_ylim((-.1, dim +.01))
        ax.set_yticks(np.arange(.5, dim))
        if not hide_lbls:
            ax.tick_params(labelleft=True) 
            fs = lbl_size*config.FONTS if lbl_size else config.FONTS
            if not pivot:
                ax.set_yticklabels(lbls, x=.5, fontsize=fs)
            else:
                ax.set_yticklabels(lbls, rotation=45, ha='right', x=1,
                                   fontsize=fs, rotation_mode='anchor')
        if colors:
            ax.bar(0, 1, width=1, color=colors, bottom=np.arange(len(lbls)))

def _check_args(trg, smp, metric, differential, 
                hide_distance_bar=None, reorder_to_distance_bar=None,
                distance_bar_range=None, cluster_hmx=None, 
                display_markergenes=False):
    """General purpose plot argument checker; returns (modified) input values"""
    def check_metric(metric, trg, smp, diff):
        # check if the samples and targets have equivalent data to compare
        if smp._type_name != 'samples':
            logger.error('The passed `samples` are not of type DPre.smaples.')
            sys.exit(1)
        if metric is None:
            if trg._has_expr and smp._has_expr:
                metric = 'cosine'
            elif trg._has_diff and smp._has_diff:
                metric = 'intersect'
            else:
                logger.error('Either initiate targets and samples with '
                             'expression or with markergenes and diff genes.')
                sys.exit(1)
        msg = 'The {} were initiated without {} data. Cannot use `{}` similarity.'
        if metric not in ('euclid', 'cosine', 'pearson', 'intersect'):
            logger.error('Invalid `metric` input: `{}`. Valid are `euclid`, '
                         '`cosine`, `pearson`, and `intersect`'.format(metric))
            sys.exit(1)
        if metric in ['euclid', 'cosine', 'pearson']:
            if not trg._has_expr:
                logger.error(msg.format('targets', 'expression', metric))
                sys.exit(1)
            elif not smp._has_expr:
                logger.error(msg.format('samples', 'expression', metric))
                sys.exit(1)
                if diff and not smp._ctrl:
                    logger.error('To plot the changes in transcriptional similarity '
                    'with metric = `{}`, the samples must be initiated '
                    'with a control. For absolute, pass differential = ' 
                    'False.'.format(metric))
                    sys.exit(1)
        elif metric == 'intersect':
            if not trg._has_diff:
                logger.error(msg.format('targets', 'merker gene', metric))
                sys.exit(1)
            elif not smp._has_diff:
                logger.error(msg.format('samples', 'diff genes', metric))
                sys.exit(1)
        
        return metric

    # checks for all plots
    metric = check_metric(metric, trg, smp, differential)
    if metric == 'intersect' and not differential:
        differential = True
        logger.warning('For the `intersect` similarity metric, '
                       'differential cannot be False. Was set to True.')

    # checks for 2 heatmaps
    if metric != 'intersect' and not hide_distance_bar and not smp._ctrl:
        hide_distance_bar = True
        logger.warning('`hide_distance_bar` must be True '
                    'for metric = `{}` if the samples data is '
                    'initialized without a control. Set to True.'.format(metric))
    if reorder_to_distance_bar and hide_distance_bar:
        reorder_to_distance_bar = False
        logger.warning('When `reorder_to_distance_bar` is True, '
                        '`hide_distance_bar` cannot be True. Set '
                        'to False.')
    if reorder_to_distance_bar and cluster_hmx:
        cluster_hmx = False
        logger.warning('Both `reorder_to_distance_bar` and '
                        '`cluster_genes` were set as True. '
                        '`cluster_genes` will be ignored.')                    
    if not differential and distance_bar_range is not None:
        distance_bar_range = None
        logger.warning('The argument `distance_bar_range` is invalid '
                       'and ignored when differential = False. To apply'
                       ' a custom range, please use "heatmap_range".')
            
    if display_markergenes is not False:
        # checks for target_sim and ranked_sim plots
        val = ['mean', 'up', 'down']
        if display_markergenes not in val:
            logger.warning('Invalid input for display_markergenes: `{}`. ' 
                           'Valid are {}. Set to default `{}`'
                           .format(display_markergenes, val, val[0]))
            display_markergenes = val[0] 

        if display_markergenes == val[2] and not trg._down_mgs: 
            logger.error('Cannot display down markergene similarity because'
                         ' the targets were not initiated with down '
                         'markergenes.')
            sys.exit(1)

    return metric, differential, hide_distance_bar, reorder_to_distance_bar, \
           distance_bar_range, cluster_hmx, display_markergenes

def plot_color_legend(labels, colors, ncolumns=1, filename='color_legend'):
    """Plot a custom color legend.
    
       Takes a list of labels and colors and links them to produce a color 
       legend. Useful for marking  sub-groups in samples/ targets elements.

    Args:
        labels (list): the list of labels in the legend
        colors (list): the list of colors correspoding to the labels. Colors 
            must be interpretable by matplotlib: for example, 'w', #ffffff, 
            (1,1,1) all refer to white.
        filename (str, optional): the filename to save the legend. Defaults to
            './color_legend.' + config.SAVE_FORMAT
        ncolumns (int, optional): the number of columns in the legend. Defaults 
            to 1.
    """
    spacer.info('\n\n')
    assert len(colors) == len(labels), 'colors and labels differ in length'
    inv_cols = [c for c in colors if not is_color_like(c)]
    if inv_cols:
        logger.error('The following colors are not recognized as colors by '
                     'matplotlib: {}'.format(inv_cols))
        sys.exit(1)
    filename, pp = _open_file(filename)
    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
    _clean_axes(np.array([ax]))
    ax.legend(handles=[Patch(color=colors[i], label=labels[i]) 
                for i in range(len(colors))], loc='center', ncol=ncolumns)
    _save_file(fig, filename=filename, pp=pp, close_pp=True)

    logger.info('Color legend generated and saved at {}/{}'
                .format(os.path.abspath(os.curdir), filename))

def add_diff_genes_from_z(samples, diff_z_threshold=2):
    if not samples._ctrl:
        logger.error('The samples `{}` were not initialized with a control.'
                     'To generate a list of differential genes, a control is '
                     'required.'.format(samples.name))
        sys.exit(1)
    expr = samples._expr.xs('z', 1, 1, False)
    expr = expr.apply(lambda smp: smp - expr.loc(1)[(samples._ctrl, 'z')])

    up = expr.mask(~(expr>diff_z_threshold), False).astype(bool)
    up.columns = pd.MultiIndex.from_product([['up'], up.columns.unique(0)])
    down = expr.mask(~(expr<-diff_z_threshold), False).astype(bool)
    down.columns = pd.MultiIndex.from_product([['down'], down.columns.unique(0)])

    samples._diff = pd.concat((up, down), axis=1)
    samples._has_diff = True
    spacer.info('\n')
    n = samples._diff.sum().unstack(0).reindex(samples.names).to_string()
    logger.info('Differential genes were added to the sample. Number of marker '
                'genes:\n{}\n'.format(n))
""" Utility module with various helper functions and subplot generaters"""
import pandas as pd
import numpy as np
import os
import sys
import re 
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import pyplot as plt
from matplotlib.patches import Patch
from matplotlib.colors import is_color_like
from scipy.cluster.hierarchy import distance
from scipy.cluster.hierarchy import linkage 
from scipy.cluster.hierarchy import dendrogram 
from scipy.cluster.hierarchy import optimal_leaf_ordering
from scipy.spatial.distance import pdist

import DPre.main.config as config
from DPre.main._logger import logger, spacer


def _add_mg_types(data, down):
    """Add markergene type index (up and down) to columns at level 0"""
    orig_order = data.columns.unique(0)
    updown_idx = ['up']*data.shape[1]
    if down:
        updown_idx.extend(['down']*data.shape[1])
        data = pd.concat((data, data), axis=1)
    data.columns = _add_level(data.columns, updown_idx)
    return data.reindex(orig_order, axis=1, level=1)

def _bool_to_int_genes(_diff, trans_updown=True, return_merged=False):
    """Take _diff input and convert up-genes to +1, down-genes to -1, optionally 
       transfer up- and down values to each other makeing up- and down subframes
       equal"""
    int_diff = _diff.astype(int)
    if 'down' in _diff.columns.unique(0):
        int_diff['down'] *= -1
    if trans_updown and 'down' in _diff.columns.unique(0):
        int_diff['up'] = int_diff['up'].mask(_diff['down'], -1)
        int_diff['down'] = int_diff['down'].mask(_diff['up'], 1)
    if not return_merged:
        return int_diff
    else:
        return int_diff.xs('up', 1, 0)

def _add_mgtmean(agg):
    """Prodcue the mean between aggregated up- and down mg similarity values"""
    agg_mean = agg.groupby(axis=1, level=1, sort=False).mean()
    agg_mean.columns = _add_level(agg_mean.columns, 'mean')
    return pd.concat([agg, agg_mean], axis=1)

def _add_log2_z(expr, rowwise_sd=False):
    """Compute log2 and z-transformed expression data. Substitute read count in
       expression. Optionally, compute the standad deviation row(gene)-wise. 
       Is used for large datasets like reference transcriptome libraries. 
    """
    expr = np.log2(expr +1)
    expr.columns = _add_level(expr.columns, 'log2', at=1)

    m = expr.values.mean()
    s = expr.values.std() if not rowwise_sd else expr.std(1)
        
    z_expr = expr.apply(lambda c: (c-m) /s)
    z_expr.columns = _add_level(z_expr.columns, 'z', 1)
    return pd.concat((expr, z_expr), axis=1).reindex(expr.columns.unique(0), 
                                                     axis=1, level=0)

def _add_level(index, label, at=0, replace=False, name=''):
    """Add a level with labels 'label' to a pd.MultiIndex"""
    index = pd.DataFrame(index=index)
    if replace:
        index.reset_index(level=at, drop=True, inplace=True)
    index[name] = label
    order = list(range(index.index.nlevels))
    order.insert(at, -1)
    return index.set_index(name, append=True).reorder_levels(order).index

def _get_gene_ann(species):
    """Open the gene annotation reference file (mouse/ human) and return it"""
    path = os.path.dirname(__file__)
    if species == 'mouse':
        return pd.read_pickle(path + '/../gene_ann/mg_ensembl96_GRCm38.p6.gzip')
    elif species == 'human':
        return pd.read_pickle(path + '/../gene_ann/hg_GRCh38.p12.gzip')
    else:
        logger.info('')
        logger.error('Invalid input for species: `{}`. Valid are `mouse` and '
                     '`human`'.format(species))
        sys.exit(1)

def annotate(ensgs, species):
    """ Annotate mouse or human ensg keys. Return the gene names.

    DPre references the ensembl gene annotation v.96 located at 
    DPre/gene_ann. 

    Args:
        ensgs (list, pandas.Index): The collection of ensg keys to annotate
        species (str): The origin species of the genes, 'mouse' or 'human'.
    
    Returns:
        annotated pandas.Index
    """
    ref = _get_gene_ann(species)
    try:
        return pd.Index(ref.reindex(ensgs).name.values)
    except Exception as e:
        logger.error('{}\nDPre references the ensembl gene annotaiton v.96. '
                     'Differently annotated datasets may cause problems.'
                     .format(e))
        sys.exit(1)

def get_ensgs(names, species):
    """ Return the ensg keys for a list of gene names.

    DPre references the ensembl gene annotation v.96 located at 
    DPre/gene_ann. If a gene name has multiple ensg keys, this gene will appear 
    last in the DataFrame regardless of the input order.
    
    Args:
        names (list pandas.Index): The collection of names to return ensg keys 
            for
        species (str): The origin species of the genes, 'mouse' or 'human'.
    
    Returns:
        pandas.Index of ensg keys 
    """
    ref = _get_gene_ann(species)
    try:
        ann = ref.reindex(ref.index[ref.name.isin(names)]).reset_index()
        if ann.name.duplicated().any():
            dupl = pd.Index(ann.name).duplicated()
            ann_dr = ann[~dupl]
            ann_du = ann[dupl]
            ann_dr = ann_dr.set_index('name').reindex(names).reset_index()
            ann_dr.rename({'index': 'name'}, axis=1, inplace=1)
            ann = ann_dr.append(ann_du, sort=False)
            ann.index = np.arange(ann.shape[0])
        else:
            ann = ann.set_index('name').reindex(names).reset_index()
            ann.rename({'index': 'name'}, axis=1, inplace=1)
        return ann
    except Exception as e:
        logger.error('{}\nDPre references the ensembl gene annotaiton v.96. '
                     'Differently annotated datasets may cause problems.'
                     .format(e))
        sys.exit(1)

def _align_indices(data, order, axis=1):
    """Align the indices/ columns in a collection of pandas objects to order"""
    for i in range(len(data)):
        if data[i] is not None:
            data[i] = data[i].reindex(order, axis=axis)
    return data


def _init_figure(fig_widths, fig_heights, nplts, spacers):
    """Calculate the size proportion of each plot element, create figure"""
    width, height = sum(fig_widths), sum(fig_heights)
    ratio = {'width_ratios': list(map(lambda w: w/width, 
                                        fig_widths[1:-2])),
             'height_ratios': list(map(lambda h: h/height, 
                                        fig_heights[1:-2]))}
    # init figure
    fig, axes = plt.subplots(*nplts, figsize=(width, height), 
                                gridspec_kw=ratio)
    if not isinstance(axes, np.ndarray):
        axes = np.array([axes])
    axes = _clean_axes(axes)
    
    wspace_prop = spacers[0] /np.array(fig_widths[1:-2]).mean()
    hspace_prop = spacers[1] /np.array(fig_heights[1:-2]).mean()
    adj_args = {'left': fig_widths[0] /width, 
                'wspace': wspace_prop,
                'right': 1 - fig_widths[-1] /width,
                'top': 1 - fig_heights[0] /height, 
                'hspace': hspace_prop,
                'bottom': fig_heights[-1] /height}
    fig.subplots_adjust(**adj_args)
    return fig, axes

def _open_file(filename):
    """Open a file based on the filename ending or if not present
       on config.SAVE_FORMAT. Must be supporte by matplotlib."""
    valid = plt.figure().canvas.get_supported_filetypes()
    if not any([filename.endswith(val_format) for val_format in valid]):
        if config.SAVE_FORMAT in valid:
            filename += '.' + config.SAVE_FORMAT
        else:
            logger.error('The value for config.SAVE_FORMAT `{}` is not '
                         'supported by matplotlib. Valid formats are:\n{}'
                         .format(config.SAVE_FORMAT, ', '.join(list(valid.keys()))))
            sys.exit(1)
        
    if filename.endswith('.pdf'):
        return filename, PdfPages(filename)
    else:
        return filename, None
def _save_file(fig, filename=None, pp=None, close_pp=False):
    """Save pdf if pp is passed, otherwise use filename to save as 
       config.SAVE_FORMAT"""
    if pp:
        fig.savefig(pp, format='pdf')
        if close_pp:
            pp.close()
            plt.close(fig)
    elif filename:
        replace = ['$\\mathit{', '}$']
        for repl in replace:
            filename = filename.replace(repl, '')
        fig.savefig(filename)
        plt.close(fig)

def _clean_axes(axes):
    """Remove all spines, ticks and tickalabels"""
    np.array([axes])
    for ax in axes.flatten():
        [s.set_visible(False) for s in ax.spines.values()]
        ax.tick_params(bottom=False, left=False, labelbottom=False, 
                       labelleft=False)
    return axes

def _make_title(differential, metric, el1, el2, pref='', postf=''):
    """Produce the plot title based on plot parmaters, pref and posf are used
       for plot specific adjustments; return the title string"""

    metric_title = 'metric: '
    if metric == 'euclid':
        metric_title += 'L1 Euclidean distance'
    elif metric == 'pearson':
        metric_title += 'Pearson correlation' 
    elif metric == 'cosine':
        metric_title += 'cosine similarity' 
    elif metric == 'intersect':
         metric_title += 'marker gene intersect'

    if differential:
        dtype = 'Change in '
    else:
        dtype = 'Absolute '
    if pref:
        dtype = dtype.lower()
    title = ('{}{}{}transcriptional similarity \nof {} & {}\n{}'
              .format(pref, dtype, postf, el1, el2, metric_title))
    return title[0].upper() + title[1:]

def _heatmap_cluster(dat, where, ax, metric):
    """Cluster the columns or index with scipy; return the new order"""
    ax.set_visible(True)
    d = dat.T if metric == 'columns' else dat 
    Y = pdist(d, metric='euclidean')
    Z = linkage(Y, method='complete', metric='euclidean')
    order = dendrogram(Z,
                        count_sort = True,
                        no_labels = True,
                        orientation = where, 
                        labels = d.index, 
                        above_threshold_color = config.dendrogram_colors[0],
                        ax = ax)['ivl']
    if metric == 'rows':
        # for some reason reversed?
        order = order[::-1]
    return order

def _plot_distance_bar(axes, data, ctrl_lbl, bar_args, draw_colorbar=False, 
                       cb_lbl=None, fig=None, pivot=None, w=None, h=None):
    """Draw the distance bar on top of the heatmap"""
    # set ylabel on the left 
    axes[0].tick_params(labelleft=True)
    axes[0].set_ylim(0, 1)
    axes[0].set_yticks((.5,))
    axes[0].set_yticklabels((ctrl_lbl,), x=.5)
    # draw the heatmap
    ax = axes[1]
    [s.set_visible(True) for s in ax.spines.values()]
    im = ax.imshow(data.values, aspect='auto', **bar_args)

    # setup the colorbar legend    
    if draw_colorbar:
        at = (config.CB_LEFT_SEC/w, 1- config.CB_TOP/h, config.CB_WIDTH/w, 
              config.CB_HEIGHT/h)
        cb = ax.figure.colorbar(im, cax=fig.add_axes(at), alpha =.3,
                                orientation='horizontal')
        bar_ticks = (bar_args['vmin'], bar_args['vmax'])
        cb.set_ticks(bar_ticks)
        cb.ax.set_xticklabels(bar_ticks)
        if pivot:
            cb.ax.tick_params(labelrotation=90)
        cb.ax.set_xlabel(cb_lbl)
        cb.ax.get_xaxis().set_label_position('top')

def _setup_heatmap_xy(x_y, ax, lbls, pivot, hide_lbls, lbl_size, colors):
    """Setting all paramters for the x- and y axis of the two heatmap plots"""
    dim = len(lbls)
    if x_y == 'x':
        # X-axis setup, colorbar bottom
        ax.set_xlim(0, dim)
        ticks = np.arange(.5, dim)
        ax.set_xticks(ticks)
        if not hide_lbls:
            ax.tick_params(labelbottom=True)
            fs = lbl_size*config.FONTS if lbl_size else config.FONTS
            if not pivot:
                ax.set_xticklabels(lbls, rotation=45, ha='right', fontsize=fs, 
                                   rotation_mode='anchor', y=-.5)
            else:
                ax.set_xticklabels(lbls, rotation=90, ha='right', va='center', 
                                   fontsize=fs, rotation_mode='anchor', y=-.5)
        if colors:
            ax.bar(ticks, 1, 1, color=colors)
    
    elif x_y == 'y':
        ax.set_ylim((-.1, dim +.01))
        ax.set_yticks(np.arange(.5, dim))
        if not hide_lbls:
            ax.tick_params(labelleft=True) 
            fs = lbl_size*config.FONTS if lbl_size else config.FONTS
            if not pivot:
                ax.set_yticklabels(lbls, x=.5, fontsize=fs)
            else:
                ax.set_yticklabels(lbls, rotation=45, ha='right', x=1,
                                   fontsize=fs, rotation_mode='anchor')
        if colors:
            ax.bar(0, 1, width=1, color=colors, bottom=np.arange(len(lbls)))

def _check_args(trg, smp, metric, differential, 
                hide_distance_bar=None, reorder_to_distance_bar=None,
                distance_bar_range=None, cluster_hmx=None, 
                display_markergenes=False):
    """General purpose plot argument checker; returns (modified) input values"""
    def check_metric(metric, trg, smp, diff):
        # check if the samples and targets have equivalent data to compare
        if smp._type_name != 'samples':
            logger.error('The passed `samples` are not of type DPre.smaples.')
            sys.exit(1)
        if metric is None:
            if trg._has_expr and smp._has_expr:
                metric = 'cosine'
            elif trg._has_diff and smp._has_diff:
                metric = 'intersect'
            else:
                logger.error('Either initiate targets and samples with '
                             'expression or with markergenes and diff genes.')
                sys.exit(1)
        msg = 'The {} were initiated without {} data. Cannot use `{}` similarity.'
        if metric not in ('euclid', 'cosine', 'pearson', 'intersect'):
            logger.error('Invalid `metric` input: `{}`. Valid are `euclid`, '
                         '`cosine`, `pearson`, and `intersect`'.format(metric))
            sys.exit(1)
        if metric in ['euclid', 'cosine', 'pearson']:
            if not trg._has_expr:
                logger.error(msg.format('targets', 'expression', metric))
                sys.exit(1)
            elif not smp._has_expr:
                logger.error(msg.format('samples', 'expression', metric))
                sys.exit(1)
                if diff and not smp._ctrl:
                    logger.error('To plot the changes in transcriptional similarity '
                    'with metric = `{}`, the samples must be initiated '
                    'with a control. For absolute, pass differential = ' 
                    'False.'.format(metric))
                    sys.exit(1)
        elif metric == 'intersect':
            if not trg._has_diff:
                logger.error(msg.format('targets', 'merker gene', metric))
                sys.exit(1)
            elif not smp._has_diff:
                logger.error(msg.format('samples', 'diff genes', metric))
                sys.exit(1)
        
        return metric

    # checks for all plots
    metric = check_metric(metric, trg, smp, differential)
    if metric == 'intersect' and not differential:
        differential = True
        logger.warning('For the `intersect` similarity metric, '
                       'differential cannot be False. Was set to True.')

    # checks for 2 heatmaps
    if metric != 'intersect' and not hide_distance_bar and not smp._ctrl:
        hide_distance_bar = True
        logger.warning('`hide_distance_bar` must be True '
                    'for metric = `{}` if the samples data is '
                    'initialized without a control. Set to True.'.format(metric))
    if reorder_to_distance_bar and hide_distance_bar:
        reorder_to_distance_bar = False
        logger.warning('When `reorder_to_distance_bar` is True, '
                        '`hide_distance_bar` cannot be True. Set '
                        'to False.')
    if reorder_to_distance_bar and cluster_hmx:
        cluster_hmx = False
        logger.warning('Both `reorder_to_distance_bar` and '
                        '`cluster_genes` were set as True. '
                        '`cluster_genes` will be ignored.')                    
    if not differential and distance_bar_range is not None:
        distance_bar_range = None
        logger.warning('The argument `distance_bar_range` is invalid '
                       'and ignored when differential = False. To apply'
                       ' a custom range, please use "heatmap_range".')
            
    if display_markergenes is not False:
        # checks for target_sim and ranked_sim plots
        val = ['mean', 'up', 'down']
        if display_markergenes not in val:
            logger.warning('Invalid input for display_markergenes: `{}`. ' 
                           'Valid are {}. Set to default `{}`'
                           .format(display_markergenes, val, val[0]))
            display_markergenes = val[0] 

        if display_markergenes == val[2] and not trg._down_mgs: 
            logger.error('Cannot display down markergene similarity because'
                         ' the targets were not initiated with down '
                         'markergenes.')
            sys.exit(1)

    return metric, differential, hide_distance_bar, reorder_to_distance_bar, \
           distance_bar_range, cluster_hmx, display_markergenes

def plot_color_legend(labels, colors, ncolumns=1, filename='color_legend'):
    """Plot a custom color legend.
    
       Takes a list of labels and colors and links them to produce a color 
       legend. Useful for marking  sub-groups in samples/ targets elements.

    Args:
        labels (list): the list of labels in the legend
        colors (list): the list of colors correspoding to the labels. Colors 
            must be interpretable by matplotlib: for example, 'w', #ffffff, 
            (1,1,1) all refer to white.
        filename (str, optional): the filename to save the legend. Defaults to
            './color_legend.' + config.SAVE_FORMAT
        ncolumns (int, optional): the number of columns in the legend. Defaults 
            to 1.
    """
    spacer.info('\n\n')
    assert len(colors) == len(labels), 'colors and labels differ in length'
    inv_cols = [c for c in colors if not is_color_like(c)]
    if inv_cols:
        logger.error('The following colors are not recognized as colors by '
                     'matplotlib: {}'.format(inv_cols))
        sys.exit(1)
    filename, pp = _open_file(filename)
    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
    _clean_axes(np.array([ax]))
    ax.legend(handles=[Patch(color=colors[i], label=labels[i]) 
                for i in range(len(colors))], loc='center', ncol=ncolumns)
    _save_file(fig, filename=filename, pp=pp, close_pp=True)

    logger.info('Color legend generated and saved at {}/{}'
                .format(os.path.abspath(os.curdir), filename))

def add_diff_genes_from_z(samples, diff_z_threshold=2):
    if not samples._ctrl:
        logger.error('The samples `{}` were not initialized with a control.'
                     'To generate a list of differential genes, a control is '
                     'required.'.format(samples.name))
        sys.exit(1)
    expr = samples._expr.xs('z', 1, 1, False)
    expr = expr.apply(lambda smp: smp - expr.loc(1)[(samples._ctrl, 'z')])

    up = expr.mask(~(expr>diff_z_threshold), False).astype(bool)
    up.columns = pd.MultiIndex.from_product([['up'], up.columns.unique(0)])
    down = expr.mask(~(expr<-diff_z_threshold), False).astype(bool)
    down.columns = pd.MultiIndex.from_product([['down'], down.columns.unique(0)])

    samples._diff = pd.concat((up, down), axis=1)
    samples._has_diff = True
    spacer.info('\n')
    n = samples._diff.sum().unstack(0).reindex(samples.names).to_string()
    logger.info('Differential genes were added to the sample. Number of marker '
                'genes:\n{}\n'.format(n))
Utility module with various helper functions and subplot generaters Add markergene type index (up and down) to columns at level 0 Take _diff input and convert up-genes to +1, down-genes to -1, optionally 
        transfer up- and down values to each other makeing up- and down subframes
        equal Prodcue the mean between aggregated up- and down mg similarity values Compute log2 and z-transformed expression data. Substitute read count in
        expression. Optionally, compute the standad deviation row(gene)-wise. 
        Is used for large datasets like reference transcriptome libraries. Add a level with labels 'label' to a pd.MultiIndex Open the gene annotation reference file (mouse/ human) and return it Annotate mouse or human ensg keys. Return the gene names.
 
     DPre references the ensembl gene annotation v.96 located at 
     DPre/gene_ann. 
 
     Args:
         ensgs (list, pandas.Index): The collection of ensg keys to annotate
         species (str): The origin species of the genes, 'mouse' or 'human'.
     
     Returns:
         annotated pandas.Index Return the ensg keys for a list of gene names.
 
     DPre references the ensembl gene annotation v.96 located at 
     DPre/gene_ann. If a gene name has multiple ensg keys, this gene will appear 
     last in the DataFrame regardless of the input order.
     
     Args:
         names (list pandas.Index): The collection of names to return ensg keys 
             for
         species (str): The origin species of the genes, 'mouse' or 'human'.
     
     Returns:
         pandas.Index of ensg keys Align the indices/ columns in a collection of pandas objects to order Calculate the size proportion of each plot element, create figure # init figure Open a file based on the filename ending or if not present
        on config.SAVE_FORMAT. Must be supporte by matplotlib. Save pdf if pp is passed, otherwise use filename to save as 
        config.SAVE_FORMAT Remove all spines, ticks and tickalabels Produce the plot title based on plot parmaters, pref and posf are used
        for plot specific adjustments; return the title string Cluster the columns or index with scipy; return the new order # for some reason reversed? Draw the distance bar on top of the heatmap # set ylabel on the left # draw the heatmap # setup the colorbar legend Setting all paramters for the x- and y axis of the two heatmap plots # X-axis setup, colorbar bottom General purpose plot argument checker; returns (modified) input values # check if the samples and targets have equivalent data to compare # checks for all plots # checks for 2 heatmaps # checks for target_sim and ranked_sim plots Plot a custom color legend.
     
        Takes a list of labels and colors and links them to produce a color 
        legend. Useful for marking  sub-groups in samples/ targets elements.
 
     Args:
         labels (list): the list of labels in the legend
         colors (list): the list of colors correspoding to the labels. Colors 
             must be interpretable by matplotlib: for example, 'w', #ffffff, 
             (1,1,1) all refer to white.
         filename (str, optional): the filename to save the legend. Defaults to
             './color_legend.' + config.SAVE_FORMAT
         ncolumns (int, optional): the number of columns in the legend. Defaults 
             to 1.