import nixio as nix
import os
from IPython import embed
#from utility import *
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import matplotlib.mlab as ml
import scipy.integrate as si
from scipy.ndimage import gaussian_filter
from IPython import embed
from myfunctions import *
#from axes import label_axes, labelaxes_params
from myfunctions import auto_rows
from myfunctions import default_settings
from myfunctions import remove_tick_marks
from myfunctions import remove_tick_ymarks
import matplotlib.gridspec as gridspec
from rotated import ps_df

def plot_single_cells(ax, data = ['2019-10-21-aa-invivo-1','2019-11-18-af-invivo-1','2019-10-28-aj-invivo-1']):
    colors = ['#BA2D22', '#F47F17', '#AAB71B', '#3673A4', '#53379B']
    # labelaxes_params(xoffs=-3, yoffs=0, labels='A', font=dict(fontweight='bold'))
    #baseline = pd.read_pickle('data_baseline.pkl')
    #data_beat = pd.read_pickle('data_beat.pkl')
    data_all = pd.read_pickle('beat_results_smoothed.pkl')
    #data = np.unique(data_all['dataset'])[0]
    #data = np.unique(data_all['dataset'])
    end = ['original', '005','05', '2' ]

    end = ['original','005']
    y_sum = [[]]*len(data)*len(end)
    counter = 0
    for dd,set in enumerate(data):
        for ee, e in enumerate(end):
            d = data_all[data_all['dataset'] == set]
            #x = d['delta_f'] / d['eodf'] + 1
            #embed()
            #y = d['result_frequency_' + e]
            y = d['result_amplitude_max_' + e]
            #y2 = d['result_amplitude_max_' + e]
            y_sum[counter] = np.nanmax(y)
            counter += 1
            #print(np.nanmax(y))
    #embed()
    lim = np.max(y_sum)
    hd = 0.3
    ws = 0.35
    rows = 1
    cols = 3
    #embed()
    main_grid = gridspec.GridSpec(1, 2,hspace=0.4, width_ratios = [1,7])
    filter_grid = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=main_grid[0], hspace=0.4)
    upper_filter_g = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=filter_grid[0], wspace=ws, hspace=0.4)  # ,
    grid1 = gridspec.GridSpecFromSubplotSpec(rows,cols,subplot_spec=upper_filter_g[0], wspace=ws, hspace=0.4)#,
    wish_df = 150
    fc = 'lightgrey'
    ec = 'grey'
    sampling_rate = 40000
    data_beat = pd.read_pickle('data_beat.pkl')
    df, p, f, db = ps_df(data_beat, d=set, wish_df=wish_df, window='no', sampling_rate=sampling_rate)
    ax = {}
    ax_nr = 0
    colors = ['brown']
    #embed()
    ax[ax_nr] = plt.subplot(grid1[0])
    ax[ax_nr].spines['right'].set_visible(False)
    ax[ax_nr].spines['left'].set_visible(False)
    ax[ax_nr].spines['top'].set_visible(False)
    ax[ax_nr].spines['bottom'].set_visible(False)
    ax[ax_nr].set_yticks([])
    ax[ax_nr].set_xticks([])
    ax[ax_nr].plot(p, f, color=colors[0])
    eodf = d['eodf'].iloc[0]
    mult = 1.1
    ax[ax_nr].fill_between([np.min(p), np.max(p)], [eodf / 2, eodf / 2], color=fc, edgecolor=ec)
    ax[ax_nr].set_ylim([0, eodf *0.6])
    ax[ax_nr].set_xlim(ax[ax_nr].get_xlim()[::-1])
    upper_filter_g = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=filter_grid[1], wspace=ws, hspace=0.4)  # ,
    grid1 = gridspec.GridSpecFromSubplotSpec(rows,cols,subplot_spec=upper_filter_g[0], wspace=ws, hspace=0.4)#,
    i = 0
    sigma = 0.0005
    df, p, f, db = ps_df(data_beat, d=set, wish_df=wish_df, window=sigma * sampling_rate,
                                                     sampling_rate=sampling_rate)
    ax_nr = 1
    sigmaf = 1 / (2 * np.pi * sigma)
    gauss = np.exp(-(f ** 2 / (2 * sigmaf ** 2)))
    stepsize = np.abs(f[0]-f[1])
    wide = 2
    scale = 1
    prev_height = np.max((p[int(abs(df) / stepsize) - wide:int(abs(df) / stepsize) + wide]) * scale)
    now_height = np.max((p[int(abs(df) / stepsize) - wide:int(abs(df) / stepsize) + wide]) *scale)

    ax[ax_nr] = plt.subplot(grid1[1])
    ax[ax_nr].spines['right'].set_visible(False)
    ax[ax_nr].spines['left'].set_visible(False)
    ax[ax_nr].spines['top'].set_visible(False)
    ax[ax_nr].set_yticks([])
    ax[ax_nr].set_xticks([])
    ax[ax_nr].spines['bottom'].set_visible(False)
    ax[ax_nr].fill_between(max(p) * gauss ** 2, f, facecolor=fc, edgecolor=ec)
    ax[ax_nr].plot([prev_height, now_height+440],[np.abs(df), np.abs(df)], color = 'black')
    ax[ax_nr].set_xlim(ax[ax_nr].get_xlim()[::-1])
    grid = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=main_grid[1], hspace=0.4)


    #grid = gridspec.GridSpecFromSubplotSpec(2,1,subplot_spec=grid[1], hspace = 0.4)

    grid1 = gridspec.GridSpecFromSubplotSpec(rows,cols,subplot_spec=grid[0], wspace=ws, hspace=0.4)#,


    end = ['original']

    color_modul = 'steelblue'
    color_mpf = 'red'

    y_sum1 = plot_single(lim, data, end, data_all, grid1, color_mpf, color_modul,title = True,xlabel = False, label =False,remove =True)

    grid2 = gridspec.GridSpecFromSubplotSpec(rows, cols, subplot_spec=grid[1], wspace=ws, hspace=0.4)#,

    end = ['05']
    #y_sum = [[]] * len(data)
    color_modul = 'steelblue'
    color_mpf = 'red'
    y_sum2 = plot_single(lim, data, end, data_all, grid2, color_mpf, color_modul,title = False,  label = True,remove =False)

    #for dd, d in enumerate(data):
        # embed()
        #embed()
        #ax[1].set_ylim([0, np.nanmax([y_sum1,y_sum2])])
        #ax[1, dd].set_ylim([0, 350])
    plt.subplots_adjust(wspace = 0.4,left = 0.17, right = 0.96,bottom = 0.2)

