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 myfunctions import auto_rows
from myfunctions import remove_tick_ymarks
from myfunctions import remove_tick_marks
import string
from matplotlib import gridspec
from plot_eod_chirp import power_func, find_beats, find_dev
#from beatsnumerics import plot_beats
from thunderfish.tabledata import TableData
from plotstyle import plot_style, spines_params
from scipy.signal import hilbert, butter, filtfilt, lfilter
from myfunctions import default_settings

def plot_amp(name,grid,ax, mean1, nrs = [3,4,5],nr_size = 12,name_nr = 'amp',nr = 1, wide = 'gainsboro', rem = 'yes', middle = 'darkgrey', narrow = 'black'):
    np.unique(mean1['type'])
    all_means = mean1[mean1['type'] == name_nr +' mean']
    original = all_means[all_means['dev'] == 'original']
    #m005 = all_means[all_means['dev'] == '005']
    m05 = all_means[all_means['dev'] == '05']
    m2 = all_means[all_means['dev'] == '2']
    # fig, ax = plt.subplots(nrows=4, ncols = 3, sharex=True)
    versions = [m05] #m005,
    dev = '05'
    lim = [[]]*len(versions)
    for i in range(len(versions)):
        keys = [k for k in versions[i]][2::]
        try:
            data = np.array(versions[i][keys])[0]
        except:
            break
        axis = np.arange(0, len(data), 1)
        axis_new = axis * 1
        similarity = [keys, data]
        sim = np.argsort(similarity[0])
        # similarity[sim]
        all_means = mean1[mean1['type'] == name_nr+' std']
        std = all_means[all_means['dev'] == dev]
        std = np.array(std[keys])[0]
        #ax[1, 1].set_ylabel('Modulation depth')
        #ax[nr,i].set_title(dev[i] + ' ms')
        all_means = mean1[mean1['type'] == name_nr+' 95']
        std95 = all_means[all_means['dev'] == dev]
        std95 = np.array(std95[keys])[0]
        all_means = mean1[mean1['type'] == name_nr+' 05']
        std05 = all_means[all_means['dev'] == dev]
        std05 = np.array(std05[keys])[0]
        #ax['mod'+str(i)] = plt.subplot(nrow, ncol, i * 2 + 2)
        ax[name] = plt.subplot(grid[nr])
        if nrs:
            ax[name].text(left, 1.1, string.ascii_uppercase[nrs[i]], transform=ax[name].transAxes,
                   size=nr_size, weight='bold')

        ax[name].fill_between(np.array(keys)[sim], list(std95[sim]), list(std05[sim]),
                              color=wide)
        ax[name].fill_between(np.array(keys)[sim], list(data[sim] + std[sim]), list(data[sim] - std[sim]),
                              color=middle)

        ax[name].spines['right'].set_visible(False)
        ax[name].spines['top'].set_visible(False)
        ax[name].set_xlim([0, 5])
        ax[name].set_xlabel('EOD multiples')
        #if i != 0:
        #    ax['mod'+str(i)] = remove_tick_ymarks(ax['mod'+str(i)])
        #if (i != 2) and (rem != 'no'):
        #    ax[name] = remove_tick_marks(ax[name])
        # ax[i].plot(data_tob.ff, data_tob.fe, color='grey', linestyle='--', label='AMf')
        ax[name].plot(np.array(keys)[sim], data[sim], color=narrow)
        lim[i] = np.nanmax(std95[sim])
        # ax[0].plot(data1.x, data1.freq20, color=colors[1], label='20 %')
    for i in range(len(versions)):
        ax[name].set_ylim(0,np.nanmax(lim))
    #embed()
    return ax

def create_beat_corr(hz_range, eod_fr):
    beat_corr = hz_range%eod_fr
    beat_corr[beat_corr>eod_fr/2] = eod_fr[beat_corr>eod_fr/2] - beat_corr[beat_corr>eod_fr/2]
    return beat_corr

