import numpy as np 
import matplotlib.pyplot as plt
import matplotlib
import pandas as pd
idx = pd.IndexSlice
import pickle as pkl
import os
import spatint_utils as su

datapath = os.getcwd()+'/../data/'
with open(datapath+'lgn_spat_profile.pkl','rb') as r:
    lgn_spat_profile = pkl.load(r)
    
with open(datapath+'lgn_ori_examples.pkl','rb') as r:
    lgn_ori_examples = pkl.load(r)
    
with open(datapath + 'v1_raw_muaRFs.pkl','rb') as r:
    v1_raw_muaRFs = pkl.load(r)
    
with open(datapath + 'v1_muaRFs.pkl','rb') as r:
    v1_muaRFs = pkl.load(r)
    
with open(datapath + 'lgn_raw_muaRFs.pkl','rb') as r:
    lgn_raw_muaRFs = pkl.load(r)
    
with open(datapath + 'lgn_muaRFs.pkl','rb') as r:
    lgn_muaRFs = pkl.load(r)

def agg_overlap_bins(data, start, end, nbins, binsize, agg_op, boot_value, filtr = False):
    """computes aggregate values for bins along a given axis, aggregate operation is defined by agg_op"""
    
    bins_data = []
    print('Averaging and bootstrapping in overlapping bins')
    for bin_start in np.linspace(start, end, nbins):
        bin_data = data.loc[(data['d']>bin_start) & (data['d']<(bin_start+binsize))]
        
        if filtr != False and len(bin_data)<filtr:
            continue
        booties = []
        for i in range(1000):
            bt_bin = bin_data[boot_value][np.random.randint(0,len(bin_data),len(bin_data))]
            booties.append(np.mean(bt_bin))
        booties = np.sort(booties)
        lower = booties[24]
        upper = booties[974]
        boot_std = np.std(booties)
        bin_agg = bin_data.agg(agg_op)
        bin_agg['lower'] = lower
        bin_agg['upper'] = upper
        bin_agg['boot'] = boot_std
        bin_agg['d'] = bin_start
        bins_data.append(bin_agg)
    return pd.concat(bins_data, axis =1)


#bin_mean = agg_overlap_bins(lgn_spat_profile,0,100,31,15,agg_op='mean',boot_value='gain',filtr=5)


def plot_fold_change(ax=None, data=lgn_spat_profile):#, bin_mean=bin_mean):
    """Plot mean fold change values of all units as scatter plot"""
    
    if ax is None:
        fig, ax = plt.subplots()
    
    # Clip low and high values for better visibility
    data.loc[data['gain']>2,'gain'] = 2.35
    data.loc[data['gain']<-2,'gain'] = -2.35
    
    ax.plot(data['d'],data['gain'],c='k', alpha=0.7, linestyle='', fillstyle='none', clip_on = False, marker='o',mec='k',markersize=1)
    examples = data.reindex([('Ntsr1Cre_2015_0080',3,4027),
                             ('Ntsr1Cre_2018_0003',2,22,3),
                            ('Ntsr1Cre_2015_0080',4,3049)])
    ax.plot([-100,100],[0,0],c='k', lw=0.35)
    ax.plot(bin_mean.loc['d'][:15]+7.5,bin_mean.loc['gain'][:15], lw=3,alpha=0.5, c='#1090cfff')
    ax.plot(bin_mean.loc["d"][8:15]+7.5,bin_mean.loc['gain'][8:15], lw=3,alpha=1, c ='#1090cfff')
    ax.set(xlim=[0,60], ylim=[-2.5,2.5])
    ax.set_xlabel('Distance of dLGN RFs to V1 RFs at injection site $(^{\circ})$',labelpad=5)
    ax.set_xticks([0,20,40,60])
    ax.set_ylabel('Fold change',labelpad=0)
    ax.minorticks_off()
    ax.set_yticks(list(np.linspace(-2,2,5))+[-2.35,2.35])
    ax.set_yticklabels([str("{0:.2f}".format(round(ytick,2)))for ytick in np.geomspace(0.25,4,5)]+['< 0.25','> 4.00'])
    ax.spines['left'].set_bounds(-2,2)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    
