
from matplotlib.colors import LinearSegmentedColormap
#from plot_eod_chirp  import find_times
import matplotlib.pyplot as plt
import numpy as np
from IPython import embed
#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
from functionssimulation import single_stim
from plot_eod_chirp  import rectify,find_dev, conv,global_maxima,integrate_chirp,find_periods,find_lm,pl_eods
from plot_eod_chirp import find_beats,snip, power_func
import os
from thunderfish.tabledata import TableData
import string
from plotstyle import plot_style, spines_params
from myfunctions import remove_tick_ymarks
from matplotlib import gridspec


def plot_beats(ax, f0, freqs, fexpected, ffirst, 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 = 'black', linestyle = 'dashed',linewidth = 0.8)
    if ffirst is not None:
        ffirst[sel] = np.nan
        ax.plot(freqs, ffirst/f0, color = 'gold')
    if fmax is not None:
        fmax[sel] = np.nan
        ax.plot(freqs, fmax/f0, color = 'orange')
    ax.set_xlim(0, 5)
    ax.set_ylim(0, 0.7)
    ax.set_xticks_delta(1.0)

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

def plot_power(beats,  results,  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,bef_c = 1.1,aft_c =-0.1, share = False):
    #embed()
    plt.rcParams['figure.figsize'] = (6.27, 8)
    plt.rcParams["legend.frameon"] = False
    colors = ['black','blue','purple','magenta','pink','orange', 'brown', 'red', 'green','lime']
    #colors = ['brown']*len(name)
    default_settings([0], intermediate_width=6.29, intermediate_length=8, ts=10, ls=9, fs=9)

    all_max = [[]] * len(results)
    first = ['hilbert>first_f','threshold>first_f','rectified>first_f','square>first_f','threshold cubed>first_f']
    max = ['hilbert>max_f', 'threshold>max_f', 'rectified>max_f', 'square>max_f', 'threshold cubed>max_f']
    maxa = ['hilbert>max_a', 'threshold>max_a', 'rectified>max_a', 'square>max_a', 'threshold cubed>max_a']

    name = ['Hilbert','Thresholded','Rectified','Squaring','Threshold cubed']
    colors = ['brown'] * len(name)
    #fig, ax = plt.subplots(nrows=len(name), ncols=2)
    ax = {}
    grid0 = gridspec.GridSpec(len(name), 1)
    nrs1 = [0,2,4,6,8]
    nrs2 = [1,3,5,7,9]
    for i in range(len(name)):

        data = TableData('highbeatspecs10.dat')

        #ax['analytic_threhold'] = plt.subplot(axis[0])
        nr_size = 11
        left = -0.1
        grid = gridspec.GridSpecFromSubplotSpec(2,2,subplot_spec=grid0[i])
        #embed()
        ax[0] = plt.subplot(grid[0])
        ax[0].text(left, 1.1, string.ascii_uppercase[nrs1[i]], transform=ax[0].transAxes,
                                     size=nr_size, weight='bold')
        f0 = data[data[:, 'freq>norm'] == 1.0, 'freq>freq']
        f0 = f0[len(f0) // 2]
        plot_beats(ax[0], f0, data[:, 'freq>freq'], data[:, 'freq>expected'],
                   data[:,first[i]], data[:, max[i]], name[i]+' analytic')
        ax[1] = plt.subplot(grid[1])

        #ax[i, 1].set_ylabel(results['type'][i], rotation=0, labelpad=40, color=colors[i])
        ax[1].plot(beats / eod_fr +1, np.array(results['result_frequency' ][i]) / eod_fr, color=colors[i])
        # plt.title(results['type'][i])
        #ax[i, 1].plot(beats / eod_fr + 1, np.array(results['amp'][i]), color=colors[i])
        ax[1].set_title(name[i]+' numerical')
        ax[1].set_xticks_delta(1.0)
        ax[1].plot(data[:, 'freq>freq']/ f0, data[:, 'freq>expected'] / f0, color='black', linestyle='dashed',linewidth = 0.8)
        ax[1].set_xlim(0, 5)
        ax[1].set_ylim(0, 0.7)
        ax[1] = remove_tick_ymarks(ax[1])

        ax[1].text(left, 1.1, string.ascii_uppercase[nrs2[i]], transform=ax[1].transAxes,
                                     size=nr_size, weight='bold')

        ax[2] = plt.subplot(grid[2])
        #embed()
        ax[2].plot(data[:, 'freq>freq'] / f0, (data[:, maxa[i]]-np.nanmin(data[:, maxa[i]]))/np.nanmax((data[:, maxa[i]]-np.nanmin(data[:, maxa[i]]))), color='blue')
        ax[2].set_ylabel('r. amp.')
        ax[3] = plt.subplot(grid[3])
        ax[3].plot(beats / eod_fr + 1, np.array(results['result_amplitude_max'][i])/(np.max(results['result_amplitude_max'][i])), color='navy')
        ax[3].set_xlim(0, 5)
        ax[3] = remove_tick_ymarks(ax[3])
        ax[2].set_xlim(0, 5)
        ax[1] = remove_tick_marks(ax[1])
        ax[0] = remove_tick_marks(ax[0])
        if i != (len(name)-1):
            ax[2] = remove_tick_marks(ax[2])
            ax[3] = remove_tick_marks(ax[3])
        for i in range(4):
            ax[i].spines['right'].set_visible(False)
            ax[i].spines['top'].set_visible(False)
        #ax[i, 1].plot(beats / eod_fr + 1, np.array(results['result_amplitude_max'][i]), color=colors[i])
        #ax[i, 2].plot(beats / eod_fr + 1, np.array(results['amp'][i]), color=colors[i])
        #ax[i,1].set_ylim([1,1.6])
        all_max[i] = np.max(np.array(results['result_amplitude_max'][i]))
    ax[3].set_xlabel('EOD multiples')
    ax[2].set_xlabel('EOD multiples')
    #for i in range(len(results)):
    #    ax[i, 2].set_ylim([0, np.max(all_max)])
    plt.subplots_adjust(left=0.13, hspace = 0.45, top = 0.94, right = 0.97)
    #ii, jj = np.shape(ax)
    #ax[0, 0].set_title('Most popular frequency')
    #ax[0, 1].set_title('Modulation depth')
    #ax[0, 2].set_title('Modulation depth (same scale)')
    #for i in range(ii):#
#
 #       for j in range(jj):
   #         ax[-1, j].set_xlabel('EOD multiples')
  #          ax[i, j].spines['right'].set_visible(False)
    #        ax[i, j].spines['top'].set_visible(False)

    def onclick(event):
        #embed()
        eod_fe = [(event.xdata-1)*eod_fr]
        nfft = 4096
        e = 0
        #sampling = 121000
        left_b = [bef_c * sampling] * len(beat_corr)
        right_b = [aft_c * sampling] * len(beat_corr)
        p = 0
        s = 0
        d = 0
        time_fish,sampled, cut, cubed, 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, chirp=False)
        #cubed, 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 = onclick_func(e,nfft,beats,  results,        disc, win, deviation_s, sigma,  d_ms, beat_corr, size, phase_zero, delta_t, a_fr, a_fe, eod_fr, eod_fe, deviation, show_figure = show_figure, plot_dist = plot_dist, save = save,bef_c = bef_c,aft_c =aft_c, sampling = sampling)

        nfft = 4096

        results = [[]] * 1
        name = ['cubed']
        var = [cubed]
        var = [maxima_interp_b]
        samp = [sampling]
        nfft = int((4096 * samp[0] / 10000) * 2)
        i = 0
        pp, f = ml.psd(var[i] - np.mean(var[i]), Fs=samp[i], NFFT=nfft,
                       noverlap=nfft / 2)
        plt.figure()
        plt.subplot(1,2,1)
        plt.plot(f, pp)
        plt.xlim([0,2000])
        #plt.subplot(1,3,2)
        #plt.plot(time_b, cubed)
        plt.subplot(1,2,2)
        plt.plot(time_b, maxima_interp_b)
        plt.show()
    if share == True:
        cid = fig.canvas.mpl_connect('button_press_event', onclick)
    plt.savefig('numerical_compar_both.pdf')
    plt.show()

