import string
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.transforms import BboxTransformTo

def hide_ticks(ax, side='bottom', ticks=True):
    axis = 'x' if side in ['top', 'bottom'] else 'y'
    params = {side: ticks, 'label' + side: False}
    ax.tick_params(axis=axis, which='both', **params)
    return None

def hide_axis(ax, side='bottom'):
    ax.spines[side].set_visible(False)
    params = {side: False, 'label' + side: False}
    ax.tick_params(axis='x' if side in ['top', 'bottom'] else 'y',
                   which='both', **params)
    return None

def get_trans_artist(artist):
    artist_type = type(artist).__name__
    if artist_type == 'Axes':
        return artist.transAxes
    elif artist_type == 'Figure':
        return artist.transFigure
    elif artist_type == 'Subfigure':
        return artist.transSubfigure
    elif hasattr(artist, 'bbox'):
        return BboxTransformTo(artist.bbox)
    renderer = artist.get_figure(root=True).canvas.get_renderer()
    if hasattr(artist, 'get_window_extent'):
        return BboxTransformTo(artist.get_window_extent(renderer))
    elif hasattr(artist, 'get_tightbbox'):
        return BboxTransformTo(artist.get_tightbbox(renderer))
    raise ValueError('Artist does not have a bounding box to use as transform.')

def title_subplot(artist, title, x=0.5, y=1.0, xref=None, yref=None, ref=None,
                  ha='center', va='bottom', fontsize=16, fontweight='normal', **kwargs):
    trans_artist = get_trans_artist(artist)
    if xref is not None or yref is not None:
        transform = get_trans_artist(ref) + trans_artist.inverted()
        if xref is not None:
            x = transform.transform((xref, 0))[0]
        if yref is not None:
            y = transform.transform((0, yref))[1]
    return artist.text(x, y, title, transform=trans_artist, ha=ha, va=va,
                       fontsize=fontsize, fontweight=fontweight, **kwargs)

def letter_subplot(artist, label, x=None, y=None, xref=None, yref=None, ref=None,
                   ha='left', va='bottom', fontsize=16, fontweight='bold', **kwargs):
    trans_artist = get_trans_artist(artist)
    if x is None or y is None:
        transform = get_trans_artist(ref) + trans_artist.inverted()
        if x is None:
            x = transform.transform([xref, 0])[0]
        if y is None:
            y = transform.transform([0, yref])[1]
    return artist.text(x, y, label, transform=trans_artist, ha=ha, va=va,
                       fontsize=fontsize, fontweight=fontweight, **kwargs)

def letter_subplots(artists, labels=None, x=None, y=None, xref=None, yref=None, ref=None,
                    ha='left', va='bottom', fontsize=16, fontweight='bold', **kwargs):
    if labels is None:
        labels = string.ascii_lowercase
    handles = []
    for artist, label in zip(artists, labels):
        handles.append(letter_subplot(artist, label, x, y, xref, yref, ref,
                       ha=ha, va=va, fontsize=fontsize, fontweight=fontweight, **kwargs))
    return handles

def xlimits(time, ax=None, minval=None, maxval=None, pad=0.05):
    limits = [minval, maxval]
    if minval is None:
        limits[0] = time[0]
    if maxval is None:
        limits[1] = time[-1]
    span = limits[1] - limits[0]
    if pad and minval is None:
        limits[0] -= span * pad
    if pad and maxval is None:
        limits[1] += span * pad
    if ax is not None:
        return ax.set_xlim(limits)
    return limits

def ylimits(signal, ax=None, minval=None, maxval=None, pad=0.05):
    limits = [minval, maxval]
    if minval is None:
        limits[0] = signal.min()
    if maxval is None:
        limits[1] = signal.max()
    span = limits[1] - limits[0]
    if pad and minval is None:
        limits[0] -= span * pad
    if pad and maxval is None:
        limits[1] += span * pad
    if ax is not None:
        return ax.set_ylim(limits)
    return limits

def xlabel(ax, label, x=None, y=-0.1, fontsize=20, transform=None, **kwargs):
    if x is None:
        x = 0.5  
        if transform is not None:
            x = (ax.transAxes + transform.inverted()).transform((x, 0))[0]
    ax.xaxis.set_label_coords(x, y, transform=transform)
    return ax.set_xlabel(label, fontsize=fontsize, **kwargs)

def ylabel(ax, label, x=-0.2, y=None, fontsize=20, transform=None, **kwargs):
    if y is None:
        y = 0.5
        if transform is not None:
            y = (ax.transAxes + transform.inverted()).transform((0, y))[1]
    ax.yaxis.set_label_coords(x, y, transform=transform)
    return ax.set_ylabel(label, fontsize=fontsize, **kwargs)

def super_xlabel(label, fig, left_ax, right_ax, y=0.005,
                 left_fig=None, right_fig=None, **kwargs):
    left_x = left_ax.get_position().x0
    right_x = right_ax.get_position().x1
    if left_fig is not None or right_fig is not None:
        trans_fig = get_trans_artist(fig)
        if left_fig is not None:
            transform = get_trans_artist(left_fig) + trans_fig.inverted()
            left_x = transform.transform((left_x, 0))[0]
        if right_fig is not None:
            transform = get_trans_artist(right_fig) + trans_fig.inverted()
            right_x = transform.transform((right_x, 0))[0]
    return fig.supxlabel(label, x=(left_x + right_x) / 2, y=y, **kwargs)