def plot_modulation_histogram(ax = None, data = lgn_spat_profile):
    
    if ax is None:
        fig = plt.figure()
        fig.set_figheight(2.576)
        fig.set_figwidth(4.342)
        ax = fig.add_axes((0.15,0.25,0.75,0.75))
        
    
    bins = np.linspace(0,60,13)
    data = data.sort_values('d')
    data['bin'] = data.apply(lambda x: np.searchsorted(bins, x['d'], side='right')-1, axis=1)

    ratios = data.groupby('bin').apply(lambda x: pd.DataFrame({'sup' :[len(x[x['modlab']==-1])/len(x)],
                                                      'fac' :[len(x[x['modlab']== 1])/len(x)],
                                                     'snull':[len(x[(x['modlab']==0) & (x['gain']<0)])/len(x)],
                                                     'fnull' :[len(x[(x['modlab']==0) & (x['gain']>0)])/len(x)]}))
    ratios.index = ratios.index.set_names('dummy', level=1)
    ratios = ratios.reset_index("dummy").drop(columns='dummy')
    ratios['binpos'] = bins[ratios.index]+2.5


    ax.bar(ratios['binpos'],-ratios['sup'], bottom = -ratios['snull'], color = 'blue', width=4)
    ax.bar(ratios['binpos'], ratios['fac'], bottom = ratios['fnull'], color = 'orange',width = 4)
    ax.bar(ratios['binpos'], -ratios['snull'], bottom = 0, color='lightskyblue', width=4)
    ax.bar(ratios['binpos'], ratios['fnull'], bottom = 0, color = 'moccasin', width=4)
    ax.set_xlabel('Distance of dLGN RF to \n mean V1 RF at injection site ($\degree$)')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_bounds(high=0.75,low=-0.75)
    ax.plot([0,70],[0,0],c='k',lw=0.35)
    ax.set(xlim=[2.5,60],ylim=[-0.9,0.9])
    ax.set_yticks([-0.75,-0.25,0,0.25,0.75])
    ax.set_yticklabels(['75','25','','25','75'])
    ax.set_xticks([0,20,40,60])
    ax.set_ylabel('Proportion (%)')
    ax.yaxis.set_label_coords(-0.1,0.55)
    ax.annotate('enhanced', xy=(1,0.15), xytext=(1,0.15),xycoords='data',rotation='vertical',color='darkorange')
    ax.annotate('suppressed', xy=(1,-0.8), xytext=(1,-0.8),xycoords='data',rotation='vertical',color='blue')
    