def onclick_func(e,nfft, beats,  results,        disc, win, deviation_s, sigma,  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,bef_c = 1.1,aft_c =-0.1,sampling = 100000):
    left_b = [bef_c * sampling] * len(beat_corr)
    right_b = [aft_c * sampling] * len(beat_corr)
    p = 0
    s = 0
    time_fish,sampled, cut, cubed, 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, chirp=False)
    return pp, f, cubed, cubed, 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





delta_t = 0.014
interest_interval = delta_t * 1.2
sigma = delta_t / math.sqrt((2 * math.log(10)))  # width of the chirp
phase_zero = np.arange(0,2*np.pi,2*np.pi/10) #phase_zero = [0]  # phase when the chirp occured (vary later) / zero at the peak o a beat cycle
eod_fr = 637# eod fish reciever
eod_fr = 537
eod_fr = 1435
eod_fr = 734
factor = 200
sampling_fish = 500
step = 500
win = 'w2'
d = 1
x = [ 1.5]#x = [ 1.5, 2.5,0.5,]
time_range = 200 * delta_t
sampling = 112345
deviation_ms, deviation_s, deviation_dp = find_dev(x, sampling)
start = 5
end = 3500
step = 25
eod_fe, beat_corr, beats = find_beats(start,end,step,eod_fr)
delta_t = 1
load = True
size = [120]
a_fr = 1
a_fe = 0.5
bef_c = -2
aft_c = -1
load = False
restrict =  np.array([0,2,3,4,5])
if (load == True) or (not os.path.exists('numerical.pkl')):
    results = power_func(restrict = restrict, a_fr = a_fr, a_fe = a_fe, eod_fr = eod_fr, eod_fe = eod_fe, win = 'w2', deviation_s = deviation_s, sigma = sigma,  sampling = sampling,  deviation_ms = deviation_ms, beat_corr = beat_corr, size = size,phase_zero =  [phase_zero[0]], delta_t = delta_t,deviation_dp = deviation_dp, show_figure = True, plot_dist = False, save = False,bef_c = bef_c,aft_c = aft_c)
    results = pd.DataFrame(results)
    results.to_pickle('numerical.pkl')
    np.save('numerical.npy', results)
else:
    results = pd.read_pickle('numerical.pkl')
#embed()

plot_power(beats, results,'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,bef_c = bef_c,aft_c =aft_c)
