# plot ramp protocol and responses of each model to ramp
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.transforms import Bbox
import string

#### from https://gist.github.com/dmeliza/3251476  #####################################################################
from matplotlib.offsetbox import AnchoredOffsetbox
class AnchoredScaleBar(AnchoredOffsetbox):
    def __init__(self, transform, sizex=0, sizey=0, labelx=None, labely=None, loc=4,
                 pad=0.1, borderpad=0.1, sep=2, prop=None, barcolor="black", barwidth=None,
                 **kwargs):
        """
        Draw a horizontal and/or vertical  bar with the size in data coordinate
        of the give axes. A label will be drawn underneath (center-aligned).
        - transform : the coordinate frame (typically axes.transData)
        - sizex,sizey : width of x,y bar, in data units. 0 to omit
        - labelx,labely : labels for x,y bars; None to omit
        - loc : position in containing axes
        - pad, borderpad : padding, in fraction of the legend font size (or prop)
        - sep : separation between labels and bars in points.
        - **kwargs : additional arguments passed to base class constructor
        """
        from matplotlib.patches import Rectangle
        from matplotlib.offsetbox import AuxTransformBox, VPacker, HPacker, TextArea, DrawingArea
        bars = AuxTransformBox(transform)
        if sizex:
            bars.add_artist(Rectangle((0, 0), sizex, 0, ec=barcolor, lw=barwidth, fc="none"))
        if sizey:
            bars.add_artist(Rectangle((0, 0), 0, sizey, ec=barcolor, lw=barwidth, fc="none"))

        if sizex and labelx:
            self.xlabel = TextArea(labelx)
            bars = VPacker(children=[bars, self.xlabel], align="center", pad=0, sep=sep)
        if sizey and labely:
            self.ylabel = TextArea(labely)
            bars = HPacker(children=[self.ylabel, bars], align="center", pad=0, sep=sep)

        AnchoredOffsetbox.__init__(self, loc, pad=pad, borderpad=borderpad,
                                   child=bars, prop=prop, frameon=False, **kwargs)


def add_scalebar(ax, matchx=True, matchy=True, hidex=True, hidey=True, **kwargs):
    """ Add scalebars to axes
    Adds a set of scale bars to *ax*, matching the size to the ticks of the plot
    and optionally hiding the x and y axes
    - ax : the axis to attach ticks to
    - matchx,matchy : if True, set size of scale bars to spacing between ticks
                    if False, size should be set using sizex and sizey params
    - hidex,hidey : if True, hide x-axis and y-axis of parent
    - **kwargs : additional arguments passed to AnchoredScaleBars
    Returns created scalebar object
    """

    def f(axis):
        l = axis.get_majorticklocs()
        return len(l) > 1 and (l[1] - l[0])

    if matchx:
        kwargs['sizex'] = f(ax.xaxis)
        kwargs['labelx'] = str(kwargs['sizex'])
    if matchy:
        kwargs['sizey'] = f(ax.yaxis)
        kwargs['labely'] = str(kwargs['sizey'])

    sb = AnchoredScaleBar(ax.transData, **kwargs)
    ax.add_artist(sb)

    if hidex: ax.xaxis.set_visible(False)
    if hidey: ax.yaxis.set_visible(False)
    if hidex and hidey: ax.set_frame_on(False)

    return sb
########################################################################################################################


def plot_ramp_V(ax, model='RS Pyramidal'):  # , stop=750
    model_ramp = pd.read_csv('./Figures/Data/model_ramp.csv')
    ax.plot(model_ramp['t'], model_ramp[model], 'k', linewidth=0.25)
    ax.set_ylabel('V')
    ax.set_xlabel('Time [s]')
    ax.set_ylim(-80, 60)
    ax.axis('off')
    ax.set_title(model)

#% plot setup
fig = plt.figure(figsize=(20,10))

gs0 = fig.add_gridspec(3, 2, wspace=0.1)
gs00 = gs0[:,0].subgridspec(7, 2, wspace=0.6, hspace=1)
gs01 = gs0[:,1].subgridspec(7, 2, wspace=0.6, hspace=1)

ax1_ramp = fig.add_subplot(gs00[0,0:2])
ax2_ramp = fig.add_subplot(gs01[0,0:2])
ax3_ramp = fig.add_subplot(gs00[1,0:2])
ax4_ramp = fig.add_subplot(gs01[1,0:2])
ax5_ramp = fig.add_subplot(gs00[2, 0:2])
ax6_ramp = fig.add_subplot(gs01[2, 0:2])
ax7_ramp = fig.add_subplot(gs00[3,0:2])
ax8_ramp = fig.add_subplot(gs01[3,0:2])
ax9_ramp = fig.add_subplot(gs00[4,0:2])
ax10_ramp = fig.add_subplot(gs01[4,0:2])
ax11_ramp = fig.add_subplot(gs00[5,0:2])
ax12_ramp = fig.add_subplot(gs01[5,0:2])

ramp_axs = [ax1_ramp, ax2_ramp, ax3_ramp, ax4_ramp, ax5_ramp,ax6_ramp, ax7_ramp, ax8_ramp,
             ax9_ramp, ax10_ramp, ax11_ramp, ax12_ramp]

# order of models
models = ['Cb stellate','RS Inhibitory','FS', 'RS Pyramidal','RS Inhibitory +$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$',
          'Cb stellate +$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$', 'FS +$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$',
          'RS Pyramidal +$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$', 'STN +$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$',
          'Cb stellate $\Delta$$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$',
          'STN $\Delta$$\mathrm{K}_{\mathrm{V}}\mathrm{1.1}$', 'STN']

# plot ramps
for i in range(len(models)):
    plot_ramp_V(ramp_axs[i], model=models[i])

# add scalebar
add_scalebar(ax11_ramp, matchx=False, matchy=False, hidex=True, hidey=True, sizex=1000, sizey=25, labelx='1 s',
                 labely='25 mV', loc=3, pad=-2, borderpad=0, barwidth=2, bbox_to_anchor=Bbox.from_bounds(-0.05, 0.1, 1, 1),
                          bbox_transform=ax11_ramp.transAxes)
# add_scalebar(ax12_ramp, matchx=False, matchy=False, hidex=True, hidey=True, sizex=1000, sizey=25, labelx='1 s',
#                  labely='25 mV', loc=3, pad=-2, borderpad=0, barwidth=2, bbox_to_anchor=Bbox.from_bounds(-0.05, 0.1, 1, 1),
#                           bbox_transform=ax12_ramp.transAxes)

# add subplot labels
for i in range(0,len(models)):
    ramp_axs[i].text(-0.05, 1.08, string.ascii_uppercase[i], transform=ramp_axs[i].transAxes, size=16, weight='bold')

fig.savefig('./Figures/ramp_firing.pdf')
plt.show()