def plot_ori_examples(ax = None, data = lgn_ori_examples):
    
    def compute_ori_curve(df,opto,xs):
    
        return su.sum_of_gaussians(xs,*df['tun_pars'][opto])
    
    xs = np.linspace(0,360,361)
    ctrl_curves = data.apply(compute_ori_curve, axis=1, opto=0,xs=xs)
    opto_curves = data.apply(compute_ori_curve, axis=1, opto=1,xs=xs)
    
    if ax is None:
        gridspec_kw = {'left':0.175,'right':0.95,'bottom':0.225,'top':0.95,'hspace':0.1}
        fig, axes = plt.subplots(1,3,figsize=(6.327,1.971), gridspec_kw = gridspec_kw) 
    
    for i,ax in enumerate(axes):
        ax.plot(xs, ctrl_curves.iloc[i],c='k',ms=2)
        ax.plot(xs, opto_curves.iloc[i],c='#1090cfff',ms=2)
        ax.plot(np.linspace(0,330,12),data.iloc[i]['tun_mean'][:,0],mfc='k',ms=8,ls='',marker='.',mew=0)
        ax.plot(np.linspace(0,330,12),data.iloc[i]['tun_mean'][:,1],mfc='#1090cfff',ms=8,ls='',marker='.',mew=0)
        ax.errorbar(np.linspace(0,330,12), data.iloc[i]['tun_mean'][:,0], yerr=data.iloc[i]['tun_sem'][:,0],fmt='none', ecolor='k', elinewidth=1.5)
        ax.errorbar(np.linspace(0,330,12), data.iloc[i]['tun_mean'][:,1], yerr=data.iloc[i]['tun_sem'][:,1],fmt='none', ecolor='#1090cfff', elinewidth=1.5)
        ax.plot(np.linspace(0,330,12), 12*[data.iloc[i]['tun_spon_mean'][0]],c='k',lw=0.5)
        ax.plot(np.linspace(0,330,12), 12*[data.iloc[i]['tun_spon_mean'][1]],c='#1090cfff',lw=0.5)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(True)
        ax.spines['bottom'].set_bounds(0,360)
        xticks = [0,180,360]
        ax.set_xticks(xticks)
        ax.set_yticks([]) 
        ax.set_xticklabels([])
        
    axes[0].spines['left'].set_visible(True)
    axes[0].spines['bottom'].set_visible(True)
    axes[0].set(ylim=[0,100])
    axes[0].spines['left'].set_bounds(0,80)
    axes[0].set_yticks([0,40,80])
    
    axes[1].set(ylim=[0,62.5])
    axes[1].spines['left'].set_bounds(0,50)
    axes[1].set_yticks([0,25,50])

    axes[2].set(ylim=[0,62.5])
    axes[2].spines['left'].set_bounds(0,50)
    axes[2].set_yticks([0,25,50])

    axes[0].set_ylabel('Firing rate (sp/s)')
    axes[0].yaxis.set_label_coords(-0.33,0.4)
    axes[0].set_xlabel('Direction ($\degree$)',labelpad=0)
    axes[0].set_xticks([0,180,360])
    axes[0].set_xticklabels(['0','180','360'])
    axes[0].xaxis.set_tick_params(length=1)
    
    
def plot_v1_raw_muaRFs(ax = None, data = v1_raw_muaRFs):
    
    if ax is None:
        gridspec_kw = {'left':0.1,'right':0.975,'bottom':0.15,'top':0.9}
        fig,ax = plt.subplots(3,1,gridspec_kw=gridspec_kw)
        fig.set_figheight(1.984)
        fig.set_figwidth(1.332)

    cmap = matplotlib.cm.get_cmap('Greens')
    cmap = matplotlib.colors.LinearSegmentedColormap.from_list('custom',
                                                                [(0,'white'),
                                                                (1,'green')])
    for axi, (index, row) in zip(ax,data.iterrows()):

        left = row['ti_axes'][:,0].min()-2.4
        right = row['ti_axes'][:,0].max()+2.5
        bottom = row['ti_axes'][:,1].min()-2.5
        top = row['ti_axes'][:,1].max()+2.5
        extent = (left,right,bottom,top)

        axi.imshow(np.rot90(row['rfs'].mean(axis=(0,3))),cmap='gray', extent = extent)
        plot_ellipse(row,axi,cmap, linewidth=1.5)
        axi.axis('scaled')
        axi.set(xlim=[left,right],ylim=[bottom,top])
        axi.set_xticks([])
        axi.set_yticks([])

    else:
        xticks = [left,right]
        yticks = [bottom,top]
        axi.set_xticks(xticks)
        axi.set_yticks([])
        ax[0].set_yticks(yticks)
        axi.set_xticklabels([str("{0:.0f}".format(round(xtick,0))) +'$\degree$' for xtick in xticks])
        ax[0].set_yticklabels([str("{0:.0f}".format(round(ytick,0))) +'$\degree$' for ytick in yticks],
                            rotation=90)
        plt.setp(ax[0].yaxis.get_majorticklabels(), va='center')
        plt.setp(axi.xaxis.get_majorticklabels(), ha='center')
        axi.tick_params(length=2)
        ax[0].tick_params(axis='y', pad=0, length=2)
        axi.tick_params(axis='x', pad=1)
        
