import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import string
from plotstyle import plot_style

def plot_sqrt(ax, a=1, b=0.2, c=100, d=0):
    x = np.linspace(0, 1, 10000)
    y = c*np.sqrt(a*(x - b)) + d
    ax.plot(x, y)
    ax.set_xlabel('Current [nA]')
    ax.set_ylabel('Frequency [Hz]')
    ax.set_xlim(0,1)
    ax.set_ylim(0, ax.get_ylim()[1])

def plot_AUC(ax, a=1, b=0.2, c=100, d=0, width=0.2):
    x = np.linspace(0, 1, 10000)
    y = c*np.sqrt(a*(x - b)) + d
    ax.plot(x, y, 'k')
    ax.set_xlabel('Current [nA]')
    ax.set_ylabel('Frequency [Hz]')
    ax.fill_between(x,y, where=(x<=b+width), color='c')
    ax.set_xlim(0,1)
    ax.set_ylim(0, ax.get_ylim()[1])

def plot_diff_sqrt(ax, a=1, b=0.2, c=100, d=0, a2=1, b2=0.2, c2=100, d2=0):
    x = np.linspace(0, 1, 10000)
    y = c*np.sqrt(a*(x - b)) + d
    y2 = c2 * np.sqrt(a2 * (x - b2)) + d2
    ax.plot(x, y, 'k')
    ax.plot(x, y2, 'b')
    ax.set_xlim(0,1)
    ax.set_ylim(0, ax.get_ylim()[1])
    ax.set_xticks([])
    ax.set_yticks([])

def plot_quadrant(ax):
    ax.plot(ax.get_xlim(),[0,0], '--k', linewidth=1.25)
    ax.plot([0,0],ax.get_ylim(), '--k', linewidth=1.25)
    ax.tick_params(axis=u'both', which=u'both', length=0)
    ax.set_xticks([-0.55, 0, 0.55])
    a = ax.get_xticks().tolist()
    a[0] = '\u2212'
    a[1] = 0
    a[2] = '+'
    ax.set_xticklabels(a)
    ax.set_yticks([-0.55, 0, 0.55])
    b = ax.get_xticks().tolist()
    b[0] = '\u2212'
    b[1] = 0
    b[2] = '+'
    ax.set_yticklabels(b)

    ax.tick_params(labelsize=16)
    # ax.set_yticks([0])
    # ax.axis('off')


plot_style()

fig = plt.figure(figsize=(4, 6))
gs = gridspec.GridSpec(3,2, top=0.95, bottom=0.1, left=0.15, right = 0.95, hspace=0.8, wspace=0.6)
ax1 = fig.add_subplot(gs[0,:])
plot_AUC(ax1, width=0.2)
ax3 = fig.add_subplot(gs[1:, :])

# add panel letter labels
ax_list = fig.axes
i =0
for a in ax_list:
    a.text(-0.18, 1.08, string.ascii_uppercase[i], transform=a.transAxes,size=16, weight='bold')
    i += 1

ax3.set_ylabel('$\Delta$ AUC')
ax3.set_xlabel('$\Delta$ rheobase')
ax3.set_xlim(-1, 1)
ax3.set_ylim(-1, 1)
plot_quadrant(ax3) # plot delineation into quadrants


