import matplotlib.pyplot as plt
import matplotlib as cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
from matplotlib.mlab import specgram
import os
import glob
import IPython
import numpy as np
import DataLoader as dl
from IPython import embed
from scipy.optimize import curve_fit
from jar_functions import parse_dataset
from jar_functions import parse_infodataset
from jar_functions import mean_traces
from jar_functions import mean_noise_cut
from jar_functions import norm_function
from jar_functions import step_response
from jar_functions import sort_values
from jar_functions import average

base_path = 'D:\\jar_project\\JAR'

#nicht: -5Hz delta f, 19-aa, 22-ae, 22-ad (?)
datasets = [#'2020-06-19-aa',   #-5Hz delta f, horrible fit
            #'2020-06-19-ab',   #-5Hz delta f, bad fit
            #'2020-06-22-aa',   #-5Hz delta f, bad fit
            #'2020-06-22-ab',   #-5Hz delta f, bad fit
            #'2020-06-22-ac',   #-15Hz delta f, good fit
            #'2020-06-22-ad',   #-15Hz delta f, horrible fit
            #'2020-06-22-ae',   #-15Hz delta f, horrible fit
            #'2020-06-22-af',    #-15Hz delta f, good fit
            '2020-07-21-ak'      #sin
            ]

#dat = glob.glob('D:\\jar_project\\JAR\\2020*\\beats-eod.dat')
#infodat = glob.glob('D:\\jar_project\\JAR\\2020*\\info.dat')

time_all = []
freq_all = []

ID = []
col = ['dimgrey', 'grey', 'darkgrey', 'silver', 'lightgrey', 'gainsboro', 'whitesmoke']
labels = zip(ID, datasets)

for infodataset in datasets:
    infodataset = os.path.join(base_path, infodataset, 'info.dat')
    i = parse_infodataset(infodataset)
    identifier = i[0]
    ID.append(identifier)


for idx, dataset in enumerate(datasets):
    datapath = os.path.join(base_path, dataset)
    for info, key, time, data in dl.iload_traces(datapath, repro='Beats', before=0.0, after=0.0):

        #plt.plot(time, data[0])  # V(t)
        #plt.show()
        nfft = 2**18
        spec, freqs, times = specgram(data[0], Fs=1.0/(time[1]-time[0]), detrend='mean', NFFT=nfft, noverlap=nfft*0.8)

        dbspec = 10.0*np.log10(spec)  # in dB

        power = dbspec[:, 10]

        fish_p = power[(freqs > 400) & (freqs < 1000)]
        fish_f = freqs[(freqs > 400) & (freqs < 1000)]
        index = np.argmax(fish_p)
        eodf = fish_f[index]
        eodf4 = eodf * 4
        print(eodf4)
        lim0 = eodf4-20
        lim1 = eodf4+20

        plt.imshow(dbspec, cmap='jet', origin='lower', extent=(times[0], times[-1], freqs[0], freqs[-1]), aspect='auto', vmin=-80, vmax=-30)

        # control of plt.imshow
        df = freqs[1] - freqs[0]
        ix = int(np.round(eodf4/df))
        ix0 = int(np.round(lim0/df))
        ix1 = int(np.round(lim1/df))
        spec4 = dbspec[ix0:ix1, :]
        freq4 = freqs[ix0:ix1]
        jar4 = freq4[np.argmax(spec4, axis=0)]
        jar = jar4 / 4

        plt.plot(times, jar4, '-k')

        #plt.plot(freqs, dbspec[:,100], label = '0')
        #plt.legend()
        plt.ylim(lim0, lim1)
        plt.show()
        embed()