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\\2020lepto04'

datasets = ['2020-07-22-ab',
            '2020-07-22-ac',
            '2020-07-22-ad',
            '2020-07-22-ae',
            '2020-07-22-af',
            '2020-07-22-ag',
            '2020-07-23-ab',
            '2020-07-23-ac',
            '2020-07-23-ad',
            '2020-07-23-ae']

time_all = []
freq_all = []

amfrequencies = []
gains = []
files = []

ID = []

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):
        file_name = []

        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:-1])

        # file_name
        file_name.append(ID[0])

        amfreq = import_amfreq(datapath)
        print(amfreq)
        file_name.append(str(amfreq))

        file_name.append(str(d))
        files.append(file_name)

        # spectogram
        if float(amfreq) < 0.01:
            spec, freqs, times = specgram(dat, Fs=1/dt, detrend='mean', NFFT=nfft, noverlap=nfft * 0.8)
        else:
            spec, freqs, times = specgram(dat, Fs=1/dt, detrend='mean', NFFT=nfft, noverlap=nfft * 0.95)

        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

        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)   # data we take
        cut_times = times[:len(jar4)]

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

        # save data
        np.save('%s time' % file_name, cut_times)
        np.save('%s' % file_name, jar4)

# save filenames for this fish
np.save('%s files' %ID[0], files)
embed()

# running average over on AM-period?