def plot_beats(ax, f0, freqs, fexpected, ffirst,color, fmax, title, ylabel = 'yes'):
    freqs /= f0
    sel = np.abs(freqs - ((freqs+0.1)//0.5)*0.5) < 0.01
    ax.set_title(title)
    ax.plot(freqs, fexpected/f0, color = 'steelblue')
    if ffirst is not None:
        ffirst[sel] = np.nan
        ax.plot(freqs, ffirst/f0, color = 'palegoldenrod')
    if fmax is not None:
        fmax[sel] = np.nan
        ax.plot(freqs, fmax/f0, color = color)
    ax.set_xlim(0, 3.5)
    ax.set_ylim(0, 0.7)
    ax.set_xticks_delta(1.0)

    ax.set_yticks_delta(0.5)
    #ax.set_xlabel('Frequency [f0]')
    if ylabel == 'yes':
        ax.set_ylabel('Frequency [f0]')
    else:
        ax = remove_tick_ymarks(ax)
    ax = remove_tick_marks(ax)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

def plot_mean_cells():
    mean1 = pd.read_pickle('mean.pkl')
    default_settings([0],intermediate_width = 6.29,intermediate_length = 8, ts = 10, ls = 9, fs = 9)

    nr_size = 12
    lp = 50
    colors_tob = ['orange','brown','firebrick','red']
    colors_mod = ['blue','navy','steelblue','steelblue']

    grid00 = gridspec.GridSpec(2, 1, height_ratios = [3,0.8], hspace=0.2, left=0.2)

    grid0  = gridspec.GridSpecFromSubplotSpec(6, 2, height_ratios = [0.01,1,0.01,1,0.01,1],
                                             subplot_spec=grid00[0],wspace=0.3, hspace=0.75)
    grid0_sub  = gridspec.GridSpecFromSubplotSpec(3, 1, height_ratios = [1,1,1],
                                             subplot_spec=grid00[0],wspace=0.25, hspace=0.4)


    eod_fr = 670
    start = 5
    end = eod_fr*5
    step = 25
    eod_fe, beat_corr, beats = find_beats(start, end, step, eod_fr)
    sampling = 100000
    deviation_ms, deviation_s, deviation_dp = find_dev([1.5], sampling)
    left = -0.5
    ax = {}

    #axis_sub = gridspec.GridSpecFromSubplotSpec(1, 1,
    #                                         subplot_spec=grid0_sub[0],wspace=0, hspace=0.7)
    axis_sub = gridspec.GridSpecFromSubplotSpec(1, 1,
                                             subplot_spec=grid0[0,:],wspace=0, hspace=0.7)
    ax['sub'] = plt.subplot(axis_sub[0])
    ax['sub'].spines['top'].set_visible(False)
    ax['sub'].spines['bottom'].set_visible(False)
    ax['sub'].spines['left'].set_visible(False)
    ax['sub'].spines['right'].set_visible(False)
    ax['sub'].set_yticks([])
    ax['sub'].set_xticks([])
    plt.title('Analytic', pad = -lp,fontsize = 10,  fontweight='bold')
    #legend_elements = [Line2D([0], [0], color=color_tob[0], lw=1, label='Maximal Frequency'),
    #                   Line2D([0], [0], color=color_mod[0], label='Modulation')]
    #ax['sub'].legend(handles=legend_elements, loc='center')
    #embed()


    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                             subplot_spec=grid0[1,0],wspace=0, hspace=0.7)

    data = TableData('highbeatspecs10.dat')
    ax['analytic_threhold'] = plt.subplot(axis[0])
    ax['analytic_threhold'].text(left, 1.1, string.ascii_uppercase[0], transform=ax['analytic_threhold'].transAxes,
                            size=nr_size, weight='bold')
    f0 = data[data[:,'freq>norm'] == 1.0,'freq>freq']
    f0 = f0[len(f0)//2]
    plot_beats(ax['analytic_threhold'] , f0, data[:,'freq>freq'], data[:,'freq>expected'],
               data[:,'threshold>first_f'], colors_tob[0], data[:,'threshold>max_f'], 'Thresholded')
    ax['analytic_threhold'].set_xlim([0, 5])
    ax['analytic_threhold_a'] = plt.subplot(axis[1])
    ax['analytic_threhold_a'] .spines['right'].set_visible(False)
    ax['analytic_threhold_a'] .spines['top'].set_visible(False)
    plt.plot(data[:,'freq>freq']/f0,data[:, 'threshold>max_a'], color = colors_mod[0])
    ax['analytic_threhold_a'] = remove_tick_marks(ax['analytic_threhold_a'])
    plt.xlim([0, 5])
    data = TableData('highbeatspecs10.dat')

    #embed()

    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                             subplot_spec=grid0[1,1],wspace=0, hspace=0.5)
    ax['analytic_cube'] = plt.subplot(axis[0])
    f0 = data[data[:,'freq>norm'] == 1.0,'freq>freq']
    f0 = f0[len(f0)//2]
    plot_beats(ax['analytic_cube'] , f0, data[:,'freq>freq'], data[:,'freq>expected'],
               data[:,'threshold cubed>first_f'],colors_tob[0], data[:,'threshold cubed>max_f'], 'Threshold cubed', ylabel = 'no')
    ax['analytic_cube'].set_xlim([0,5])
    ax['analytic_cube_a'] = plt.subplot(axis[1])
    plt.plot(data[:,'freq>freq']/f0, data[:,'threshold cubed>max_a'], color = colors_mod[0])
    ax['analytic_cube_a'].spines['right'].set_visible(False)
    ax['analytic_cube_a'].spines['top'].set_visible(False)
    ax['analytic_cube_a'] = remove_tick_marks(ax['analytic_cube_a'])
    ax['analytic_cube_a'] = remove_tick_ymarks(ax['analytic_cube_a'])
    plt.xlim([0,5])
    #embed()
    #grid1  = gridspec.GridSpecFromSubplotSpec(1, 2,
    #                                         subplot_spec=grid00[1],wspace=0, hspace=0.7)

    axis_sub = gridspec.GridSpecFromSubplotSpec(1, 1,
                                             subplot_spec=grid0[2,:],wspace=0, hspace=0.7)
    ax['sub'] = plt.subplot(axis_sub[0])
    ax['sub'].spines['top'].set_visible(False)
    ax['sub'].spines['bottom'].set_visible(False)
    ax['sub'].spines['left'].set_visible(False)
    ax['sub'].spines['right'].set_visible(False)
    ax['sub'].set_yticks([])
    ax['sub'].set_xticks([])
    plt.title('Numerical',pad = -lp, fontsize = 10,  fontweight='bold')

    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                            subplot_spec=grid0[3, 0], wspace=0, hspace=0.5)
    threshold = power_func(restrict=[2], a_fr=1, a_fe=0.2, eod_fr=eod_fr, eod_fe=eod_fe, win='w2', sigma=0,
                           sampling=sampling, deviation_dp=deviation_dp, deviation_s=deviation_s, beat_corr=beat_corr,
                           size=[120],
                           phase_zero=[0], delta_t=1, show_figure=True,
                           plot_dist=False, save=False, bef_c=-2, aft_c=-1)
    ax['t_f'] = plt.subplot(axis[0])
    plt.title('Threshold')
    ax['t_f'].plot(beats / eod_fr + 1, np.array(threshold['result_frequency'][0]) / eod_fr, color=colors_tob[1])
    ax['t_f'].spines['right'].set_visible(False)
    ax['t_f'].spines['top'].set_visible(False)
    ax['t_f'] = remove_tick_marks(ax['t_f'])
    ax['t_f'].set_xlim([0, 5])
    ax['t_f'].text(left, 1.1, string.ascii_uppercase[1], transform=ax['t_f'].transAxes,
                   size=nr_size, weight='bold')
    ax['t_a'] = plt.subplot(axis[1])
    ax['t_a'].plot(beats / eod_fr + 1, threshold['result_amplitude_max'][0], color=colors_mod[1])
    ax['t_a'].spines['right'].set_visible(False)
    ax['t_a'].spines['top'].set_visible(False)
    ax['t_a'] = remove_tick_marks(ax['t_a'])
    ax['t_a'].set_xlim([0, 5])
    ax['t_f']= plt.subplot(axis[0])



    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                            subplot_spec=grid0[3, 1], wspace=0, hspace=0.5)
    cubed =power_func(restrict = [5],a_fr=1, a_fe=0.2, eod_fr=eod_fr, eod_fe=eod_fe, win='w2', sigma=0,
                         sampling=sampling, beat_corr=beat_corr, size=[120],
                         phase_zero=[0], delta_t=1, show_figure=True,
                         plot_dist=False, save=False, bef_c=-2, aft_c=-1)
    ax['t_f'] = plt.subplot(axis[0])
    plt.title('Threshold cubed')
    ax['t_f'].plot(beats/eod_fr+1,np.array(cubed['result_frequency'][0])/eod_fr, color = colors_tob[1])
    ax['t_f'].spines['right'].set_visible(False)
    ax['t_f'].spines['top'].set_visible(False)
    ax['t_f'] = remove_tick_marks(ax['t_f'])
    ax['t_f'] = remove_tick_ymarks(ax['t_f'])
    ax['t_f'].set_xlim([0, 5])
    ax['t_a'] = plt.subplot(axis[1])
    ax['t_a'] = remove_tick_marks(ax['t_a'])
    ax['t_a'] = remove_tick_ymarks(ax['t_a'])
    ax['t_a'].plot(beats/eod_fr+1,cubed['result_amplitude_max'][0], color = colors_mod[1])
    ax['t_a'].spines['right'].set_visible(False)
    ax['t_a'].spines['top'].set_visible(False)
    ax['t_a'].set_xlim([0, 5])






    #ax = plot_amp('data_mod',axis, ax, data, nrs=[], nr_size=nr_size, name_nr='tob max', nr=2, wide='gainsboro',
    #              middle='grey', narrow='black')


    axis_sub = gridspec.GridSpecFromSubplotSpec(1, 1,
                                             subplot_spec=grid0[4,:],wspace=0, hspace=0.7)
    ax['sub'] = plt.subplot(axis_sub[0])
    ax['sub'].spines['top'].set_visible(False)
    ax['sub'].spines['bottom'].set_visible(False)
    ax['sub'].spines['left'].set_visible(False)
    ax['sub'].spines['right'].set_visible(False)
    ax['sub'].set_yticks([])
    ax['sub'].set_xticks([])
    plt.title('Leaky Integrate and Fire - P-unit model',pad = -lp, fontsize = 10,  fontweight='bold')

    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                             subplot_spec=grid0[5,0], wspace=0, hspace=0.5)
    data = pd.read_pickle('mean' + 'simulation' + '.pkl')
    ax = plot_hz('data_hz',axis, 0, ax, data,nrs = [2],left = left, color = colors_tob[2])
    plt.title('Threshold cubed')
    ax = plot_amp('data_mod',axis, ax, data, nrs=[], nr_size=nr_size, name_nr='amp max', nr=1, wide='gainsboro',
                  middle='lightblue', narrow=colors_mod[2])

    #ax = plot_amp('data_mod',axis, ax, data, nrs=[], nr_size=nr_size, name_nr='tob max', nr=2, wide='gainsboro',
    #              middle='grey', narrow='black')


    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                             subplot_spec=grid0[5,1], wspace=0, hspace=0.5)
    data = pd.read_pickle('mean' + 'sinz' + '.pkl')
    ax = plot_hz('data_hz',axis, 0, ax, data,nr_size = nr_size,color = colors_tob[2])
    plt.title('Threshold cubed')
    ax = plot_amp('data_mod',axis, ax, data, nrs=[], nr_size=nr_size, name_nr='amp max', nr=1, wide='gainsboro',
                  middle='lightblue', narrow=colors_mod[2],rem = 'no')
    #ax = plot_amp('data_mod',axis, ax, data, nrs=[], nr_size=nr_size, name_nr='tob max', nr=2, wide='gainsboro',
    #              middle='grey', narrow='black',rem = 'no')
    ax['data_mod'] = remove_tick_ymarks(ax['data_mod'])
    ax['data_hz'] = remove_tick_ymarks(ax['data_hz'])


    ax_intermediate = gridspec.GridSpecFromSubplotSpec(1, 3,
                                             subplot_spec=grid00[1], width_ratios = [1,2,1], wspace=0, hspace=0.5)

    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                             subplot_spec=ax_intermediate[1], wspace=0, hspace=0.5)
    data = pd.read_pickle('mean' + '' + '.pkl')
    ax = plot_hz('data_hz',axis, 0, ax, data, nrs = [3],left = left, color = colors_tob[2])
    plt.title('P-units', fontweight = 'bold')

    ax = plot_amp('data_mod',axis, ax, data, nrs=[], nr_size=nr_size, name_nr='amp max', nr=1, wide='gainsboro',
                  middle='lightblue', narrow='steelblue')


    plt.subplots_adjust(bottom = 0.1, top = 0.96)
    plt.savefig('compare.pdf')
    try:
        plt.savefig('../highbeats_pdf/compare.pdf')
    except:
        a = 0
    plt.show()
    #fig.tight_layout()f
    # fig.label_axes()
    return fig

