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 differentcells import default_settings
#from differentcells import plot_single_cells
import matplotlib.gridspec as gridspec
from functionssimulation import single_stim
import math
from functionssimulation import find_times
from functionssimulation import rectify
from functionssimulation import global_maxima
from functionssimulation import integrate_chirp
from functionssimulation import find_periods
from myfunctions import default_settings, load_cell
from axes import labelaxes_params,label_axes
from mpl_toolkits.axes_grid1 import host_subplot
import mpl_toolkits.axisartist as AA
import string



def plot_single_cells(ax, colors = ['#BA2D22', '#F47F17', '#AAB71B', '#3673A4', '#53379B'], data = ['2019-10-21-aa-invivo-1','2019-11-18-af-invivo-1','2019-10-28-aj-invivo-1'], var = '05'):
    y_sum = [[]] * len(data)
    axis = {}
    for ee, e in enumerate([var]):
        #d = data_all[data_all['dataset'] == ]
        d = load_cell(data[0], fname='singlecellexample5_smoothed'+data[0], big_file='beat_results_smoothed')
        eod = d['eodf'].iloc[0]
        #embed()
        x = d['delta_f'] / d['eodf'] + 1
        xx = d['delta_f']
        y = d['result_frequency_' + e]
        y2 = d['result_amplitude_max_' + e]
        y_sum[0] = np.nanmax(y)
        axis[1] = plt.subplot(ax[0])
        axis[1].plot(x, y, zorder = 1,color=colors[0])
        axis[1].set_ylabel('[Hz]')
        axis[1].set_xlim([0, 4])
        labels = [item.get_text() for item in axis[1].get_xticklabels()]
        empty_string_labels = [''] * len(labels)
        axis[1].set_xticklabels(empty_string_labels)

        #axis[2] = host_subplot(ax[1], axes_class=AA.Axes)
        axis[2] = plt.subplot(ax[1])
        #aixs[3] = axis[2].twinx()
        #par2 = axis[2].twinx()
        #host = host_subplot(ax[1], axes_class=AA.Axes)
        #host.spines['right'].set_visible(False)
        #host.spines['top'].set_visible(False)
        #axis[2] = host.twiny()
        axis[2].plot(xx, y2,  label="Beats [Hz]", zorder = 2,color=colors[0])
        axis[2].set_ylabel('Modulation ')
        axis[1].spines['right'].set_visible(False)
        axis[1].spines['top'].set_visible(False)
        axis[2].spines['right'].set_visible(False)
        axis[2].spines['top'].set_visible(False)
        axis[1].set_xlim([0, np.max(x)])
        axis[2].set_xlim([-eod, np.max(xx)])

        nr_size = 10
        axis[2].text(-0.02, 1.1, string.ascii_uppercase[4],
                                          transform=axis[2].transAxes,
                                          size=nr_size, weight='bold')
        axis[1].text(-0.02, 1.1, string.ascii_uppercase[3],
                                          transform=axis[1].transAxes,
                                          size=nr_size, weight='bold')
        #axis[4] = axis[2].twiny()
        axis[3] = axis[2].twiny()
        axis[3].set_xlabel('EOD multiples')
        offset = -40
        #new_fixed_axis = axis[3].get_grid_helper().new_fixed_axis
        #axis[3].axis["bottom"] = new_fixed_axis(loc="bottom", axes=axis[3],
        #                                    offset=(0,offset))
        axis[3].xaxis.set_ticks_position("bottom")
        axis[3].xaxis.set_label_position("bottom")

        axis[3].spines["bottom"].set_position(("axes", -1.2))
        axis[3].spines['right'].set_visible(False)
        axis[3].spines['top'].set_visible(False)
        #axis[3].axis["bottom"].toggle(all=True)
        axis[2].set_xlabel("Difference frequency [Hz]")
        #par2.set_xlim([np.min(xx), np.max(xx)])
        axis[3].set_xlim([0, np.max(x)])
        #p1, = host.plot([0, 1, 2], [0, 1, 2], label="Density")
        #p2, = par1.plot([0, 1, 2], [0, 3, 2], label="Temperature")
        p3, = axis[3].plot(x, y2,color = 'grey',zorder = 1)
        #p3, = axis[4].plot(x, y2, color='grey', zorder=1)
        #embed()
        axis[2].set_xticks(np.arange(-eod,np.max(xx),eod/2))
        #ax['corr'].set_yticks(np.arange(eod_fe[0] - eod_fr, eod_fe[-1] - eod_fr, eod_fr / 2))

    axis[2].set_ylim([0, np.nanmax(y_sum)])
    plt.subplots_adjust(wspace = 0.4,left = 0.17, right = 0.96,bottom = 0.2)
    return axis,y2

