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
datasets = ["2018-11-13-ad-invivo-1", "2018-11-13-aj-invivo-1", \
            "2018-11-13-ak-invivo-1", "2018-11-14-ad-invivo-1"]
# dataset = "2018-11-20-af-invivo-1"

fig = plt.figure()
# figsize=(5 / inch_factor, 2.5 / inch_factor))
fig.set_size_inches((35/inch_factor, 15/inch_factor))
axes = []
axes.append(fig.add_subplot(221))
axes.append(fig.add_subplot(222))
axes.append(fig.add_subplot(223))
axes.append(fig.add_subplot(224))

for dataset, ax in zip(datasets, axes):
    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)
                modulation = np.std(smoothed_data)
                rates[deltaf][x] = modulation
                break

    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_xlim(-450, 800)
    ax.set_ylim(0.1, 0.3)
    ax.set_xticks(np.arange(-400, 810, 100))
    ax.yaxis.set_tick_params(labelsize=18)
    ax.xaxis.set_tick_params(labelsize=18)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

axes[0].set_ylabel('Firing rate modulation', fontsize=22)
axes[0].yaxis.set_label_coords(-0.15, -.125)
axes[2].set_xlabel('$\Delta$f [Hz]', fontsize=22)
axes[3].set_xlabel('$\Delta$f [Hz]', fontsize=22)
axes[1].set_yticklabels([])
axes[3].set_yticklabels([])
axes[0].set_xticklabels([])
axes[1].set_xticklabels([])
axes[2].set_xticklabels(['-400', '', '-200', '', '0', '', '200', '', '400', '', '600', '', '800'], rotation=45)
axes[3].set_xticklabels(['-400', '', '-200', '', '0', '', '200', '', '400', '', '600', '', '800'], rotation=45)

fig.subplots_adjust(left=0.09, bottom=0.175, right=0.975, top=0.95)
fig.savefig('spikes_beat_20af.png')