def super_ylabel(label, fig, low_ax, high_ax, x=0.005,
                 high_fig=None, low_fig=None, **kwargs):
    low_y = high_ax.get_position().y0
    high_y = low_ax.get_position().y1
    if low_fig is not None or high_fig is not None:
        trans_fig = get_trans_artist(fig)
        if low_fig is not None:
            transform = get_trans_artist(low_fig) + trans_fig.inverted()
            low_y = transform.transform((0, low_y))[1]
        if high_fig is not None:
            transform = get_trans_artist(high_fig) + trans_fig.inverted()
            high_y = transform.transform((0, high_y))[1]
    return fig.supylabel(label, x=x, y=(low_y + high_y) / 2, **kwargs)

def plot_line(ax, time, signal, ymin=None, ymax=None, xmin=None, xmax=None,
              xpad=None, ypad=0.05, yloc=None, xloc=None, **kwargs):
    handles = ax.plot(time, signal, **kwargs)
    xlimits(time, ax=ax, minval=xmin, maxval=xmax, pad=xpad)
    ylimits(signal, ax=ax, minval=ymin, maxval=ymax, pad=ypad)
    if xloc is not None:
        ax.xaxis.set_major_locator(plt.MultipleLocator(xloc))
    if yloc is not None:
        ax.yaxis.set_major_locator(plt.MultipleLocator(yloc))
    return handles

def plot_barcode(ax, time, binary, offset=0.5, xmin=None, xmax=None, **kwargs):
    if binary.ndim == 1:
        binary = binary[:, None]
    lower, upper, handles = 0, 1, []
    for i in range(binary.shape[1]):
        h = ax.fill_between(time, lower, upper, where=binary[:, i], **kwargs)
        handles.append(h)
        if i < binary.shape[1] - 1:
            lower += offset + 1
            upper += offset + 1
    xlimits(time, ax=ax, minval=xmin, maxval=xmax, pad=0)
    ax.set_ylim(0, upper)
    hide_axis(ax, 'bottom')
    hide_axis(ax, 'left')
    return handles

def indicate_zoom(fig, high_ax, low_ax, zoom_abs, **kwargs):
    y0 = low_ax.get_position().y0
    y1 = high_ax.get_position().y1
    transform = low_ax.transData + fig.transFigure.inverted()
    x0 = transform.transform((zoom_abs[0], 0))[0]
    x1 = transform.transform((zoom_abs[1], 0))[0]
    return fig.add_artist(plt.Rectangle((x0, y0), x1 - x0, y1 - y0,
                                        transform=fig.transFigure, **kwargs))

def assign_colors(handles, types, colors):
    for handle, type_id in zip(handles, types):
        handle.set_color(colors[str(int(type_id))])
    return None

def reorder_traces(handles, signal, zlow=2, zhigh=2.5):
    inds = np.argsort(signal.std(axis=0))
    zorders = np.linspace(zlow, zhigh, len(inds))[::-1]
    for ind, z in zip(inds, zorders):
        handles[ind].set_zorder(z)
    return None

def strip_zeros(num, right_digits=5):
    if isinstance(num, int):
        return num
    num = f'{num:.{right_digits}f}'
    left, right = num.split('.')
    right = right.rstrip('0')
    if right:
        return f'{left}.{right}'
    return left

def time_bar(ax, dur, y0=0.9, y1=0.95, xshift=0.5, parent=None,
             text_pos=None, text_str=None, text_kwargs={}, **kwargs):
    if parent is None:
        parent = ax
    trans_parent = get_trans_artist(parent)
    transform = ax.transData + trans_parent.inverted()
    t0 = ax.get_xlim()[0]
    x0 = transform.transform((t0, 0))[0]
    x1 = transform.transform((t0 + dur, 0))[0]
    dur = x1 - x0
    x0 = (1 - dur) * xshift
    rect = parent.add_artist(plt.Rectangle((x0, y0), dur, y1 - y0,
                                           transform=trans_parent, **kwargs))
    if text_pos is not None:
        trans_bar = get_trans_artist(rect)
        text_pos = (trans_bar + trans_parent.inverted()).transform(text_pos)
        if text_str is None:
            text_str = f'{dur:.2f} s'
        t = parent.text(*text_pos, text_str, transform=trans_parent, **text_kwargs)
        return rect, t
    return rect

def zoom_inset(ax, inset, handle, x0=None, x1=None, y0=None, y1=None, ref='x',
               transform=None, 
               low_left=False, up_left=False, low_right=False, up_right=False,
               props=['c', 'lw', 'ls', 'zorder', 'alpha'], **kwargs):
    if not kwargs:
        kwargs = dict(edgecolor='k', alpha=1, lw=2)

    if transform is not None:
        transform = transform + ax.transData.inverted()
    xlims = ax.get_xlim()
    ylims = ax.get_ylim()  
    if x0 is None:
        x0 = xlims[0]
    elif transform is not None:
        x0 = transform.transform((x0, 0))[0]
    if x1 is None:
        x1 = xlims[1]
    elif transform is not None:
        x1 = transform.transform((x1, 0))[0]
    if y0 is None:
        y0 = ylims[0]
    elif transform is not None:
        y0 = transform.transform((0, y0))[1]
    if y1 is None:
        y1 = ylims[1]
    elif transform is not None:
        y1 = transform.transform((0, y1))[1]
    inset.set_xlim(x0, x1)
    inset.set_ylim(y0, y1)

    x = handle.get_xdata()
    y = handle.get_ydata()
    if ref == 'x':
        zoom_inds = (x >= x0) & (x <= x1)
    elif ref == 'y':
        zoom_inds = (y >= y0) & (y <= y1)
    x = x[zoom_inds]
    y = y[zoom_inds]

    inset_handle = inset.plot(x, y)[0]
    inset_handle.set(**{prop: plt.getp(handle, prop) for prop in props})
    elements = ax.indicate_inset_zoom(inset, **kwargs)

    visibility = low_left, up_left, low_right, up_right
    [l.set_visible(v) for l, v in zip(elements.connectors, visibility)]
    return inset_handle, elements.rectangle, elements.connectors