import matplotlib.pyplot as plt
import numpy as np
from read_chirp_data import *
from read_baseline_data import *
from utility import *
from IPython import embed

# define data path and important parameters
data_dir = "../data"
sampling_rate = 40 #kHz
cut_window = 100
cut_range = np.arange(-cut_window * sampling_rate, 0, 1)
window = 1

'''
# norm: -150, 150, 300  aa, #ac, aj??
data = ["2018-11-13-aa-invivo-1", "2018-11-13-ad-invivo-1", "2018-11-13-ah-invivo-1", "2018-11-13-ai-invivo-1",
        "2018-11-13-ak-invivo-1", "2018-11-13-al-invivo-1"]

# norm: -50
data = ["2018-11-20-aa-invivo-1", "2018-11-20-ab-invivo-1", "2018-11-20-ac-invivo-1","2018-11-20-ad-invivo-1", 
        "2018-11-20-ae-invivo-1", "2018-11-20-af-invivo-1", "2018-11-20-ag-invivo-1", "2018-11-20-ah-invivo-1",
        "2018-11-20-ai-invivo-1"]
'''

data = ["2018-11-14-ad-invivo-1", "2018-11-14-ak-invivo-1", "2018-11-14-am-invivo-1"]

#data = ["2018-11-14-ac-invivo-1", "2018-11-14-ad-invivo-1", "2018-11-14-ak-invivo-1"]
#data = ["2018-11-09-ad-invivo-1", "2018-11-14-af-invivo-1"]

#data = ["2018-11-20-ad-invivo-1", "2018-11-13-ad-invivo-1"]
#data = ["2018-11-09-ad-invivo-1"]

rates = {}

for dataset in data:
    rates[dataset] = {}
    print(dataset)
    # read baseline spikes
    base_spikes = read_baseline_spikes(os.path.join(data_dir, dataset))
    base_spikes = base_spikes[1000:2000]
    spikerate = len(base_spikes)/base_spikes[-1]
    print(spikerate)

    # read spikes during chirp stimulation
    spikes = read_chirp_spikes(os.path.join(data_dir, dataset))
    df_map = map_keys(spikes)

    # iterate over df
    for df in df_map.keys():
        '''
        if df == 50:
            pass
        else:
            continue
        '''

        #print(df)
        rep_rates = []
        beat_duration = int(abs(1 / df) * 1000)
        beat_window = 0
        while beat_window + beat_duration <= cut_window/2:
            beat_window = beat_window + beat_duration
        for rep in df_map[df]:
            for phase in spikes[rep]:
                # get spikes 40 ms before the chirp first chirp
                spikes_to_cut = np.asarray(spikes[rep][phase])
                spikes_cut = spikes_to_cut[(spikes_to_cut > -cut_window) & (spikes_to_cut < 0)]
                spikes_idx = np.round(spikes_cut * sampling_rate)
                # also save as binary, 0 no spike, 1 spike
                binary_spikes = np.isin(cut_range, spikes_idx) * 1
                smoothed_data = smooth(binary_spikes, window, 1 / sampling_rate)
                #train = smoothed_data[window:beat_window+window]
                #norm_train = train*1000/spikerate
                #df_rate = np.std(norm_train)
                #rates[dataset][df] = [df_rate]
                rep_rates.append(np.std(smoothed_data))#/spikerate)
                '''
                if df in rates[dataset].keys():
                    rates[dataset][df].append(np.std(norm_train))
                else:
                    rates[dataset][df] = [np.std(norm_train)]
                '''
                break
            #break
        df_rate = np.mean(rep_rates)
        #df_rate = rep_rates
        rates[dataset][df] = df_rate

        #embed()
        #exit()
        '''
        if df in rates.keys():
            rates[dataset][df].append(df_rate)
        else:
            rates[dataset][df] = [df_rate]
        '''

colors = ['royalblue', 'red', 'green', 'violet', 'orange', 'black', 'gray']

fig, ax = plt.subplots()
for i, cell in enumerate(rates.keys()):
    for j, df in enumerate(sorted(rates[cell].keys())):
        ax.plot(df, rates[cell][df], 'o', color=colors[i])
#ax.legend(sorted(rates.keys()), loc='upper left', bbox_to_anchor=(1.04, 1))
fig.tight_layout()
plt.show()

'''
fig, ax = plt.subplots()
for i, cell in enumerate(rates.keys()):
    for j, df in enumerate(sorted(rates[cell].keys())):
        ax.plot(np.ones(len(rates[cell][df]))*df, rates[cell][df], 'o', color=colors[i])
#ax.legend(sorted(rates.keys()), loc='upper left', bbox_to_anchor=(1.04, 1))
fig.tight_layout()
plt.show()
'''