from plot_eod_chirp import find_lm,find_periods,integrate_chirp,global_maxima,rectify,find_dev,power_func,power_parameters
from plot_eod_chirp import rectify, find_beats,find_dev,find_times,global_maxima, find_lm, conv
from functionssimulation import single_stim,pl_eods
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.pyplot as plt
import numpy as np
from IPython import embed
import matplotlib as matplotlib
import math
import scipy.integrate as integrate
from scipy import signal
from scipy.interpolate import interp1d
from scipy.interpolate import CubicSpline
import scipy as sp
import pickle
from scipy.spatial import distance
from myfunctions import *
import time
from matplotlib import gridspec
from matplotlib_scalebar.scalebar import ScaleBar
import matplotlib.mlab as ml
import scipy.integrate as si
import pandas as pd
from myfunctions import default_settings
import os
import string


def snip(left_c,right_c,e,g,sampling, deviation_s,d,eod_fr, a_fr, eod_fe,phase_zero,p, size,s, sigma,a_fe,deviation,beat_corr, chirp = True):
    time, time_cut, cut = find_times(left_c[g], right_c[g], sampling, deviation_s[d])
    eod_fish_e, eod_fish_r, period_fish_r, period_fish_e = find_periods(a_fe, time, eod_fr, a_fr, eod_fe, e)
    #embed()
    if chirp == False:
        eod_fe_chirp = eod_fish_e
    else:
        eod_fe_chirp = integrate_chirp(a_fe, time, eod_fe[e], phase_zero[p], size[s], sigma)
    eod_rec_down, eod_rec_up = rectify(eod_fish_r, eod_fe_chirp)  # rectify
    eod_overlayed_chirp = (eod_fish_r + eod_fe_chirp)[cut:-cut]
    threshold_cube = (eod_rec_up)**3
    #embed()
    maxima_values, maxima_index, maxima_interp = global_maxima(period_fish_e, period_fish_r,
                                                               eod_rec_up[cut:-cut])  # global maxima
    index_peaks, value_peaks, peaks_interp = find_lm(eod_rec_up[cut:-cut])  # local maxima
    middle_conv, eod_conv_down, eod_conv_up, eod_conv_downsampled = conv(eod_fr,sampling, cut, deviation[d], eod_rec_up,
                                                                         eod_rec_down)  # convolve
    eod_fish_both = integrate_chirp(a_fe, time, eod_fe[e] - eod_fr, phase_zero[p], size[s], sigma)
    am_corr_full = integrate_chirp(a_fe, time_cut, beat_corr[e], phase_zero[p], size[s],
                                   sigma)  # indirect am calculation
    _, time_fish, cut_f = find_times(left_c[g], right_c[g], eod_fr, deviation_s[d])  # downsampled through fish EOD
    am_corr_ds = integrate_chirp(a_fe, time_fish, beat_corr[e], phase_zero[p], size[s], sigma)
    am_df_ds = integrate_chirp(a_fe, time_fish, eod_fe[e] - eod_fr, phase_zero[p], size[s],
                              sigma)  # indirect am calculation
    return threshold_cube , time_cut, eod_conv_up, am_corr_full, peaks_interp, maxima_interp, am_corr_ds,am_df_ds,eod_fish_both,eod_overlayed_chirp