def plot_v1_rf_ellipses(ax = None, data = v1_muaRFs):
    
    
    if ax is None:
        fig = plt.figure(figsize=(7,5))
        ax = fig.add_axes((0.15,0.15,0.7,0.7))
    
    cmap = matplotlib.cm.get_cmap('Greens')

    data = data.reset_index('ch')
    
    one_series = data.loc['Ntsr1Cre_2015_0080',2]
    one_series = one_series[one_series['sigma_x_mix']<15]
    maxchan,minchan = one_series['ch'].max(), one_series['ch'].min()
    V1_x = data.groupby('m').agg('mean').loc['Ntsr1Cre_2015_0080']['x_mix']
    V1_y = data.groupby('m').agg('mean').loc['Ntsr1Cre_2015_0080']['y_mix']

    one_series.apply(lambda x: plot_ellipse(x,ax=ax,cmap=cmap,linewidth=1,maxchan=maxchan),axis=1)

    ax.scatter(V1_x,V1_y,marker='+',s=100,lw=2,c='k',zorder=3)

    ax.set(xlim=[-20,80],ylim=[-35,50])
    ax.set_xlabel('Azimuth ($\degree$)'
                 ,labelpad=0)
    ax.set_ylabel('Elevation ($\degree$)'
                 ,labelpad=0)

    xticks = [tick.get_text() for tick in ax.get_xticklabels()]

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    
def plot_lgn_raw_muaRFs(ax = None, data = lgn_raw_muaRFs):
    
    if ax is None:
        gridspec_kw = {'left':0.1,'right':0.975,'bottom':0.1,'top':0.93}
        fig,ax = plt.subplots(5,2,gridspec_kw=gridspec_kw)
        fig.set_figheight(3.552)
        fig.set_figwidth(2.190)
        ax_column1 = ax[:,0]
        ax_column2 = ax[:,1]
    
    cmap = matplotlib.colors.LinearSegmentedColormap.from_list('custom',
                                                            [(0,'xkcd:off white'),
                                                            (1,'xkcd:neon purple')])
    cmap2 = matplotlib.colors.LinearSegmentedColormap.from_list('custom',
                                                               [(0,'xkcd:off white'),
                                                               (1,'xkcd:saffron')]) #deep tea
    cols = {5:'#00555a',9:'blue'}

    ser1 = data.loc['Ntsr1Cre_2015_0080',9].iloc[14:19]
    ser2 = data.loc['Ntsr1Cre_2015_0080',5].iloc[7:12]

    for axi, (index, row) in zip(ax_column1,ser1.iterrows()):

        left = row['ti_axes'][:,0].min()-2.4
        right = row['ti_axes'][:,0].max()+2.5
        bottom = row['ti_axes'][:,1].min()-2.5
        top = row['ti_axes'][:,1].max()+2.5
        extent = (left,right,bottom,top)

        axi.imshow(np.rot90(row['rfs'].mean(axis=(0,3))),cmap='gray', extent = extent)
        plot_ellipse(row,axi,cmap, linewidth =1.5)
        axi.axis('scaled')
        axi.set(xlim=[left,right],ylim=[bottom,top])
        axi.set_xticks([])
        axi.set_yticks([])

    else:
        xticks = [left,right]
        yticks = [bottom,top]
        axi.set_xticks(xticks)
        axi.set_yticks([])
        ax_column1[0].set_yticks(yticks)
        axi.set_xticklabels([str("{0:.0f}".format(round(xtick,0))) +'$\degree$' for xtick in xticks])
        ax_column1[0].set_yticklabels([str("{0:.0f}".format(round(ytick,0))) +'$\degree$' for ytick in yticks],
                            rotation=90)
        plt.setp(ax_column1[0].yaxis.get_majorticklabels(), va='center')
        plt.setp(axi.xaxis.get_majorticklabels(), ha='center')
        ax_column1[0].tick_params(axis='y', pad=0)

    for axi, (index, row) in zip(ax_column2,ser2.iterrows()):

        left = row['ti_axes'][:,0].min()-2.4
        right = row['ti_axes'][:,0].max()+2.5
        bottom = row['ti_axes'][:,1].min()-2.5
        top = row['ti_axes'][:,1].max()+2.5
        extent = (left,right,bottom,top)

        axi.imshow(np.rot90(row['rfs'].mean(axis=(0,3))),cmap='gray', extent = extent)
        plot_ellipse(row,axi,cmap2, linewidth=1.5)
        axi.axis('scaled')
        axi.set(xlim=[left,right],ylim=[bottom,top])

        axi.set_xticks([])
        axi.set_yticks([])

    else:
        xticks = [left,right]
        yticks = [bottom,top]
        axi.set_xticks(xticks)
        axi.set_yticks([])
        ax_column2[0].set_yticks(yticks)
        axi.set_xticklabels([str("{0:.0f}".format(round(xtick,0))) +'$\degree$' for xtick in xticks]
            )
        ax_column2[0].set_yticklabels([str("{0:.0f}".format(round(ytick,0))) +'$\degree$' for ytick in yticks],
                            rotation=90)
        plt.setp(ax_column2[0].yaxis.get_majorticklabels(), va='center')
        plt.setp(axi.xaxis.get_majorticklabels(), ha='center')
        ax_column2[0].tick_params(axis='y', pad=0)
        