def plot_beat_corr(ax,lower,beat_corr_col = 'red',df_col = 'pink',ax_nr = 3,multiple = 3):
    eod_fr = 500
    eod_fe = np.arange(0, eod_fr * multiple, 5)
    beats = eod_fe - eod_fr
    beat_corr = eod_fe % eod_fr
    beat_corr[beat_corr > eod_fr / 2] = eod_fr - beat_corr[beat_corr > eod_fr / 2]
    #gs0 = gridspec.GridSpec(3, 1, height_ratios=[4, 1, 1], hspace=0.7)

    #plt.figure(figsize=(4.5, 6))
    style = 'dotted'
    color_v = 'black'
    color_b = 'silver'
    # plt.subplot(3,1,1)
    ax['corr'] = plt.subplot(lower[ax_nr])
    np.max(beats) / eod_fr

    ax['corr'].set_xticks(np.arange((eod_fe[0]-eod_fr)/eod_fr+1, (eod_fe[-1]-eod_fr)/eod_fr+1,(eod_fr/2)/eod_fr+1))
    ax['corr'].set_yticks(np.arange((eod_fe[0]-eod_fr)/eod_fr+1, (eod_fe[-1]-eod_fr)/eod_fr+1,(eod_fr/2)/eod_fr+1))
    ax['corr'].set_xticks(np.arange(0,10,0.5))
    ax['corr'].set_yticks(np.arange(0, 10, 0.5))
    # plt.axvline(x = -250, Linestyle = style,color = color_v)
    # plt.axvline(x = 250, Linestyle = style,color = color_v)
    # plt.axvline(x = 750, Linestyle = style,color = color_v)
    # plt.axvline(x = 1500, Linestyle = style)
    # plt.subplot(3,1,2)
    plt.xlabel('Beats [Hz]')
    plt.ylabel('Difference frequency [Hz]')
    #plt.subplot(gs0[1])
    if beat_corr_col != 'no':
        plt.plot(beats/eod_fr+1, beat_corr/(eod_fr+1), color=beat_corr_col, alpha = 0.7)
        plt.ylim([0,np.max(beat_corr/(eod_fr+1))*1.4])
        plt.xlim([(beats/eod_fr+1)[0],(beats/eod_fr+1)[-1]])
    if df_col != 'no':
        plt.plot(beats/eod_fr+1, np.abs(beats)/(eod_fr+1), color=df_col, alpha = 0.7)
    #plt.axvline(x=-250, Linestyle=style, color=color_v)
    #plt.axvline(x=250, Linestyle=style, color=color_v)
    #plt.axvline(x=750, Linestyle=style, color=color_v)
    plt.xlabel('EOD adjusted beat [Hz]')
    ax['corr'] .spines['right'].set_visible(False)
    ax['corr'] .spines['top'].set_visible(False)
    ax['corr'] .spines['left'].set_visible(True)
    ax['corr'] .spines['bottom'].set_visible(True)
    # plt.axvline(x = 1250, Linestyle = style,color = color_v)
    # plt.axvline(x = 1500, Linestyle = style,color = color_v)
    mult = np.array(beats) / eod_fr + 1
    # plt.subplot(3,1,3)
    plt.xlabel('EOD multiples')
    plt.ylabel('EOD adj. beat [Hz]', fontsize = 10)
    plt.grid()
    #plt.subplot(gs0[2])
    #plt.plot(mult, beat_corr, color=color_b)
    # plt.axvline(x = 0, Linestyle = style)
    #plt.axvline(x=0.5, Linestyle=style, color=color_v)
    # plt.axvline(x = 1, Linestyle = style)
    #plt.axvline(x=1.5, Linestyle=style, color=color_v)
    #plt.axvline(x=2.5, Linestyle=style, color=color_v)
    #plt.xlabel('EOD multiples')
    #plt.ylabel('EOD adj. beat [Hz]', fontsize = 10)
    return ax

