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
inch_factor = 2.54
#dataset = "2018-11-13-ad-invivo-1"
#dataset = "2018-11-13-aj-invivo-1"
#dataset = "2018-11-13-ak-invivo-1" #al
#dataset = "2018-11-14-ad-invivo-1"
dataset = "2018-11-20-af-invivo-1"

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)

rates = {}
# iterate over df
for deltaf in df_map.keys():
    rates[deltaf] = {}
    beat_duration = int(abs(1 / deltaf) * 1000)
    beat_window = 0
    while beat_window + beat_duration <= cut_window/2:
        beat_window = beat_window + beat_duration
    for x, repetition in enumerate(df_map[deltaf]):
        for phase in spikes[repetition]:
            # get spikes some ms before the chirp first chirp
            spikes_to_cut = np.asarray(spikes[repetition][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*sampling_rate:beat_window*sampling_rate+window*sampling_rate]
            modulation = np.std(smoothed_data)
            rates[deltaf][x] = modulation
            break

plt.close()
fig = plt.figure()
#figsize=(5 / inch_factor, 2.5 / inch_factor))
fig.set_size_inches((5, 2.5))
ax = fig.add_subplot(111)
for i, df in enumerate(sorted(rates.keys())):
    max_rep = len(sorted(rates[df].keys()))-1
    for j, rep in enumerate(rates[df].keys()):
        if j == max_rep:
            farbe = 'royalblue'
            gro = 8
        else:
            farbe = 'k'
            gro = 6
        ax.plot(df, rates[df][rep], marker='o', color=farbe, ms=gro)

#ax.set_xlabel('difference frequency [Hz]', fontsize=20)
ax.set_ylabel('firing modulation', fontsize=16)
ax.set_xlabel('$\Delta$f [Hz]', fontsize=16)

ax.yaxis.set_tick_params(labelsize=14)
ax.xaxis.set_tick_params(labelsize=14)

ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
fig.tight_layout()
fig.savefig('spikes_beat_20af.png')