import matplotlib
from matplotlib import rcParams
matplotlib.rcParams['text.usetex'] = True
import matplotlib.pyplot as plt
import pylab as pl
from copy import copy


### Format for figures
import json
P= json.load(open("cfg/PLOSmpl.json")) # import plot parameter

matplotlib.rcParams.update([(j,k) for (j,k) in P.items()
                            if j in matplotlib.rcParams.keys()])

matplotlib.rcParams['lines.linewidth']=0.065
fig_wide=matplotlib.rcParams["figure.figsize"][0]
fig_height=matplotlib.rcParams["figure.figsize"][1]


### includes scale bar in the figures
import matplotlib.offsetbox
from matplotlib.lines import Line2D

class AnchoredHScaleBarY(matplotlib.offsetbox.AnchoredOffsetbox):
    """ size: length of bar in data units
        extent : height of bar ends in axes units """
    def __init__(self, size=1, extent = 0.03, label="", loc=2, ax=None,
                 pad=0.4, borderpad=0.5, ppad = 0, sep=2, prop=None, 
                 frameon=True, linekw={}, **kwargs):
        if not ax:
            ax = plt.gca()
        trans = ax.get_yaxis_transform()
        size_bar = matplotlib.offsetbox.AuxTransformBox(trans)
        line = Line2D([0,0],[0,size], **linekw)
        size_bar.add_artist(line)
        txt = matplotlib.offsetbox.TextArea(label, minimumdescent=False,textprops={'fontsize':matplotlib.rcParams['ytick.labelsize']})
        self.vpac = matplotlib.offsetbox.HPacker(children=[size_bar,txt],  
                                 align="center", pad=ppad, sep=sep) 
        matplotlib.offsetbox.AnchoredOffsetbox.__init__(self, loc, pad=pad, 
                 borderpad=borderpad, child=self.vpac, prop=prop, frameon=frameon,
                 **kwargs)

        
### Format for figures
#######################################################################
########################## Creating nice figure
def fancy_fig_4_stimulation(sim,title=[],check_fit=False,add_scale=[],fwidealt=0,popt=[],tad=[]):
    
    t_v=[float(i_n) for i_n in range(0,92)]
    lt=int(len(sim.a_Results.t))
    if fwidealt==0:
        f1 = plt.figure(facecolor="1",figsize=(fig_wide,fig_height*0.7))
    else:
        f1 = plt.figure(facecolor="1",figsize=(fwidealt,fig_height*0.7))
    ax1 = plt.subplot2grid((5,1), (0, 0), colspan=2,rowspan=1)
    axp = plt.subplot2grid((5,1), (1, 0), colspan=2,rowspan=4,sharex=ax1)
    ##### Plot Input
    ### Python 3 doesn't work with this...
    ax1.plot(sim.a_Results.t[::20]/1000.0, sim.c_neuron.noisy_current[::20],color="black")
    ax1.set_ylim([min(sim.c_neuron.noisy_current)-0.01,max(sim.c_neuron.noisy_current)+0.45])
    ax1.get_xaxis().set_visible(False)
    ax1.get_yaxis().set_visible(False)
    #ax2.set_ylabel("Input",color="white",labelpad=10)
    ax1.set_ylabel(r'I [uA]',labelpad=10)
    ax1.set_title('input', loc='left')
    locatory1 = MaxNLocator(nbins=2) # with 3 bins you will have 4 ticks
    ax1.yaxis.set_major_locator(locatory1)
    ax1.spines['bottom'].set_color('white')
    ax1.spines['top'].set_color('white')
    ax1.spines['left'].set_color('white')
    ax1.spines['right'].set_color('white')
    if len(add_scale)<1:
        add_scale=[0.2,'0.2nA']
        
    box1 =ax1.get_position()
    xxpos=1.12
    yypos=0.7
    ob = AnchoredHScaleBarY(ax=ax1,size=add_scale[0], label=add_scale[1],bbox_to_anchor=(xxpos,yypos), bbox_transform=ax1.transAxes, loc="upper right", frameon=False,
                            pad=0.6,sep=4, linekw=dict(color="k", linewidth=0.8))

    ax1.add_artist(ob)
    for t in ax1.xaxis.get_ticklines(): t.set_color('white')
    #### Plot complete trace
    axp.plot(sim.a_Results.t[0:lt]/1000.0,sim.a_Results.V[0:lt],color=[0,0,0])
#     ### show fit
#     sim_exp=sim
#     v_t=sim.a_Results.t
#     sp_ampl,sp_V_max,sp_V_min=extract_spike_ampl(sim_exp.a_Results.V, sim_exp.fr[2][0:-1])
#     tad, popt,tpeak,perr=doub_expon_fit_2(v_t[sim_exp.fr[2][0:-1]],sp_V_max,std_error=True)
#     mqe=mean_squared_error(sp_V_max[len(sp_V_max)-len(doub_expon_fun_2(tad,*popt)):], doub_expon_fun_2(tad,*popt))
#     # f1,axp=fancy_fig_4_stimulation(sim,title='40Hz stim '+s_Cell+' msqe='+str(mqe),check_fit=True)
#     c=0
#     while c < 4 and mqe>1.5:
#         print('Trying improvement!')
#         poptx=popt
#         # poptx[0]=popt[0]/10
#         tad, popt,tpeak,perr=doub_expon_fit_2(v_t[sim_exp.fr[2][0:-1]],sp_V_max,p=poptx,std_error=True)
#         mqe=mean_squared_error(sp_V_max[len(sp_V_max)-len(doub_expon_fun_2(tad,*popt)):], doub_expon_fun_2(tad,*popt))
#         c+=1
    
    if len(popt)!=0:
        sim=sim_exp
        sp_ampl,sp_V_max,sp_V_min=extract_spike_ampl(sim_exp.a_Results.V, sim_exp.fr[2][0:-1])
        tad, popt,tpeak,perr=doub_expon_fit_2(v_t[sim_exp.fr[2][0:-1]],sp_V_max,std_error=True)
        mqe=mean_squared_error(sp_V_max[len(sp_V_max)-len(doub_expon_fun_2(tad,*popt)):], doub_expon_fun_2(tad,*popt))
        # f1,axp=fancy_fig_4_stimulation(sim,title='40Hz stim '+s_Cell+' msqe='+str(mqe),check_fit=True)
        c=0
        while c < 4 and mqe>1.5:
            print('Trying improvement!')
            poptx=popt
            # poptx[0]=popt[0]/10
            tad, popt,tpeak,perr=doub_expon_fit_2(v_t[sim_exp.fr[2][0:-1]],sp_V_max,p=poptx,std_error=True)
            mqe=mean_squared_error(sp_V_max[len(sp_V_max)-len(doub_expon_fun_2(tad,*popt)):], doub_expon_fun_2(tad,*popt))
            c+=1
        axp.plot(tad/1000.0,doub_expon_fun_2(tad,*popt),'--',linewidth=2,color='gray')

    axp.set_ylabel(r'V (mV)',labelpad=5)
    axp.set_xlim([-10/1000.0,max(sim.a_Results.t[0:lt])/1000.0])
    axp.set_ylim([-100,80])
    axp.set_xlabel('time (sec)',labelpad=5)
    locatory11 = MaxNLocator(nbins=3) # with 3 bins you will have 4 ticks
    axp.yaxis.set_major_locator(locatory11)
    #f1.subplots_adjust(bottom=0.17)
    if title ==[]:
        pass
    else:
        f1.suptitle(title)
    if check_fit:
        return f1,axp
    else:
        return f1