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 functionssimulation import find_times
from functionssimulation import find_periods
from functionssimulation import integrate_chirp
from functionssimulation import rectify, find_beats,find_dev
from functionssimulation import global_maxima, find_lm, conv

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):
    time, time_cut, cut = find_times(left_c[g], right_c[g], sampling, deviation_s[d])
    eod_fish_r, period_fish_r, period_fish_e = find_periods(time, eod_fr, a_fr, eod_fe, e)
    #embed()
    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]

    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 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_func(bef_c, aft_c, win, deviation_s, sigma, sampling, d_ms, beat_corr, size, phase_zero, delta_t, a_fr, a_fe, eod_fr, eod_fe, deviation, show_figure = False, plot_dist = False, save = False):
    results = [[]]*7
    for d in range(len(deviation)):
        bef_c = 0.3
        aft_c = -0.1
        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)):
                    left_b = [-0.3*sampling]*len(beat_corr)
                    right_b = [-0.1 * sampling]*len(beat_corr)
                    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
                    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]
                    #pp, f = ml.psd(eod_overlayed_chirp - np.mean(eod_overlayed_chirp), Fs=sampling, NFFT=nfft, noverlap=nfft / 2)
                    for i in range(len(results)):


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

                        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()
                        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_fr])]])
                            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_fr])*np.abs(f[1]-f[0]))])
                        else:
                            results[i]['f'].extend(list([f[np.argmax(pp[f < 0.5 *eod_fr])]]))
                            #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_fr]) * np.abs(f[1] - f[0]))]))
                #if save:
                #    results = pd.DataFrame(results)
                #    results.to_pickle('../data/power_simulation.pkl')
                #    np.save('../data/Ramona/power_simulation.npy', results)
    return results


def plot_power(results):
    plt.rcParams['figure.figsize'] = (3, 3)
    plt.rcParams["legend.frameon"] = False
    colors = ['black', 'magenta', 'pink', 'orange', 'moccasin', 'red', 'green', 'silver']
    colors = ['red','pink']
    results = [results[5]]
    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])
        ax[0].plot(beats / eod_fr + 1, np.array(results[i]['f']) / eod_fr + 1, color=colors[i])
        # 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[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('MPF')
    ax[1].set_ylabel('Modulation depth')
    #ax[0, 2].set_title('Modulation depth (same scale)')
    for i in range(len(ax)):
        ax[1].set_xlabel('EOD multiples')
        ax[i].spines['right'].set_visible(False)
        ax[i].spines['top'].set_visible(False)
    plt.subplots_adjust(bottom = 0.2)
    plt.savefig('localmaxima.pdf')
    plt.savefig('../highbeats_pdf/localmaxima.pdf')
    plt.show()

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)
eod_fr = 500  # eod fish reciever
a_fr = 1  # amplitude fish reciever
amplitude = a_fe = 0.2  # amplitude fish emitter
factor = 200
sampling = eod_fr * factor
sampling_fish = 500
#start = 0
#end = 2500
#step = 10
start = 510
end = 3500
step = 500
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)
start = 5
end = 2500
step = 25
eod_fe, beat_corr, beats = find_beats(start,end,step,eod_fr)
results = power_func( bef_c, aft_c, 'w2', deviation_s, sigma, sampling, deviation_ms, beat_corr, size, [phase_zero[0]], delta_t, a_fr, a_fe, eod_fr, eod_fe, deviation_dp, show_figure = True, plot_dist = False, save = True)
plot_power(results)