import matplotlib.pyplot as plt
import matplotlib as cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import os
import glob
import IPython
import numpy as np
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\\step\\step_2018lepto98'

#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-13-ad',
'2020-07-13-ae',
'2020-07-13-af',
'2020-07-13-ag',
'2020-07-13-ah',
'2020-07-13-ai',
'2020-07-13-aj',
#'2020-07-13-ak',
#'2020-07-13-al',
'2020-07-13-am',
#'2020-07-13-an',
#'2020-07-13-ao'
]

#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):
    dataset = os.path.join(base_path, dataset, 'beats-eod.dat')
    #input of the function
    frequency, time, amplitude, eodf, deltaf, stimulusf, duration, pause = parse_dataset(dataset)
    dm = np.mean(duration)
    pm = np.mean(pause)
    timespan = dm + pm
    start = np.mean([t[0] for t in time])
    stop = np.mean([t[-1] for t in time])

    norm = norm_function(frequency, time, onset_point=dm - dm, offset_point=dm)  # dm-dm funktioniert nur wenn onset = 0 sec

    mf, tnew = mean_traces(start, stop, timespan, norm, time) # maybe fixed timespan/sampling rate

    cf, ct = mean_noise_cut(mf, n=1250)

    cf_arr = np.array(cf)
    ct_arr = np.array(ct)

    freq_all.append(cf_arr)
    time_all.append(ct_arr)

    plt.plot(ct_arr, cf_arr, label='fish=%s' % datasets[idx]) #, color = col[idx]

    sv, sc = curve_fit(step_response, ct_arr[ct_arr < dm], cf_arr[ct_arr < dm], [1.0, 1.0, 5.0, 50.0], bounds=(0.0, np.inf))  # step_values and step_cov

    # sorted a and tau
    values = sort_values(sv)

    # fit for each trace 
    #plt.plot(ct_arr[ct_arr < dm], step_response(ct_arr[ct_arr < dm], *sv), label='fit: a1=%.2f, a2=%.2f, tau1=%.2f, tau2=%.2f' % tuple(values))
    plt.plot(ft, step_response(ft, *sv), color='orange', label='fit: a1=%.2f, a2=%.2f, tau1=%.2f, tau2=%.2f' % tuple(values))

    print('fish: a1, a2, tau1, tau2', values)


# average over all fish
mf_all, tnew_all, values_all = average(freq_all, time_all, start, stop, timespan, dm)


#const_line = plt.axhline(y = 0.632)
stimulus_duration = plt.hlines(y = -0.25, xmin = 0, xmax = 100, color = 'r', label = 'stimulus_duration')
base_line = plt.axhline(y = 0, color = 'black', ls = 'dotted', linewidth = '1')

plt.xlim([-10,220])
plt.xlabel('time [s]')
plt.ylabel('rel. JAR magnitude')
plt.title('relative JAR')
plt.savefig('relative JAR')
plt.legend(loc = 'lower right')
plt.show()
embed()


# natalie fragen ob sie bei verschiedenen Amplituden messen kann (siehe tim)