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 tqdm import tqdm
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\\eigenmannia\\sin\\2015eigen8'

time_all = []
freq_all = []

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

for idx, dataset in enumerate(os.listdir(base_path)):
    if dataset == 'prerecordings':
        continue
    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):
        if len(dat) == 1:
            print(datapath)

        file_name = []
        ID = []

        # identifier for file_name
        infodatapath = os.path.join(base_path, dataset, 'info.dat')
        i = parse_infodataset(infodatapath)
        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)
        plt.plot(cut_times, jar4)
        plt.show()

        # 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()