def plot_hz(name, grid, nr, ax, mean1,nrs = [], nr_size = 12,left = -0.1, color = 'red'):
    x = np.arange(0, 2550, 50)
    corr = create_beat_corr(x, np.array([500] * len(x)))
    all_means = mean1[mean1['type'] == 'max mean']
    original = all_means[all_means['dev'] == 'original']
    m005 = all_means[all_means['dev'] == '005']
    m05 = all_means[all_means['dev'] == '05']
    m2 = all_means[all_means['dev'] == '2']
    dev = '05'#'005',
    titels = ['0.5 ms']
    versions = m05#m005,
    keys = [k for k in versions][2::]
    try:
        data = np.array(versions[keys])[0]
    except:
        a = 0
    i = 0
    # grid = gridspec.GridSpec(5, 1, hspace=0.4, left=0.2)
    # grid = gridspec.GridSpecFromSubplotSpec(1, 2, hspace = 0.63, subplot_spec=grid0[0])
    axis = np.arange(0, len(data), 1)
    axis_new = axis * 1
    similarity = [keys, data]
    sim = np.argsort(similarity[0])
    # similarity[sim]
    all_means = mean1[mean1['type'] == 'max std']
    std = all_means[all_means['dev'] == dev]
    std = np.array(std[keys])[0]
    # ax[0,i].set_title(dev[i] +' ms')
    ax[name] = plt.subplot(grid[nr])
    #ax[name].set_title(titels[i])
    # ax[i+1].set_ylabel('MPF [EODf]')
    all_means = mean1[mean1['type'] == 'max 95']
    std95 = all_means[all_means['dev'] == dev]
    std95 = np.array(std95[keys])[0]
    all_means = mean1[mean1['type'] == 'max 05']
    std05 = all_means[all_means['dev'] == dev]
    std05 = np.array(std05[keys])[0]
    # std[np.where(list(data[sim] + std[sim])>std95[sim])] = std95[sim][np.where(list(data[sim] + std[sim])>std95[sim])]
    # embed()
    middle = 'pink'  # 'LightSalmon'
    wide = 'lightpink'  # mistyrose'
    narrow = color #'crimson'
    ax[name].fill_between(np.array(keys)[sim], list(std95[sim]-1), list(std05[sim]-1),
                          color=wide, alpha=0.45)  #
    # embed()

    ax[name].fill_between(np.array(keys)[sim], list(data[sim] + std[sim]-1), list(data[sim] - std[sim]-1), color=middle)

    ax[name].plot(x / 500, corr / 500, color='grey', linestyle='--', label='AMf')
    ax[name].plot(np.array(keys)[sim], data[sim]-1, color=narrow)
    ax[name].spines['right'].set_visible(False)
    ax[name].spines['top'].set_visible(False)
    # plt.fill_between(np.array([0,1]),np.array([0,0]),np.array([1,1]))
    if nrs:
        ax[name].text(left, 1.1, string.ascii_uppercase[nrs[i]], transform=ax[name].transAxes,
                      size=nr_size, weight='bold')
    ax[name] = remove_tick_marks(ax[name])
    ax[name].set_xlim([0,5])
    # if i != 0:
    #    ax[i+1] = remove_tick_ymarks(ax[i+1])
    # ax[0].plot(data1.x, data1.freq20, color=colors[1], label='20 %')
    return ax

if __name__ == "__main__":
    #plot_style()
    fig = plot_mean_cells()
    #fig.savefig()
    plt.savefig('compare.pdf')
    try:
        plt.savefig('../highbeats_pdf/compare.pdf')
    except:
        pass
    plt.show()
    # plt.close()