def power_func2(bef_c = -2, aft_c = -1, a_fr = 1, factor = 200, sampling = 10000, phase_zero = np.arange(0, 2 * np.pi, 2 * np.pi / 10),new = True, eod_fe =  [50], beat_corr =  [50], beats = [50],eod_fr = 500,size = [120]  ,delta_t = 0.014,a_fe = 0.2, win = 'w2',deviation = [150], save = False):
    print('simulation')
    interest_interval = delta_t * 1.2
    #bef_c = interest_interval / 2
    #aft_c = interest_interval / 2
    sigma = delta_t / math.sqrt((2 * math.log(10)))  # width of the chirp
    # maximal frequency excursion during chirp / 60 or 100 here

      # eod fish reciever
     # amplitude fish reciever

    #sampling = 500
    d = 1
    x = [1.5, 2.5, 0.5, ]
    x = [1.5]
    time_range = 200 * delta_t
    deviation_ms, deviation_s, deviation_dp = find_dev(x, sampling)
    if (not os.path.exists('power_simulation.pkl')) or (new == True):
        results = [[]] * 1
        for d in range(len(deviation)):
            bef_c = bef_c
            aft_c = aft_c
            left_c, right_c, left_b, right_b, period_distance_c, period_distance_b, _, period, to_cut,exclude,consp_needed,deli,interval = period_calc(
                [float('Inf')]*len(beat_corr), 50, win, deviation_s[d], sampling, beat_corr, bef_c, aft_c, 'stim')
            save_n = win

            for s in range(len(size)):
                for p in range(len(phase_zero)):
                    beats = eod_fe - eod_fr
                    for e in range(len(eod_fe)):
                        #embed()
                        left_b = [-0.3*sampling]*len(eod_fe)
                        right_b = [-0.1 * sampling]*len(eod_fe)
                        threshold_cube, time_b, conv_b, am_corr_b, peaks_interp_b, maxima_interp_b,am_corr_ds_b, am_df_ds_b, am_df_b,eod_overlayed_chirp = snip(left_b, right_b,e,e,
                                                                                                    sampling, deviation_s,
                                                                                                    d, eod_fr, a_fr, eod_fe,
                                                                                                    phase_zero, p, size, s,
                                                                                                    sigma, a_fe, deviation,
                                                                                                    beat_corr)
                        #time_c, conv_c, am_corr_c, peaks_interp_c, maxima_interp_c,am_corr_ds_c, am_df_ds_c, am_df_c = snip(left_c, right_c,e,e,
                        #                                                                            sampling, deviation_s,
                        #                                                                            d, eod_fr, a_fr, eod_fe,
                        #                                                                            phase_zero, p, size, s,
                        #                                                                            sigma, a_fe, deviation,
                        #                                                                            beat_corr)

                        #embed()
                        nfft = 4096
                        if sampling>1000:
                            #nfft = int((4096 * sampling/ 10000) * 2)
                            nfft = int(4096 * sampling / 40000 * 2)
                        else:
                            nfft = 4096
                        #embed()
                        name = ['conv','df','df ds','corr','corr ds','global max','local max']
                        var = [conv_b,am_df_b,am_df_ds_b, am_corr_b, am_corr_ds_b,maxima_interp_b,peaks_interp_b ]
                        samp = [sampling,sampling,eod_fr,sampling,eod_fr,sampling,sampling]
                        name = ['global max']
                        var = [maxima_interp_b]
                        samp = [sampling]
                        #pp, f = ml.psd(eod_overlayed_chirp - np.mean(eod_overlayed_chirp), Fs=sampling, NFFT=nfft, noverlap=nfft / 2)

                        for i in range(len(samp)):
                            #print(i)

                            plot = False
                            pp, f = ml.psd(var[i] - np.mean(var[i]), Fs=samp[i], NFFT=nfft,
                                           noverlap=nfft / 2)

                            #plot = True
                            if plot == True:
                                plt.figure()
                                plt.subplot(1,2,1)
                                plt.plot(var[i])
                                plt.title(name[i])
                                plt.subplot(1, 2, 2)
                                plt.plot(f,pp)
                                #plt.xlim([0,2000])
                                plt.show()
                            #print(results)

                            #embed()
                            #eod_fr2 = eod_fr*3
                            eod_fr2 = eod_fr
                            try:
                                if type(results[i]) != dict:
                                    results[i] = {}
                                    results[i]['type'] = name[i]
                                    #embed()
                                    results[i]['f'] = list([f[np.argmax(pp[f < 0.5 * eod_fr2])]])
                                    results[i]['amp'] = list([np.sqrt(si.trapz(pp, f, np.abs(f[1]-f[0])))])
                                    results[i]['max'] = list([np.sqrt(np.max(pp[f < 0.5 * eod_fr2])*np.abs(f[1]-f[0]))])
                                else:
                                    results[i]['f'].extend(list([f[np.argmax(pp[f < 0.5 *eod_fr2])]]))
                                    #embed()
                                    results[i]['amp'].extend(list([np.sqrt(si.trapz(pp, f, np.abs(f[1]-f[0])))]))
                                    results[i]['max'].extend(list([np.sqrt(np.max(pp[f < 0.5 * eod_fr2]) * np.abs(f[1] - f[0]))]))
                            except:
                                embed()
        if save:
            #embed()
            results = pd.DataFrame(results)
            results.to_pickle('power_simulation.pkl')
            np.save('power_simulation.npy', results)
    else:
        results = pd.read_pickle('power_simulation.pkl')
    return results