def try_resort_automatically():
    diffs = np.diff(dfs)
    fast_sampling = dfs[np.concatenate([np.array([True]),diffs <21])]
    second_derivative = np.diff(np.diff(fast_sampling))
    first_index = np.concatenate([np.array([False]),second_derivative <0])
    second_index = np.concatenate([second_derivative > 0,np.array([False])])
    remaining = fast_sampling[np.concatenate([np.array([True]),second_derivative == 0, np.array([True])])]
    first = np.arange(0,len(first_index),1)[first_index]
    second = np.arange(0, len(second_index), 1)[second_index]-1
    residual = []
    indeces = []
    for i in range(len(first)):
        index = np.arange(first[i],second[i],2)
        index2 = np.arange(first[i], second[i], 1)
        indeces.append(index2)
        residual.append(fast_sampling[index])#first[i]:second[i]:2
    indeces = np.concatenate(indeces)
    remaining = fast_sampling[~indeces]
    residual = np.concatenate(residual)
    new_dfs = np.sort(np.concatenate([residual, remaining]))


def gap_subplot(ax, dfs, pos_left = 0, pos_right = 1, xl = -0.5, xr = 30, ax_name = 'between'):
    ax[ax_name] = plt.subplot(grid[pos_left, pos_right])
    ax[ax_name].spines['right'].set_visible(False)
    ax[ax_name].spines['top'].set_visible(False)
    ax[ax_name].spines['left'].set_visible(False)
    ax[ax_name].spines['bottom'].set_visible(False)
    ax[ax_name].set_ylim([np.min(dfs), np.max(dfs)])
    ax[ax_name].set_xlim([xl,xr])
    ax[ax_name].set_xticks([])
    ax[ax_name].set_yticks([])
    ax[ax_name].set_ylim(ax[ax_name].get_ylim()[::-1])
    return ax[ax_name]

def morane_plot(ax, dfs, grid, d,nr_size = 11, ll = 0.1, ul = 0.3, ax_name = 'scatter',gl = 0, gr =0):
    ax[ax_name] = plt.subplot(grid[gl,gr])
    ax[ax_name].spines['right'].set_visible(False)
    ax[ax_name].spines['top'].set_visible(False)
    counter = 0
    new_dfs = np.concatenate([dfs[0:25], dfs[25:40:2], dfs[40:53:2], dfs[54:-1]])
    for i in range(len(new_dfs)):
        spikes = d[d['df'] == new_dfs[i]]['spike_times']
        counter += 1
        transformed_spikes = spikes.iloc[0]-spikes.iloc[0][0]
        used_spikes = transformed_spikes[transformed_spikes>ll]
        used_spikes = used_spikes[used_spikes<ul]*1000
        ax['scatter'].scatter(used_spikes,np.ones(len(used_spikes))*new_dfs[i],s = 0.2,color = 'silver')
    ax[ax_name].set_ylim([np.min(dfs),np.max(dfs)])
    ax[ax_name].set_xlim([ll*1000,ul*1000])
    ax[ax_name].set_ylabel('Difference frequency [Hz]')
    ax[ax_name].set_xlabel('Time [ms]')
    ax[ax_name].set_ylim(ax['scatter'].get_ylim()[::-1])
    ax[ax_name].text(-0.1, 1.025, string.ascii_uppercase[0], transform=ax['scatter'].transAxes,
            size= nr_size, weight='bold')
    return ax[ax_name]