inset_ylim = (0, 100)
# top left
ax3.text(x=-0.75, y=0.25, s=r'''$\uparrow$ AUC
$\downarrow$ rheobase''')
ax3_TL = ax3.inset_axes([0, 0.8, 0.3, 0.2])
plot_diff_sqrt(ax3_TL, b2=0.1, c2=200)
ax3_TL.set_ylim(inset_ylim)
# top right
ax3.text(x=0.25, y=0.25, s=r'''$\uparrow$ AUC
$\uparrow$ rheobase''')
ax3_TR = ax3.inset_axes([0.7, 0.8, 0.3, 0.2])
plot_diff_sqrt(ax3_TR, b2=0.4, c2=200)
ax3_TR.set_ylim(inset_ylim)
# bottom left
ax3.text(x=-0.75, y=-0.35, s=r'''$\downarrow$ AUC
$\downarrow$ rheobase''')
ax3_BL = ax3.inset_axes([0, 0, 0.3, 0.2])
plot_diff_sqrt(ax3_BL, b2=0.1, c2=75)
ax3_BL.set_ylim(inset_ylim)
# bottom right
ax3.text(x=0.25, y=-0.35, s=r'''$\downarrow$ AUC
$\uparrow$ rheobase''')
ax3_BR = ax3.inset_axes([0.7, 0, 0.3, 0.2])
plot_diff_sqrt(ax3_BR, b2=0.4, c2=75)
ax3_BR.set_ylim(inset_ylim)
fig.savefig('./Figures/firing_characterization.pdf', bbox_inches='tight')
plt.show()



#%% OLD
# fig = plt.figure()
# # fig, axs = plt.subplots(2, 3, figsize=(6, 4))
# # fig.subplots_adjust(top=0.95, bottom=0.1, left=0.1, right=0.95, hspace=0.6, wspace=0.6)
# gs = gridspec.GridSpec(2,3, top=0.95, bottom=0.1, left=0.1, right = 0.95, hspace=0.6, wspace=0.6)
# ax1 = fig.add_subplot(gs[0,0])
# plot_AUC(ax1, width=0.2)
# ax2 = fig.add_subplot(gs[1, 0])
# ax3 = fig.add_subplot(gs[:, 1:])
# # ax_list = fig.axes
# # import string
# # i =0
# # for a in ax_list:
# #     a.text(-0.13, 1.05, string.ascii_uppercase[i], transform=a.transAxes,size=12, weight='bold')
# #     i += 1
# ax3.set_ylabel('$\Delta$ AUC')
# ax3.set_xlabel('$\Delta$ rheobase')
# ax3.set_xlim(-1,1)
# ax3.set_ylim(-1,1)
# plot_quadrant(ax3)
#
# # top left
# ax3.text(x=-0.75, y=0.25, s=r'''$\uparrow$ AUC
# $\downarrow$ rheobase''')
# ax3_TL = ax3.inset_axes([0,0.8,0.3,0.2])
# plot_diff_sqrt(ax3_TL, b2=0.1, c2=200)
#
# # top right
# ax3.text(x=0.25, y=0.25, s=r'''$\uparrow$ AUC
# $\uparrow$ rheobase''')
# ax3_TR = ax3.inset_axes([0.7, 0.8, 0.3, 0.2])
# plot_diff_sqrt(ax3_TR, b2=0.4, c2=200)
#
# # bottom left
# ax3.text(x=-0.75, y=-0.35, s=r'''$\downarrow$ AUC
# $\downarrow$ rheobase''')
# ax3_BL = ax3.inset_axes([0, 0, 0.3, 0.2])
# plot_diff_sqrt(ax3_BL, b2=0.1, c2=75)
#
# # bottom right
# ax3.text(x=0.25, y=-0.35, s=r'''$\downarrow$ AUC
# $\uparrow$ rheobase''')
# ax3_BR = ax3.inset_axes([0.7, 0, 0.3, 0.2])
# plot_diff_sqrt(ax3_BR, b2=0.4, c2=75)
#
# # ax3_TL.sharey(ax3_TR)
#
# # ax3_TL.get_shared_x_axes().join(ax3_TL, ax3_TR, ax3_BL, ax3_BR)
# # ax3_TL.get_shared_y_axes().join(ax3_TL, ax3_TR, ax3_BL, ax3_BR)
# ax_list = [ax3_TL, ax3_TR, ax3_BL, ax3_BR] #< your axes objects
# ax_list[0].get_shared_x_axes().join(*ax_list)
# plt.show()
#
#     # plot_sqrt(axs[0,0])
#     # plot_temperature(axs[0,1])
#     # plot_signal_n_power(axs[1,0], axs[1,1])
#     # fig.savefig('coolresult.pdf')
#