def plot_power(a_fe, a_fr,beats, beat_corr, eod_fe, sampling, ax, eod_fr,nr_size = 11, left = 0.15,nrs = [6,7], bef_c = -2, aft_c = -1):
    plt.rcParams['figure.figsize'] = (3, 3)
    plt.rcParams["legend.frameon"] = False
    colors = ['black', 'magenta', 'pink', 'orange', 'moccasin', 'red', 'green', 'silver']
    colors = ['red','pink']
    #eod_fr = 500
    #start = 5
    #end = 2500
    #step = 5
    #eod_fe, beat_corr, beats = find_beats(start, end, step, eod_fr)
    #embed()
    #beats = eod_fe - eod_fr
    factor = 200
    #sampling = eod_fr * factor,
    #(a_fr
    # = 1, a_fe = 0.2, eod_fr = eod_fr, eod_fe = eod_fe, win = 'w2', deviation_s = deviation_s, sigma = sigma, sampling =  100000, deviation_ms = deviation_ms, beat_corr = beat_corr, size =[120], phase_zero =[phase_zero[0]], delta_t = delta_t, deviation_dp = deviation_dp, show_figure = True, plot_dist = False, save = False, bef_c = -2, aft_c = -1)
    results = power_func2(bef_c = bef_c, aft_c = aft_c, a_fe = a_fe, a_fr = a_fr, sampling = sampling, phase_zero = [0], eod_fe = eod_fe, beat_corr = beat_corr, eod_fr = eod_fr, save = True)
    #results = [results.iloc[5]]
    results = [results.iloc[0]]
    #fig, ax = plt.subplots(nrows=2, ncols=1, sharex=True)
    all_max = [[]] * len(results)
    #embed()
    for i in range(len(results)):
        #embed()
        #ax[0].set_ylabel(results[i]['type'], rotation=0, labelpad=40, color=colors[i])
        #embed()
        ax[0].plot(beats / eod_fr + 1, np.array(results[i]['f']) / eod_fr + 1, color=colors[i])

        ax[0].text(-left, 1.1, string.ascii_uppercase[nrs[0]], transform=ax[0].transAxes,
                   size=nr_size, weight='bold')
        # plt.title(results['type'][i])
        #ax[1].plot(beats / eod_fr + 1, np.array(results[i]['amp']), color=colors[0])
        ax[1].plot(beats / eod_fr + 1, np.array(results[i]['max']), color=colors[1])
        ax[1].text(-0.1, 1.1, string.ascii_uppercase[nrs[1]], transform=ax[1].transAxes,
                   size=nr_size, weight='bold')
        #remove_tick_marks(ax[1])
        remove_tick_marks(ax[0])
        #ax[2].plot(beats / eod_fr + 1, np.array(results[i]['amp']), color=colors[i])
        all_max[i] = np.max(np.array(results[i]['amp']))
    #for i in range(len(results)):
    #    ax[2].set_ylim([0, np.max(all_max)])
    plt.subplots_adjust(left=0.25)
    #ii, jj = np.shape(ax)
    ax[0].set_ylabel('[Hz]')
    ax[1].set_ylabel('Modulation')
    #ax[0, 2].set_title('Modulation depth (same scale)')
    for i in [0,1]:
        ax[1].set_xlabel('EOD multiples')
        ax[i].spines['right'].set_visible(False)
        ax[i].spines['top'].set_visible(False)
        ax[i].set_xlim([0,5])
    #plt.subplots_adjust(bottom = 0.2)
    #plt.savefig('localmaxima.pdf')
    #plt.savefig('../highbeats_pdf/localmaxima.pdf')
    #plt.show()