def examples(dfs,ax, axis, eod,grid, lg = 0, rg = 2):
    new_dfs = np.concatenate([dfs[0:25], dfs[25:40:2], dfs[40:53:2], dfs[54:-1]])
    first = ['darkviolet','lightcoral','crimson',  ]
    third = ['orange', 'orangered', 'darkred']
    fourth = ['DarkGreen', 'LimeGreen', 'YellowGreen']
    fifth = ['lightblue','blue','navy']
    colors = np.concatenate([fourth, third, first, fifth])
    low_nr = 60
    from_middle = 45  # 20
    example_df = [low_nr - eod, eod / 2 - from_middle - eod,
                  low_nr, eod / 2 - from_middle,
                  low_nr + eod, eod / 2 - from_middle + eod,
                  low_nr + eod * 2, eod / 2 - from_middle + eod * 2,
                  low_nr + eod * 3, eod / 2 - from_middle + eod * 3,
                  low_nr + eod * 4 + eod * 4]
    example_df = [low_nr - eod, eod / 2 - from_middle - eod, eod - low_nr - eod, low_nr, eod / 2 - from_middle,
                  low_nr + eod,eod / 2 - from_middle+eod*2]

    #embed()
    dfs_all = [[]]*len(example_df)
    for i in range(len(example_df)):
        df = new_dfs[np.argmin(np.abs(new_dfs - example_df[i]))]
        dfs_all[i] = [df]
    example_df = np.unique(dfs_all)
    max_p = [[]] * len(example_df)
    min_p = [[]] * len(example_df)
    rows = len(example_df)
    cols = 1
    power_raster = gridspec.GridSpecFromSubplotSpec(rows, cols,
                                                    subplot_spec=grid[lg, rg], wspace=0.05, hspace=0.3)


    for i in range(len(example_df)):

        power = gridspec.GridSpecFromSubplotSpec(1, 2, width_ratios=[1, 1.7], hspace=0.2, wspace=0.2,
                                                 subplot_spec=power_raster[i])

        ax_nr = 0
        ax['scatter_small' + str(i)] = plt.subplot(power[ax_nr])
        eod_fr = np.array(500)
        eod_fr = np.array(eod)
        eod_fe = np.array([example_df[i] + eod])
        e = 0
        factor = 200
        sampling = 100000
        minus_bef = -250
        plus_bef = -200
        f_max, lims, _ = single_stim(ax, [colors[i]], 1, 1, eod_fr, eod_fe, e, power, delta_t=0.001, add='no',
                                     minus_bef=minus_bef, plus_bef=plus_bef, sampling=sampling,
                                     col_basic='silver', col_hline='no', labels=False, a_fr=1, ax_nr=ax_nr,
                                     phase_zero=[2*np.pi*0.4], shift_phase=0, df_col='no', beat_corr_col='no', size=[120],
                                     a_fe=0.8)

        ax['between'].scatter(0.12, example_df[i], zorder=2, s=25, marker='<', color=colors[i])
        ax['between'].scatter(0.12, example_df[i], zorder=2, s=25, marker='<', color=colors[i])
        ll = np.abs(plus_bef)
        ul = np.abs(minus_bef)
        #df = new_dfs[np.argmin(np.abs(new_dfs - example_df[i]))]
        df = example_df[i]
        spikes = d[d['df'] == df]['spike_times']

        tranformed_spikes = spikes.iloc[0] * 1000 - spikes.iloc[0][0] * 1000
        used_spikes = tranformed_spikes[tranformed_spikes > ll]
        used_spikes = used_spikes[used_spikes < ul]
        used_spikes = used_spikes - used_spikes[0]
        ax['scatter_small' + str(i)].scatter((used_spikes), np.ones(len(used_spikes)) * -2, zorder=2, s=2, marker='|',
                                             color='black')  # color = colors[i]
        ax['scatter_small' + str(i)].set_ylim([-2.5, 2.2])

        ax['power' + str(i)] = plt.subplot(power[1])
        nfft = 4096  #
        sampling_rate = 40000
        # embed()

        pp = [[]] * len(spikes)
        for s in range(len(spikes)):
            new_spikes = list(map(int, (spikes.iloc[s] - spikes.iloc[s][0]) * sampling_rate))
            array = np.zeros(new_spikes[-1] + 2)
            array[new_spikes] = 1
            array = array * sampling_rate
            if var == '05':
                window05 = 0.0005 * sampling_rate
                array = gaussian_filter(array, sigma=window05) * sampling_rate
            pp[s], f = ml.psd(array - np.mean(array), Fs=sampling_rate, NFFT=nfft, noverlap=nfft / 2)
        p = np.mean(pp, axis=0)
        diff = d['eodf'].iloc[0] * (df / d['eodf'].iloc[0] - int(df / d['eodf'].iloc[0]))
        if diff > d['eodf'].iloc[0] * 0.5:
            diff = diff - d['eodf'].iloc[0] * 0.5
        ref = (np.max(p))
        p = 10 * np.log10(p / ref)
        plt.plot(f, p, zorder=1, color=main_color)
        max_p[i] = np.max(p)
        min_p[i] = np.min(p)
        ax['power' + str(i)].scatter(f[np.argmax(p[f < 0.5 * eod])], max(p[f < 0.5 * eod]), zorder=2, color=colors[i],
                                     s=25)
        ax['power' + str(i)].scatter(f[f == f[np.argmin(np.abs(f - eod))]],
                                     p[f == f[np.argmin(np.abs(f - eod))]] * 0.90, zorder=2,
                                     s=25, color='darkgrey', edgecolor='black')
        ax['power' + str(i)].axvline(x=eod / 2, color='black', linestyle='dashed', lw=0.5)
        plt.xlim([-40, 1600])
        axis[3].scatter(example_df[i] / (eod) + 1, np.sqrt(np.max(p[f < 0.5 * eod]) * np.abs(f[0] - f[1])), zorder=3,
                        s=20, marker='o', color=colors[i])
        axis[1].scatter(example_df[i] / (eod) + 1, f[np.argmax(p[f < 0.5 * eod])], zorder=2, s=20, marker='o',
                        color=colors[i])

        if i != rows - 1:
            labels = [item.get_text() for item in ax['scatter_small' + str(i)].get_xticklabels()]
            empty_string_labels = [''] * len(labels)
            ax['scatter_small' + str(i)].set_xticklabels(empty_string_labels)
            labels = [item.get_text() for item in ax['power' + str(i)].get_xticklabels()]
            empty_string_labels = [''] * len(labels)
            ax['power' + str(i)].set_xticklabels(empty_string_labels)
        else:
            ax['power' + str(i)].set_xlabel('Frequency [Hz]')
            ax['scatter_small' + str(i)].set_xlabel('Time [ms]')
        ax['power' + str(i)].set_yticks([])
        ax['power' + str(i)].spines['left'].set_visible(False)
        ax['scatter_small' + str(i)].spines['left'].set_visible(False)
        ax['scatter_small' + str(i)].set_yticks([])
        ax['power' + str(i)].spines['right'].set_visible(False)
        ax['power' + str(i)].spines['top'].set_visible(False)
        ax['scatter_small' + str(i)].spines['right'].set_visible(False)
        ax['scatter_small' + str(i)].spines['top'].set_visible(False)

    for i in range(len(example_df)):
        #ax['power' + str(i)].set_ylim([0, np.max(max_p)])
        #ax['power' + str(i)].set_ylim([np.min(min_p), np.max(max_p)*1.3])
        ax['power' + str(i)].set_ylim([np.min(min_p), 4])
    #embed()
    ax['power' + str(0)].text(-0.1, 1.1, string.ascii_uppercase[2], transform=ax['power' + str(0)].transAxes,
                              size=nr_size, weight='bold')
    ax['scatter_small' + str(0)].text(-0.1, 1.1, string.ascii_uppercase[1],
                                      transform=ax['scatter_small' + str(0)].transAxes,
                                      size=nr_size, weight='bold')
    return ax, axis