def plot_single(lim, data, end, data_all, grid1, color_mpf, color_modul,label = True, xlabel = True,remove =True, title = True):
    y_sum = [[]] * len(data)
    ax = {}
    for dd,set in enumerate(data):
        for ee, e in enumerate(end):
            if title == True:
                plt.title('Cell '+str(dd))
            d = data_all[data_all['dataset'] == set]
            x = d['delta_f'] / d['eodf'] + 1
            #embed()
            y = d['result_frequency_' + e]
            y2 = d['result_amplitude_max_' + e]
            y_sum[dd] = np.nanmax(y)
            ff = d['delta_f'] / d['eodf'] + 1
            fe = d['beat_corr']
            #fig.suptitle(set)
            grid2 = gridspec.GridSpecFromSubplotSpec(2,1, subplot_spec=grid1[dd], wspace=0.02, hspace=0.2)  # ,
            ax[0] = plt.subplot(grid2[0])
            ax[0].plot(ff, fe, color='grey', linestyle='--')
            ax[0].plot(x, y, color=color_mpf)

            ax[set+e+str(1)] = plt.subplot(grid2[1])
            if (set == data[0]) and (label == True):
                ax[set+e+str(1)].set_ylabel('Modulation ')
                ax[0].set_ylabel('MPF [EODf]')
            if (set == data[0]) and (xlabel == True):
                ax[set + e + str(1)].set_xlabel('EOD multiples')
            #ax[0, dd].set_title(e + ' ms')
            ax[0].set_xlim([0, 4])

            ax[set+e+str(1)].plot(x, y2, color=color_modul)
            # ax[1,0].set_ylabel('modulation depth [Hz]')
            #ax[2, ee].plot(x, y2, color=colors[0])
            #ax[2, 0].set_ylabel('         modulation depth [Hz]')
            # ax[1,ee].annotate("", xy=(0.53, 16.83), xytext=(0.53, 17.33), arrowprops=dict(arrowstyle="->"))
            # ax[1,ee].annotate("", xy=(1.51, 16.83), xytext=(1.51, 17.33), arrowprops=dict(arrowstyle="->"))

            #ax[1, 0].set_xlabel('stimulus frequency [EODf]')

            #ax[1, 2].set_xlabel('stimulus frequency [EODf]')
            #ax[2, 3].set_xlabel('stimulus frequency [EODf]')

            ax[0].spines['right'].set_visible(False)
            ax[0].spines['top'].set_visible(False)
            ax[set+e+str(1)].spines['right'].set_visible(False)
            ax[set+e+str(1)].spines['top'].set_visible(False)
            #ax[2, ee].spines['right'].set_visible(False)
            #ax[2, ee].spines['top'].set_visible(False)
            ax[0].set_xlim([0, 5])
            #plt.tight_layout()
            # fig.label_axes()
            ax[set+e+str(1)].set_ylim([0, lim])
            #ax[0].set_ylim([0, 240])
            #ax[set+e+str(1)] = remove_tick_marks(ax[0])
            ax[0] = remove_tick_marks(ax[0])
            if set != data[0]:
                ax[0] = remove_tick_ymarks(ax[0])
                ax[set + e + str(1)] = remove_tick_ymarks(ax[set+e+str(1)] )
            if remove == True:
                ax[set+e+str(1)] = remove_tick_marks(ax[set+e+str(1)])
            #ax[0].set_ylim([0, lim])
    return y_sum

if __name__ == "__main__":
    data = ['2019-10-21-aa-invivo-1','2019-11-18-af-invivo-1','2019-10-28-aj-invivo-1']
    data = ['2019-10-21-aa-invivo-1', '2019-10-21-au-invivo-1', '2019-10-28-aj-invivo-1']
    default_settings(data,intermediate_width = 6.7,intermediate_length = 5)
    #fig, ax = plt.subplots(nrows=2, ncols=3, sharex=True)
    ax = {}
    plot_single_cells(ax, data = data)
    #fig.savefig()
    plt.savefig('differentcells_filter.pdf')
    plt.savefig('../highbeats_pdf/differentcells_filter.pdf')
    # plt.subplots_adjust(left = 0.25)
    plt.show()
    #plt.close()
