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 step_response
from jar_functions import sin_response
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 sort_values
from jar_functions import average
from jar_functions import import_data
from jar_functions import import_amfreq

base_path = 'D:\\jar_project\\JAR\\sin\\2018lepto98'

#nicht: -5Hz delta f, 19-aa, 22-ae, 22-ad (?)

#dat = glob.glob('D:\\jar_project\\JAR\\2020*\\beats-eod.dat')
#infodat = glob.glob('D:\\jar_project\\JAR\\2020*\\info.dat')
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-al',      #sin
            '2020-07-21-am',
            '2020-07-21-ak',
            #'2020-07-21-an',
            #'2020-07-21-ao'
            ]

time_all = []
freq_all = []

amfrequencies = []
gains = []

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]
    if not identifier[1:-2] in ID:
        ID.append(identifier[1:-2])


for idx, dataset in enumerate(datasets):
    datapath = os.path.join(base_path, dataset, '%s.nix' % dataset)
    print(datapath)

    data, pre_data, dt = import_data(datapath)


    nfft = 2**17

    for d, dat in enumerate(data):
        amfreq = float(import_amfreq(datapath))
        print(amfreq)
        amfrequencies.append(amfreq)

        file_name = []
        file_name.append(ID)
        file_name.append(amfreq)
        file_name.append(d)

        spec, freqs, times = specgram(dat, Fs=1/dt, detrend='mean', NFFT=nfft, noverlap=nfft*0.6)

        dbspec = 10.0*np.log10(spec)  # in dB
        power = dbspec[:, 50]

        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

        lim0 = eodf4-10
        lim1 = eodf4+15

        # control of plt.imshow
        df = freqs[1] - freqs[0]
        ix0 = int(np.floor(lim0/df))    # back to index
        ix1 = int(np.ceil(lim1/df))    # back to index
        spec4 = dbspec[ix0:ix1, :]
        freq4 = freqs[ix0:ix1]
        jar4 = freq4[np.argmax(spec4, axis=0)]      # all freqs at max specs over axis 0
        jar = jar4 / 4
        jm = jar4 - np.mean(jar4)

        cut_times = times[:len(jar4)]
        #np.save('%s.npy' % file_name, jar4)

        plt.plot(cut_times, jm, '-k')

        cf, ct = mean_noise_cut(jar4, cut_times, n = int(round(len(jar4)/((times[-1] - times [0]) * amfreq))))

        #plt.plot(ct, cf, '-k')

        #plt.imshow(spec4, cmap='jet', origin='lower', extent=(times[0], times[-1], lim0, lim1), aspect='auto', vmin=-80, vmax=-10)
        #np.save( , spec4)

        embed()
        b, a = signal.butter(4, 0.01 / 10000, 'high', analog=True)
        y = signal.filtfilt(b, a, jm)
        sinv, sinc = curve_fit(sin_response, cut_times, jm, [amfreq, 2, 0.5])
        print('frequency, phaseshift, amplitude:', sinv)
        gains.append(sinv[2])

        plt.plot(cut_times, sin_response(cut_times, *sinv), label='fit: f=%f, p=%.2f, A=%.2f' % tuple(sinv))

        #plt.legend()
        #plt.ylim(lim0, lim1)
        plt.legend()
        plt.show()

        #print(np.load('%s.npy' % file_name))
embed()