if __name__ == "__main__":
     # data = ['2019-10-21-aa-invivo-1']
    data = ['2019-11-18-ac-invivo-1']# '2019-10-28-aj-invivo-1'
    data = ['2019-10-28-aj-invivo-1']
    data = ['2019-09-23-ad-invivo-1']#6.29
    default_settings(data,intermediate_width = 6.28,intermediate_length = 7.5, ts = 6, ls = 9, fs = 9)
    fig = plt.figure()
    ax = {}
    nr_size = 10
    main_color = 'darkgrey'
    var = 'original'#var = '05'
    d = load_cell(data[0], fname='singlecellexample5_data_beat'+data[0], big_file='data_beat',new = False)
    eod = d['eodf'].iloc[0]
    dfs = np.unique(d['df'])
    x = d['df'] / d['eodf'] + 1
    grid = gridspec.GridSpec(2, 3, wspace=0.0, height_ratios=[6, 2], width_ratios=[1,0.3,3], hspace=0.24)


    #first subplot
    ax['scatter'] = morane_plot(ax, dfs, grid, d, nr_size=nr_size, ll=0.1, ul=0.3, ax_name='scatter', gl=0, gr=0)

    #gap between subplots

    ax['between'] = gap_subplot(ax,dfs, pos_left=0, pos_right=1, xl=-0.5, xr=30,  ax_name='between')


    #summary picture
    axis = gridspec.GridSpecFromSubplotSpec(2, 1,
                                             subplot_spec=grid[1,:], wspace=0, hspace=0.5)
    axis,y2 = plot_single_cells(axis, colors = [main_color], data = data,var = var)

    #examples
    ax, axis = examples(dfs,ax, axis, eod, grid, lg=0, rg=2)

    plt.subplots_adjust(left = 0.11, bottom = 0.16, top = 0.94)
    plt.savefig('singlecellexample5.pdf')
    plt.savefig('../highbeats_pdf/singlecellexample5.pdf')
    plt.savefig('singlecellexample5'+data[0]+'.pdf')
    plt.savefig('../highbeats_pdf/singlecellexample5'+data[0]+'.pdf')
    plt.show()