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

# define sampling rate and data path
sampling_rate = 40 #kHz
data_dir = "../data"
dataset = "2018-11-09-ad-invivo-1"
# parameters for binning, smoothing and plotting
chirp_size = 14 #ms
neuronal_delay = 5 #ms
chirp_start = int((-chirp_size/2+neuronal_delay+50)*sampling_rate)
chirp_end = int((chirp_size/2+neuronal_delay+51)*sampling_rate)
num_bin = 12
window = 1 #ms
time_axis = np.arange(-50, 50, 1/sampling_rate) #steps
spike_bins = np.arange(-50, 51) #ms

# read data from files
spikes = read_chirp_spikes(os.path.join(data_dir, dataset))
eod = read_chirp_eod(os.path.join(data_dir, dataset))
chirp_times = read_chirp_times(os.path.join(data_dir, dataset))

# make a delta f map for the quite more complicated keys
df_map = {}
for k in spikes.keys():
    df = k[1]
    if df in df_map.keys():
        df_map[df].append(k)
    else:
        df_map[df] = [k]

# differentiate between phases
phase_vec = np.arange(0, 1+1/num_bin, 1/num_bin)
cut_range = np.arange(-50*sampling_rate, 50*sampling_rate, 1)

# make dictionaries for spiketimes
df_phase_time = {}
df_phase_binary = {}

# iterate over delta f, repetition, phases and a single chirp
for deltaf in df_map.keys():
    df_phase_time[deltaf] = {}
    df_phase_binary[deltaf] = {}
    for rep in df_map[deltaf]:
        for phase in spikes[rep]:
            for idx in np.arange(num_bin):
                # check the phase
                if phase[1] > phase_vec[idx] and phase[1] < phase_vec[idx+1]:

                    # get spikes between 50 ms befor and after the chirp
                    spikes_to_cut = np.asarray(spikes[rep][phase])
                    spikes_cut = spikes_to_cut[(spikes_to_cut > -50) & (spikes_to_cut < 50)]
                    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

                    # add the spikes to the dictionaries with the correct df and phase
                    if idx in df_phase_time[deltaf].keys():
                        df_phase_time[deltaf][idx].append(spikes_cut)
                        df_phase_binary[deltaf][idx] = np.vstack((df_phase_binary[deltaf][idx], binary_spikes))
                    else:
                        df_phase_time[deltaf][idx] = [spikes_cut]
                        df_phase_binary[deltaf][idx] = binary_spikes

# for plotting and calculating iterate over delta f and phases
for df in df_phase_time.keys():
    for phase in df_phase_time[df].keys():

        trials_binary = df_phase_binary[df][phase]

        sr_chirp = np.zeros(len(trials_binary))
        sr_beat = np.zeros(len(trials_binary))
        for i, trial in enumerate(trials_binary):
            smoothed_trial = smooth(trial, window, 1/sampling_rate)
            sr_chirp[i] = np.mean(smoothed_trial[chirp_start:chirp_end])
            sr_beat[i] = np.mean(smoothed_trial[0:chirp_start])

        for rate_chirp in sr_chirp:
            for rate_beat in sr_beat:
                r = np.corrcoef(rate_chirp, rate_beat)
                print(r)

        embed()
        exit()

        #csi = (spikerate_chirp-spikerate_befor)/(spikerate_chirp+spikerate_befor)

        # plot
        #plot_trials = df_phase_time[df][phase]
        #plot_trials_binary = np.mean(df_phase_binary[df][phase], axis=0)

        # calculation
        #overall_spikerate = (np.sum(plot_trials_binary)/len(plot_trials_binary))*sampling_rate*1000

        '''
        smoothed_spikes = smooth(plot_trials_binary, window, 1./sampling_rate)

        fig, ax = plt.subplots(2, 1, sharex=True)
        for i, trial in enumerate(plot_trials):
            ax[0].scatter(trial, np.ones(len(trial))+i, marker='|', color='k')
        ax[1].plot(time_axis, smoothed_spikes*1000)

        ax[0].set_title(df)
        ax[0].set_ylabel('repetition', fontsize=12)

        ax[1].set_xlabel('time [ms]', fontsize=12)
        ax[1].set_ylabel('firing rate [Hz]', fontsize=12)
        plt.show()
        '''


'''
for trial in range(len(trials_binary)):
spike_rate = np.zeros(len(spike_bins)-1)
for idx in range(len(spike_bins)-1):
bin_start = spike_bins[idx]*sampling_rate
bin_end = spike_bins[idx+1]*sampling_rate
spike_rate[idx] = np.sum(trials_binary[trial][bin_start:bin_end])
embed()
exit()
'''