def plot_eod(beat_corr_upper, beats_upper,eod_fe_upper,beat_corr, minus_bef ,plus_bef,row,col,a_fe, delta_t, a_fr, size,beat_corr_org,s,p, amplitude, phase_zero, deviation_ms, eod_fe,deviation, eod_fr, sampling,factor,fs = [6.2992, 4], shift_phase = 0, bef_c = -2, aft_c = -1):
    beats = eod_fe - eod_fr
    nr_size = 11
    plt.rcParams['axes.titlesize'] = 10
    plt.rcParams['axes.labelsize'] = 10
    plt.rcParams['font.size'] = 8
    plt.rcParams["legend.frameon"] = False
    beat_corr_col = 'goldenrod'
    df_col = 'steelblue'
    interval = np.hstack([np.linspace(0.25, 0.97)])
    colors = plt.cm.Blues(interval)
    colors = plt.cm.Reds(interval)
    cmap = LinearSegmentedColormap.from_list('name', colors)
    colors = cmap(np.linspace(0, 1, len(eod_fe)))
    d = 0
    fig = plt.figure(figsize=fs)#hspace = 0.55
    outer = gridspec.GridSpec(5, 1, top = 0.9, hspace = 0, wspace = 0, height_ratios = [2,1.2,3.2,1.2,2])
    #fig, ax = plt.subplots(nrows=row, ncols=col, figsize=fs)
    ax = {}

    #shift_phase = 0
    #lower = gridspec.GridSpecFromSubplotSpec(1, 1, hspace = 0.63, subplot_spec=outer[0])

    ax['fill_up'] = plt.subplot(outer[1])
    ax['fill_up'].spines['right'].set_visible(False)
    ax['fill_up'].spines['top'].set_visible(False)
    ax['fill_up'].spines['left'].set_visible(False)
    ax['fill_up'].spines['bottom'].set_visible(False)

    ax['fill_lower'] = plt.subplot(outer[3])
    ax['fill_lower'].spines['right'].set_visible(False)
    ax['fill_lower'].spines['top'].set_visible(False)
    ax['fill_lower'].spines['left'].set_visible(False)
    ax['fill_lower'].spines['bottom'].set_visible(False)

    left = 0.15


    grid0 = gridspec.GridSpecFromSubplotSpec(3, 1, subplot_spec=outer[2], height_ratios =[2,1,1], wspace=0.02, hspace=0.2)  # ,
    ax['upper'] = plt.subplot(grid0[0])
    #start =
    #embed()
    remove_tick_marks(ax['upper'])
    #embed()
    ax['upper'].plot((eod_fe_upper - eod_fr)/eod_fr+1, (beat_corr_upper- eod_fr)/eod_fr+1, color = beat_corr_col)
    ax['upper'].set_xlabel('EOD multiples')
    ax['upper'].set_ylabel('[Hz]')
    ax['upper'].plot((eod_fe_upper - eod_fr) / eod_fr + 1,np.abs(eod_fe_upper - eod_fr) / (eod_fr), color=df_col, alpha=0.7,zorder = 1)
    #ax['upper'].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))
    #ax['upper'].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))
    ax['upper'].set_xticks(
        np.arange((eod_fe_upper[0] - eod_fr) / eod_fr + 1, (eod_fe_upper[-1] - eod_fr) / eod_fr + 1, (eod_fr / 2) / eod_fr))
    ax['upper'].set_yticks(
        np.arange((eod_fe_upper[0] - eod_fr) / eod_fr + 1, (eod_fe_upper[-1] - eod_fr) / eod_fr + 1, (eod_fr / 2) / eod_fr))
    nrs = [5]
    ax['upper'].text(-left, 1.1, string.ascii_uppercase[nrs[0]], transform=ax['upper'].transAxes,
               size=nr_size, weight='bold')




    plt.grid()
    colors = np.concatenate([['hotpink']*col,['red']*col,['darkred']*col])
    plt.xlim([(np.min(eod_fe_upper)- eod_fr)/eod_fr+1, (np.max(eod_fe_upper)- eod_fr)/eod_fr+1])
    #ax[d] = fig.add_subplot(row,col, 1)#int(np.ceil(np.sqrt(len(eod_fe))))
    #sampling_rate = sampling
    plt.ylim([-0.2,1])
    ax[0] = plt.subplot(grid0[1])
    ax[1] = plt.subplot(grid0[2])
    plot_power(a_fe, a_fr, beats_upper, beat_corr_upper, eod_fe_upper, sampling, ax, eod_fr,left = left,  nr_size = nr_size,bef_c = bef_c, aft_c = aft_c)


    upper = gridspec.GridSpecFromSubplotSpec(int(row / 2), col, hspace=0.5, subplot_spec=outer[0])
    nfft = 4096
    nfft = int(4096 * sampling / 10000 * 2)
    full = eod_fe
    colors_full = colors
    #embed()
    eod_fe = full[5:10]
    colors_plot = colors_full[5:10]
    colors = colors_full[5:10]
    colors = ['black'] * len(eod_fe)
    #colors = ['black']*len(eod_fe)
    nrs = [0, 1, 2, 3, 4, 5]
    for e in range(len(eod_fe)):
        f_max,lims,ax = single_stim(ax,colors_plot, int(row/2), col, eod_fr, eod_fe,  e, upper,s = s, p = p, d = d, df_col = df_col,  factor = factor, beat_corr_col = beat_corr_col, nfft = nfft, minus_bef = minus_bef, delta_t = delta_t, sampling = sampling,
                                 deviation = deviation,plus_bef = plus_bef, a_fr = a_fr, phase_zero = phase_zero, shift_phase = shift_phase, size = size, a_fe = a_fe)

        ax[e].text(-left, 1.1, string.ascii_uppercase[nrs[e]], transform=ax[e].transAxes,
                   size=nr_size, weight='bold')
        x = (eod_fe[e] - eod_fr)/eod_fr
        ax['upper'].scatter(x, (f_max)/eod_fr, color=colors_full[e], s=19,zorder = 2)
        ax['fill_up'].scatter(x, 1, marker = 'v',color=colors[e], s=19,zorder = 2)
        ax['fill_up'].plot([x,x], [3,1], color=colors[e], zorder = 2)
        axis = (eod_fe_upper - eod_fr) / eod_fr
        #embed()
        ax['fill_up'].set_xlim([axis[0],axis[-1]])
        ax['fill_up'].set_ylim([0.72,6])
        ax['fill_up'].set_yticks([])



    lower = gridspec.GridSpecFromSubplotSpec(int(row / 2), col, hspace=0.5, subplot_spec=outer[4])

    eod_fe = full[0:5]
    colors_plot = colors_full[0:5]

    colors_plot = colors_full[5:10]
    colors = colors_full[5:10]
    colors = ['black'] * len(eod_fe)

    nrs = [8,9, 10, 11, 12, 13, 14]
    for e in range(len(eod_fe)):
        f_max,lims,ax = single_stim(ax,colors_plot, int(row/2), col, eod_fr, eod_fe, e, lower, s = s, p = p, d = d, df_col = df_col,  factor = factor, beat_corr_col = beat_corr_col, nfft = nfft, minus_bef = minus_bef, delta_t = delta_t, sampling = sampling,
                                 deviation = deviation,plus_bef = plus_bef, a_fr = a_fr, phase_zero = phase_zero, shift_phase = shift_phase, size = size, a_fe = a_fe)
        ax[e].text(-left, 1.1, string.ascii_uppercase[nrs[e]], transform=ax[e].transAxes,
                   size=nr_size, weight='bold')
        ax['upper'].scatter((eod_fe[e] - eod_fr)/eod_fr+1, (f_max)/eod_fr, color=colors_full[e], s=19, zorder = 2)
        #ax['upper'].scatter((eod_fe[e] - eod_fr) / eod_fr, -0.17, marker = '^', color=colors[e], s=19)
        x = (eod_fe[e] - eod_fr) / eod_fr
        ax['fill_lower'].scatter(x, 4.5, marker = '^',color=colors[e], s=19, zorder = 2)
        ax['fill_lower'].plot([x,x], [1.5,4.5], color=colors[e], zorder = 2)
        axis = (eod_fe_upper - eod_fr) / eod_fr
        #embed()
        ax['fill_lower'].set_xlim([axis[0],axis[-1]])
        ax['fill_lower'].set_ylim([0.5,7.3])
        ax['fill_lower'].set_yticks([])
    #for e in range(int(len(eod_fe)/2)):
    #    f_max,lims = single_stim(factor, beats, beat_corr, plus_bef,ax,s,df_col, beat_corr_col, colors, lower, row, col, fig, nfft, p, minus_bef, delta_t, sampling,
    #                deviation, d, eod_fr, eod_fe, e, a_fr, phase_zero, shift_phase, size, sigma, a_fe)
    #    ax['upper'].scatter((eod_fe[e] - eod_fr)/eod_fr, (f_max)/eod_fr, color=colors[e], s=19)
    #plt.xlabel('Time [ms]')
    plt.legend(loc=(-5, -0.7),ncol = 4)
    plt.subplots_adjust(bottom = 0.1, top = 0.8,left = 0.13)
    #embed()
    #scalebar = ScaleBar(50,'ms')
    #plt.gca().add_artist(scalebar)
    #embed()
    # plt.xlim([-200,200])

    #plt.show()
    return fig