def plot_lgn_rf_ellipses(ax = None, data = lgn_muaRFs, v1data=v1_muaRFs):
    
    if ax is None:
        fig = plt.figure(figsize=(2.995,3.531))
        ax = fig.add_axes((0.275,0.2,0.65,0.725))

    cmap1 = matplotlib.colors.LinearSegmentedColormap.from_list('custom',
                                                                [(0,'xkcd:neon purple'),
                                                                (1,'xkcd:neon purple')])
    cmap2 = matplotlib.colors.LinearSegmentedColormap.from_list('custom',
                                                               [(0,'xkcd:saffron'),
                                                               (1,'xkcd:saffron')]) 
    cmaps = {9:cmap1,5:cmap2}
    two_series = data.loc[idx['Ntsr1Cre_2015_0080',(5,9),:]]
    V1_x = v1data.groupby('m').agg('mean').loc['Ntsr1Cre_2015_0080']['x_mix']
    V1_y = v1data.groupby('m').agg('mean').loc['Ntsr1Cre_2015_0080']['y_mix']


    two_series.groupby('s').apply(lambda x: plot_RF_ellipses(x, ax, cmaps[x.index.unique(level=1)[0]],linewidth=0.5))


    ax.scatter(V1_x,V1_y,marker='+',s=100,lw=2,c='k',zorder=3)
    ax.set(xlim=[-5,70],ylim=[-25,60])
    ax.set_xlabel('Azimuth ($\degree$)'
                 ,labelpad=0)
    ax.set_ylabel('Elevation ($\degree$)'
                 ,labelpad=0)
    xticks = [tick.get_text() for tick in ax.get_xticklabels()]
    xticks = [0,40,80]
    yticks=[-20,0,30,60]
    xticklabels = [str(xtick) for xtick in xticks]
    yticklabels = [str(ytick) for ytick in yticks]

    ax.yaxis.set_tick_params(pad=0)

    ax.set_xticks(xticks)
    ax.set_xticklabels(xticklabels)
    ax.set_yticks(yticks)
    ax.set_yticklabels(yticklabels)

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_bounds(-20,60)
    ax.spines['bottom'].set_bounds(0,80)

    norm= matplotlib.colors.Normalize(vmin=0.5,vmax=1)
        
def plot_RF_ellipses(data, ax, cmap,linewidth=1):
    """Plotting RF ellipses from fitted parameters."""
    data.apply(lambda x: plot_ellipse(x,ax,cmap,linewidth=linewidth),axis=1)
    
def plot_ellipse(data, ax, cmap, linewidth=1, maxchan=60):
    """Plot from single df row"""
    ddiff = su.degdiff(180,(data['theta_mix']*180/np.pi),180)
    params = data[['x_mix',
                   'y_mix',
                   'sigma_x_mix',
                   'sigma_y_mix']]
    x_ellipse, y_ellipse = su.calculate_ellipse(*params,ddiff)
    col = cmap(data['rsq_mix'])
    ax.plot(x_ellipse,y_ellipse, lw=linewidth,c=col)
    ax.axis('scaled')