import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as ss
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-13-al-invivo-1"
'''
data = ["2018-11-09-ad-invivo-1", "2018-11-09-ae-invivo-1", "2018-11-09-ag-invivo-1", "2018-11-13-aa-invivo-1",
    "2018-11-13-ac-invivo-1", "2018-11-13-ad-invivo-1", "2018-11-13-ah-invivo-1", "2018-11-13-ai-invivo-1",
    "2018-11-13-aj-invivo-1", "2018-11-13-ak-invivo-1", "2018-11-13-al-invivo-1", "2018-11-14-aa-invivo-1",
    "2018-11-14-ad-invivo-1", "2018-11-14-ah-invivo-1", "2018-11-14-ai-invivo-1",
    "2018-11-14-ak-invivo-1", "2018-11-14-al-invivo-1", "2018-11-14-am-invivo-1", "2018-11-14-an-invivo-1",
    "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-13-aa-invivo-1", "2018-11-13-ac-invivo-1", "2018-11-13-ad-invivo-1", "2018-11-13-ah-invivo-1",
    "2018-11-13-ai-invivo-1", "2018-11-13-aj-invivo-1", "2018-11-13-ak-invivo-1", "2018-11-13-al-invivo-1"]

# parameters for binning, smoothing and plotting
cut_window = 20
#cut_window_csi = 20 #ms
#cut_window_plot = 50 #ms
chirp_duration = 14 #ms
neuronal_delay = 5 #ms
chirp_start = int((-chirp_duration / 2 + neuronal_delay + cut_window * 2) * sampling_rate) #index
chirp_end = int((chirp_duration / 2 + neuronal_delay + cut_window * 2) * sampling_rate) #index
number_bins = 12
window = 1 #ms
time_axis = np.arange(-cut_window*2, cut_window*2, 1/sampling_rate) #steps
spike_bins = np.arange(-cut_window*2, cut_window*2) #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 = map_keys(spikes)

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

# make dictionaries for spiketimes
df_phase_time = {}
df_phase_binary = {}
#embed()
#exit()

for dataset in data:
    spikes = read_chirp_spikes(os.path.join(data_dir, dataset))
    df_map = map_keys(spikes)
    print(dataset)
    # 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]:
            chirp_size = int(rep[-1].strip('Hz'))
            #print(chirp_size)
            if chirp_size == 150:
                continue
            for phase in spikes[rep]:
                for idx in np.arange(number_bins):
                    # check the phase
                    if phase[1] > phase_vec[idx] and phase[1] < phase_vec[idx+1]:

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


    # make dictionaries for csi and beat
    csi_trains = {}
    csi_rates = {}
    beat = {}
    # for plotting and calculating iterate over delta f and phases
    for df in df_phase_time.keys():
        csi_trains[df] = []
        csi_rates[df] = []
        beat[df] = []
        beat_duration = int(abs(1/df*1000)*sampling_rate) #steps
        beat_window = 0
        # beat window is at most 20 ms long, multiples of beat_duration
        while beat_window+beat_duration <= cut_window*sampling_rate:
            beat_window = beat_window+beat_duration
        for phase in df_phase_time[df].keys():

            # csi calculation
            # trains for synchrony and rate
            trials_binary = df_phase_binary[df][phase]

            train_chirp = []
            train_beat = []
            #csi_spikerate = []
            for i, trial in enumerate(trials_binary):
                smoothed_trial = smooth(trial, window, 1/sampling_rate)
                train_chirp.append(smoothed_trial[chirp_start:chirp_end])
                train_beat.append(smoothed_trial[chirp_start-beat_window:chirp_start])
                #std_chirp = np.std(smoothed_trial[chirp_start:chirp_end])
                #std_beat = np.std(smoothed_trial[chirp_start-beat_window:chirp_start])
                #csi = (std_chirp - std_beat)/(std_chirp + std_beat)
                #csi_spikerate.append(csi)

            std_chirp = np.std(np.mean(train_chirp, axis=0))
            std_beat = np.std(np.mean(train_beat, axis=0))
            beat[df].append(std_beat)
            csi_spikerate = (std_chirp - std_beat) / (std_chirp + std_beat)

            rcs = []
            rbs = []
            for i, train in enumerate(train_chirp):
                for j, train2 in enumerate(train_chirp):
                    if i >= j:
                        continue
                    else:
                        rc, _ = ss.pearsonr(train, train2)
                        rb, _ = ss.pearsonr(train_beat[i], train_beat[j])
                        rcs.append(rc)
                        rbs.append(rb)

            r_train_chirp = np.mean(rcs)
            r_train_beat = np.mean(rbs)

            csi_train = (r_train_chirp - r_train_beat) / (r_train_chirp + r_train_beat)

            # add the csi to the dictionaries with the correct df and phase
            csi_trains[df].append(csi_train)
            csi_rates[df].append(np.mean(csi_spikerate))

            '''
            # 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()
            '''

    '''
    fig, ax = plt.subplots()
    for i, k in enumerate(sorted(csi_rates.keys())):
        ax.scatter(np.ones(len(csi_rates[k]))*i, csi_rates[k], s=20)
        #ax.plot(i, np.mean(csi_rates[k]), 'o', markersize=15)
    ax.legend(sorted(csi_rates.keys()), loc='upper left', bbox_to_anchor=(1.04, 1))
    ax.plot(np.arange(-1, len(csi_rates.keys())+1), np.zeros(len(csi_rates.keys())+2), 'silver', linewidth=2, linestyle='--')
    #ax.set_xticklabels(sorted(csi_rates.keys()))
    fig.tight_layout()
    plt.show()
    
    fig, ax = plt.subplots()
    for i, k in enumerate(sorted(csi_trains.keys())):
        ax.plot(np.ones(len(csi_trains[k]))*i, csi_trains[k], 'o')
        #ax.plot(i, np.mean(csi_trains[k]), 'o', markersize=15)
    ax.legend(sorted(csi_trains.keys()), loc='upper left', bbox_to_anchor=(1.04, 1))
    ax.plot(np.arange(-1, len(csi_trains.keys())+1), np.zeros(len(csi_trains.keys())+2), 'silver', linewidth=2, linestyle='--')
    #ax.set_xticklabels(sorted(csi_trains.keys()))
    fig.tight_layout()
    plt.show()
    '''

    fig, ax = plt.subplots()
    for i, k in enumerate(sorted(beat.keys())):
        ax.plot(np.ones(len(beat[k]))*i, beat[k], 'o')
    ax.legend(sorted(beat.keys()), loc='upper left', bbox_to_anchor=(1.04, 1))
    #ax.set_xticklabels(sorted(csi_trains.keys()))
    fig.tight_layout()
    plt.show()