if __name__ == "__main__":
    delta_t = 0.014  # ms
    interest_interval = delta_t * 1.2
    bef_c = interest_interval / 2
    aft_c = interest_interval / 2
    sigma = delta_t / math.sqrt((2 * math.log(10)))  # width of the chirp
    size = [120]  # maximal frequency excursion during chirp / 60 or 100 here
    phase_zero = [0]  # phase when the chirp occured (vary later) / zero at the peak o a beat cycle
    phase_zero = np.arange(0, 2 * np.pi, 2 * np.pi / 10)

    a_fr = 1  # amplitude fish reciever
    amplitude = a_fe = 0.5  # amplitude fish emitter
    factor = 200

    eod_fr =500 # eod fish reciever
    eod_fr = 537
    eod_fr = 734
    eod_fr = 500
    initial = np.array([60, 310])
    stim_matrix = np.transpose([initial, initial + eod_fr, initial + eod_fr*2, initial + eod_fr*3, initial + eod_fr*4])
    eod_fe = np.concatenate(stim_matrix)

    sampling = eod_fr * factor
    sampling = 112683 #122300 #500
    sampling_fish = 1000
    start = 5
    step = 10
    eod_fe_upper = np.arange(start,np.int(np.max(eod_fe)/eod_fr)*eod_fr+eod_fr,step)
    beats_upper = eod_fe_upper - eod_fr
    beat_corr_upper = eod_fe_upper % eod_fr
    beat_corr_upper[beat_corr_upper > eod_fr / 2] = eod_fr - beat_corr_upper[beat_corr_upper > eod_fr / 2]

    # start = 0
    # end = 2500
    # step = 10
    win = 'w2'

    d = 1
    x = [1.5, 2.5, 0.5, ]
    x = [1.5]
    time_range = 200 * delta_t
    deviation_ms, deviation_s, deviation_dp = find_dev(x, sampling)
    #embed()
    pp = [1, 3, 6, 9]  # [2,7]

    minus_bef = -20
    plus_bef = -10
    #minus_bef = -10
    #plus_bef = -10
    bef_c = -2
    aft_c = -1
    # embed()
    # eod_fe = np.array([initial,60+500,310+500,1010,60+1000,310+1000,1010+500,60+1000+500,310+1000+500])
    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]
    row, col = np.shape(stim_matrix)
    fig = plot_eod(beat_corr_upper, beats_upper,eod_fe_upper,beat_corr, minus_bef, plus_bef, row, col, a_fe, delta_t, a_fr, size, beat_corr, 0, 0, a_fe, phase_zero,
             deviation_ms, eod_fe, deviation_dp, eod_fr, sampling, factor, fs=[6.28, 6],
             shift_phase=(2 * np.pi) / 4, bef_c = bef_c, aft_c = aft_c)  # 6.2992,6.69
    #fig.savefig()
    plt.savefig(
        fname= 'simulationexamples.pdf')
    plt.savefig('../highbeats_pdf/simulationexamples.pdf')
    plt.show()