#%% initial loading of tools 
import os
os.chdir('C:\EAfiles\EAdetection')
#import core.ea_management as eam
import core.helpers as hf

import os
os.chdir('C:\EAfiles\MyCode')

import pickle


import matplotlib.cm as cmx # colormap module
import numpy as np # module for low-level scientific computing
import scipy as sp # library for working with NumPy arrays
from scipy import stats 
import scipy.signal # signal processing module
import pandas as pd # data manipulation tools. built on NumPy library
import matplotlib.pyplot as plt # makes matplotlib work like MATLAB. ’pyplot’ functions.
import seaborn as sns # based on matplotlib. high-level interface for visualization.

'''
abbreviations:
recID-recording ID: 'animal_session_electrode', e.g. 'PK22_pre1-3_HCc'	
sr-sampling rate
auc-area under the curve
psd-power spectral density
spec-spectrogram
dict-dictionary
ddict- data dictionary
pre- pre recording (the hour before stimulation)
stim- recording with 0.1 Hz or 0.05 Hz blue light stimulation
post- post recording (the hour after stimulation)
rec-recording
chan- channel of the recording electrode (HCi1- ipsilateral to saline/kainate injection, HCc-contralateral to saline/kainate injection)
'''

#%% define functions

# prepare variables:

sr = 10000. 

#prepare functions:

# Plot the power spectrum:
def plot_psd(f,psd,chan,chan_x, session):
    plt.subplot(2,1,chan_x+1)
    plt.subplots_adjust(hspace=0.4)
    if session == 'pre':
      plt.semilogy(f,psd,'darkgrey')
    elif session == 'stim':
      plt.semilogy(f,psd,'dodgerblue')
      #plt.semilogy(f,psd,'dimgrey') #for reference
    elif session == 'post':
      plt.semilogy(f,psd,'k')
      
    sns.despine()
    plt.xlim(0,100)
    plt.ylim(10**-6,10**-1)
    plt.yticks(size=15)
    plt.xticks(size=15)
    plt.ylabel('power [$mV^{2}/Hz$]',size=15)
    plt.xlabel('f [Hz]',size=15)
    plt.title('PSD of '+chan, size=15)


def plot_psd_average(f,psd, sem, session):
    if session == 'pre':
      plt.semilogy(f,psd,'darkgrey', linewidth=2)
    elif session == 'post':
      plt.semilogy(f,psd,'k', linewidth=2)
    elif session == 'stim':
     plt.semilogy(f,psd,'dodgerblue', linewidth=2)
     #plt.semilogy(f,psd,'dimgrey', linewidth=2) #for reference

    sns.despine()
    plt.xlim(0,120)
    plt.ylim(10**-6,10**-1)
    plt.yticks(size=20)
    plt.xticks(size=20)
    plt.ylabel('power [$mV^{2}/Hz$]',size=15)
    plt.xlabel('f [Hz]',size=15)
    
#extract area under the curve (auc) of the psd plot
def psd_auc(psd, freq_start, freq_end):
  start=freq_start*10 #the f-array is with 0.1 steps so that is why *10 is needed
  stop=freq_end*10+1
  freq_psd=psd[start:stop]
  freq_auc=np.sum(np.abs(freq_psd))
  return freq_auc

def psd_auc_dict(stim_ID, dictionary, psd):
  
  all_freq_start=0
  all_freq_end=120

  gamma_start=30
  gamma_end=120

  beta_start=12
  beta_end=30
  
  theta_start=4
  theta_end=12
  
  delta_start=1
  delta_end=4
  
  dictionary[stim_ID]['psd_auc']=psd_auc(psd, all_freq_start, all_freq_end)
  dictionary[stim_ID]['gamma_auc']=psd_auc(psd, gamma_start, gamma_end)
  dictionary[stim_ID]['beta_auc']=psd_auc(psd, beta_start, beta_end)
  dictionary[stim_ID]['theta_auc']=psd_auc(psd, theta_start, theta_end)
  dictionary[stim_ID]['delta_auc']=psd_auc(psd, delta_start, delta_end)
  
#plotting spectrograms:
def plot_spec_80szoom(datatrace,recstart_s,recstop_s,vmin=-100,vmax=-20):    
    #samp_spec = range(0,60*60*int(sr))
    samp_spec = range(0,80*int(sr))
    f, t_spec, x_spec = sp.signal.spectrogram(datatrace[samp_spec], fs=sr, \
                    window='hanning', nperseg=1*int(sr/2), nfft=10*int(sr), noverlap=int(0.25*sr), mode='psd') 
    
    #for a better quality zoom (shorter time periods), reduce the nperseg and noverlap, eg. noverlap=int(sr*0.05), nperseg=1*int(sr) 
    
    '''
    f- Array of sample frequencies.
    t_spec-Array of segment times.
    x_spec- Spectrogram of x. By default, the last axis of x_spec corresponds to the segment times.

    '''
    fmax = 120
        
    x_mesh, y_mesh = np.meshgrid(t_spec, f[f<fmax])
    t_index = np.arange(0,len(datatrace)/sr,0.5/sr) # time array  
    
    plt.pcolormesh(x_mesh+t_index[samp_spec[0]], y_mesh, 20*np.log10(x_spec[f<fmax]), cmap=cmx.jet, vmin=vmin, vmax=vmax)
    plt.ylabel('f [Hz]', size=12)
    plt.yticks(size=12)
    plt.xticks(size=12)
    #plt.colorbar()
    #plt.colorbar(label='dB')

'''now follow exacmples of these functions were used to analyze/present data in Kleis et al., 2021'''
#%% load dictionaries, where recording IDs are matched with corresponding smr file names (with .smr in the end)
filename = 'C:\EAfiles\MyDicts\IDmatch_dict'
infile = open(filename, 'rb') 
IDmatch_dict = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\preIDs_smr'
infile = open(filename, 'rb') 
preIDs_smr = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\stimIDs_smr'
infile = open(filename, 'rb') 
stimIDs_smr = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\postIDs_smr'
infile = open(filename, 'rb') 
postIDs_smr= pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\preIDs_smr_contra'
infile = open(filename, 'rb') 
preIDs_smr_contra = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\stimIDs_smr_contra'
infile = open(filename, 'rb') 
stimIDs_smr_contra = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\postIDs_smr_contra'
infile = open(filename, 'rb') 
postIDs_smr_contra= pickle.load(infile)
infile.close() 


filename = 'C:\EAfiles\MyDicts\preIDs_smr_ipsi'
infile = open(filename, 'rb') 
preIDs_smr_ipsi = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\stimIDs_smr_ipsi'
infile = open(filename, 'rb') 
stimIDs_smr_ipsi = pickle.load(infile)
infile.close() 

filename = 'C:\EAfiles\MyDicts\postIDs_smr_ipsi'
infile = open(filename, 'rb') 
postIDs_smr_ipsi= pickle.load(infile)
infile.close() 

#%% creating lists from textfiles, which contain recIDs that will be used for each type of recording
      
pathtobadrecIDs = 'C:\EAfiles\ID_Dateien\\badrecIDs.txt' 
badrecIDs = []
with open (pathtobadrecIDs,'r') as f:
    badIDs = f.read().splitlines()
    badrecIDs = badrecIDs + badIDs
    
pathto_saline_ref = 'C:\EAfiles\ID_Dateien\\saline_PACK-CA1_ref.txt' #used in example 1
saline_PACK_CA1_ref = []
with open (pathto_saline_ref,'r') as f:
    IDs = f.read().splitlines()
    saline_PACK_CA1_ref = saline_PACK_CA1_ref + IDs  
    
pathto_saline_01Hz = 'C:\EAfiles\ID_Dateien\\saline_PACK-CA1_01Hz.txt' #used in example 2
saline_PACK_CA1_01Hz = []
with open (pathto_saline_01Hz,'r') as f:
    IDs = f.read().splitlines()
    saline_PACK_CA1_01Hz = saline_PACK_CA1_01Hz+ IDs
    
pathto_saline_bPAC_01Hz = 'C:\EAfiles\ID_Dateien\\saline_bPAC-CA1_01Hz.txt' #used in example 3
saline_bPAC_CA1_01Hz = []
with open (pathto_saline_bPAC_01Hz,'r') as f:
    IDs = f.read().splitlines()
    saline_bPAC_CA1_01Hz = saline_bPAC_CA1_01Hz+ IDs

    
#%%(Example 1) plot individual PSDs and create a matrix with all PSDs (for saline PACK reference recordings)
    
sr=10000
count=0

os.chdir('C:\EAfiles\ID_Dateien')

file = open('saline_PACK-CA1.txt', 'r')
rows = 0
for line in file:
  line = line.strip('\n')
  rows += 1
psd_saline_ref1_matrix=np.zeros((rows,50001), dtype=float)
psd_saline_ref2_matrix=np.zeros((rows,50001), dtype=float)
psd_saline_ref3_matrix=np.zeros((rows,50001), dtype=float)
 

psd_saline_ref1_fmatrix=np.zeros((rows,50001), dtype=float)
psd_saline_ref2_fmatrix=np.zeros((rows,50001), dtype=float)
psd_saline_ref3_fmatrix=np.zeros((rows,50001), dtype=float)

# Load data
folder = r'I:\Piret\EEG' #add directory, where your raw data is (.smr files)


for stim_ID in saline_PACK_CA1_ref:
  
  print 'working on '+stim_ID
  plt.figure(figsize=(5,8))
  
  pre_ID = IDmatch_dict[stim_ID]['pre']
  post_ID = IDmatch_dict[stim_ID]['post']
    
  if stim_ID [-1:] == '2':
    chanlist = ['HCi1_2', 'HCc_2']
  elif stim_ID [-1:] == '3':
    chanlist = ['HCi1_3', 'HCc_3']
  elif stim_ID [-1:] == '4':
    chanlist = ['HCi1_4', 'HCc_4']
  else:
    chanlist = ['HCi1', 'HCc' ]
        
    filename_pre = preIDs_smr[pre_ID]
    filename_stim = stimIDs_smr[stim_ID]
    filename_post = postIDs_smr[post_ID]
    
    ddict_pre = hf.extract_smrViaNeo(os.path.join(folder, filename_pre), chanlist=chanlist)
    ddict_stim = hf.extract_smrViaNeo(os.path.join(folder, filename_stim), chanlist=chanlist)
    ddict_post = hf.extract_smrViaNeo(os.path.join(folder, filename_post), chanlist=chanlist)
    
    excelpath = r'C:\EAfiles\Excel\startstop_PK_HCc.xlsx'
    df = pd.read_excel(excelpath) #reading an excel file with start and stop times (s) of a recording
    df.set_index('recID', inplace=True) #recID
    pre_start = df.start[pre_ID] #start time
    pre_stop = df.stop[pre_ID] #stop time
    stim_start = df.start[stim_ID] #start time
    stim_stop = df.stop[stim_ID] #stop time
    post_start = df.start[post_ID] #start time
    post_stop = df.stop[post_ID] #stop time
      
    for chan_x, chan in enumerate(chanlist):
      print chan
          
      #get datatrace
      datatrace_pre = ddict_pre[chan]['trace']
      datatrace_stim = ddict_stim[chan]['trace']
      datatrace_post = ddict_post[chan]['trace']
          
      #cut data trace
      reccut_pre=datatrace_pre[int(pre_start*sr):int(pre_stop*sr)]
      reccut_stim=datatrace_stim[int(stim_start*sr):int(stim_stop*sr)]
      reccut_post=datatrace_post[int(post_start*sr):int(post_stop*sr)]
          
      # Compute the power spectrum
      f_pre, psd_pre = sp.signal.welch(reccut_pre, sr, nperseg=10*sr)
      f_stim, psd_stim = sp.signal.welch(reccut_stim, sr, nperseg=10*sr)
      f_post, psd_post = sp.signal.welch(reccut_post, sr, nperseg=10*sr)
          
      #plot the power spectrum
      plot_psd(f_pre,psd_pre,chan,chan_x, 'pre')
      plot_psd(f_stim,psd_stim,chan,chan_x, 'stim')
      plot_psd(f_post,psd_post,chan,chan_x, 'post')
          
      #add the power spectrum into psd_matrix:
      if chan=='HCi1_2' or chan=='HCi1':
        psd_saline_ref1_matrix[count]=psd_pre
        psd_saline_ref2_matrix[count]=psd_stim
        psd_saline_ref3_matrix[count]=psd_post
          
        psd_saline_ref1_fmatrix[count]=f_pre
        psd_saline_ref2_fmatrix[count]=f_stim
        psd_saline_ref3_fmatrix[count]=f_post
        
      elif chan=='HCi1_3' or chan=='HCi1_4':
        psd_saline_ref1_matrix[count]=psd_pre
        psd_saline_ref2_matrix[count]=psd_stim
        psd_saline_ref3_matrix[count]=psd_post
          
        psd_saline_ref1_fmatrix[count]=f_pre
        psd_saline_ref2_fmatrix[count]=f_stim
        psd_saline_ref3_fmatrix[count]=f_post
      
    count += 1
    plt.savefig('C:\EAfiles\SPECTROGRAM\PSD_PACK_reference/'+stim_ID+'_psd.png')
    #plt.close('all')   

#%% save the matrices containing PSDs:         
np.save('psd_ref1_matrix',  psd_saline_ref1_matrix)
np.save('psd_ref2_matrix',  psd_saline_ref2_matrix)
np.save('psd_ref3_matrix',  psd_saline_ref3_matrix)
          
np.save('psd_ref1_fmatrix',  psd_saline_ref1_fmatrix)
np.save('psd_ref2_fmatrix',  psd_saline_ref2_fmatrix)
np.save('psd_ref3_fmatrix',  psd_saline_ref3_fmatrix)
#%% plot the average PSDs:

psd_saline_ref1_average=np.mean(psd_saline_ref1_matrix, axis=0)
psd_saline_ref2_average=np.mean(psd_saline_ref2_matrix, axis=0)
psd_saline_ref3_average=np.mean(psd_saline_ref3_matrix, axis=0)

psd_saline_ref1_std=np.std(psd_saline_ref1_matrix, axis=0)
psd_saline_ref2_std=np.std(psd_saline_ref2_matrix, axis=0)
psd_saline_ref3_std=np.std(psd_saline_ref3_matrix, axis=0)

psd_saline_ref1_sem=stats.sem(psd_saline_ref1_matrix, axis=0)
psd_saline_ref2_sem=stats.sem(psd_saline_ref2_matrix, axis=0)
psd_saline_ref3_sem=stats.sem(psd_saline_ref3_matrix, axis=0)

f_saline_ref1_average=np.mean(psd_saline_ref1_fmatrix, axis=0)
f_saline_ref2_average=np.mean(psd_saline_ref2_fmatrix, axis=0)
f_saline_ref3_average=np.mean(psd_saline_ref3_fmatrix, axis=0)

np.save('psd_saline-PACK_ref1_average',  psd_saline_ref1_average)
np.save('psd_saline-PACK_ref2_average',  psd_saline_ref2_average)
np.save('psd_saline-PACL_ref3_average',  psd_saline_ref3_average)

plt.figure(figsize=(6,5))
plot_psd_average(f_saline_ref1_average,psd_saline_ref1_average, psd_saline_ref1_sem, 'pre')
plot_psd_average(f_saline_ref2_average,psd_saline_ref2_average, psd_saline_ref2_sem, 'stim')
plot_psd_average(f_saline_ref3_average,psd_saline_ref3_average, psd_saline_ref3_sem, 'post')

#%% (Example 2) getting the power areas of delta, theta, beta, and gamma oscillations 
#by taking the AUC of PSD plots in respective ranges
    
sr=10000
count=0

os.chdir('C:\EAfiles\ID_Dateien')

file = open('saline_bPAC-CA1_01Hz.txt', 'r')
rows = 0

psd_bPAC_01Hz_pre={}
psd_bPAC_01Hz_stim={}
psd_bPAC_01Hz_post={}

# Load data
folder = r'I:\Piret\EEG'


for stim_ID in stimIDs_smr_contra.keys():
    
    if stim_ID in saline_bPAC_CA1_01Hz:
    #else:
      print 'working on '+stim_ID
      psd_bPAC_01Hz_pre[stim_ID]={}
      psd_bPAC_01Hz_stim[stim_ID]={}
      psd_bPAC_01Hz_post[stim_ID]={}
  
      plt.figure(figsize=(5,8))
        
      pre_ID = IDmatch_dict[stim_ID]['pre']
      post_ID = IDmatch_dict[stim_ID]['post']
    
      if stim_ID [-1:] == '2':
          chanlist = ['HCc_2']
      elif stim_ID [-1:] == '3':
          chanlist = ['HCc_3']
      elif stim_ID [-1:] == '4':
          chanlist = ['HCc_4']
      else:
          chanlist = ['HCc' ]
        
      filename_pre = preIDs_smr[pre_ID]
      filename_stim = stimIDs_smr[stim_ID]
      filename_post = postIDs_smr[post_ID]
    
      ddict_pre = hf.extract_smrViaNeo(os.path.join(folder, filename_pre), chanlist=chanlist)
      ddict_stim = hf.extract_smrViaNeo(os.path.join(folder, filename_stim), chanlist=chanlist)
      ddict_post = hf.extract_smrViaNeo(os.path.join(folder, filename_post), chanlist=chanlist)
    
      excelpath = r'C:\EAfiles\Excel\startstop_controls_HCc.xlsx'
      df = pd.read_excel(excelpath) #reading an excel file with start and stop times (s) of a recording
      df.set_index('recID', inplace=True) #recID
      pre_start = df.start[pre_ID] #start time
      pre_stop = df.stop[pre_ID] #stop time
      stim_start = df.start[stim_ID] #start time
      stim_stop = df.stop[stim_ID] #stop time
      post_start = df.start[post_ID] #start time
      post_stop = df.stop[post_ID] #stop time
      
      for chan_x, chan in enumerate(chanlist):
          print chan
          
          #get datatrace
          datatrace_pre = ddict_pre[chan]['trace']
          datatrace_stim = ddict_stim[chan]['trace']
          datatrace_post = ddict_post[chan]['trace']
          
          #cut data trace
          reccut_pre=datatrace_pre[int(pre_start*sr):int(pre_stop*sr)]
          reccut_stim=datatrace_stim[int(stim_start*sr):int(stim_stop*sr)]
          reccut_post=datatrace_post[int(post_start*sr):int(post_stop*sr)]
          
          # Compute the power spectrum
          f_pre, psd_pre = sp.signal.welch(reccut_pre, sr, nperseg=10*sr)
          f_stim, psd_stim = sp.signal.welch(reccut_stim, sr, nperseg=10*sr)
          f_post, psd_post = sp.signal.welch(reccut_post, sr, nperseg=10*sr)
          
          
          #add the psds into dictionaries:
          psd_bPAC_01Hz_pre[stim_ID]['psd']=psd_pre
          psd_bPAC_01Hz_stim[stim_ID]['psd']=psd_stim
          psd_bPAC_01Hz_post[stim_ID]['psd']=psd_post
          
          #calculate and save power areas for beta, low gamma, high gamma, all frequencies:
          psd_auc_dict(stim_ID, psd_bPAC_01Hz_pre, psd_pre)
          psd_auc_dict(stim_ID, psd_bPAC_01Hz_stim, psd_stim)
          psd_auc_dict(stim_ID, psd_bPAC_01Hz_post, psd_post)
  
      count += 1
      
#save the 0.1 Hz PSD dictionaries and later convert them to excel files to extract data:

filename = 'C:\EAfiles\MyDicts\psd_bPAC_01Hz_pre'
outfile = open(filename, 'wb') #write binary: mit Erlaubnis zu (über)schreiben
pickle.dump(psd_bPAC_01Hz_pre, outfile)
outfile.close()

filename = 'C:\EAfiles\MyDicts\psd_bPAC_01Hz_stim'
outfile = open(filename, 'wb') #write binary: mit Erlaubnis zu (über)schreiben
pickle.dump(psd_bPAC_01Hz_stim, outfile)
outfile.close()

filename = 'C:\EAfiles\MyDicts\psd_bPAC_01Hz_post'
outfile = open(filename, 'wb') #write binary: mit Erlaubnis zu (über)schreiben
pickle.dump(psd_bPAC_01Hz_post, outfile)
outfile.close()
 

#%% (Example 3) plot spectrograms for saline PACK mice in sessions with 0.1 Hz stimulation (used for the figures in the paper, aligned with snippets from Spike2 LFPs)
      
folder = r'I:\Piret\EEG'

for stim_ID in IDmatch_dict.keys():

    if stim_ID in saline_PACK_CA1_01Hz: 

        plt.figure(figsize=(5,8))
        pre_ID = IDmatch_dict[stim_ID]['pre']
        post_ID = IDmatch_dict[stim_ID]['post']
        
        if stim_ID [-1:] == '2':
          chanlist = ['HCc_2']
        elif stim_ID [-1:] == '3':
          chanlist = ['HCc_3']
        elif stim_ID [-1:] == '4':
          chanlist = ['HCc_4']
        else:
          chanlist = ['HCc']
            
        filename_pre = preIDs_smr[pre_ID]
        filename_stim = stimIDs_smr[stim_ID]
        filename_post = postIDs_smr[post_ID]
        
        ddict_pre = hf.extract_smrViaNeo(os.path.join(folder, filename_pre), chanlist=chanlist)
        ddict_stim = hf.extract_smrViaNeo(os.path.join(folder, filename_stim), chanlist=chanlist)
        ddict_post = hf.extract_smrViaNeo(os.path.join(folder, filename_post), chanlist=chanlist)
        
        pre_start=1100
        pre_stop=1180
        

        stim_start=1100
        stim_stop=1180
        
        post_start=1100
        post_stop=1180

        for chan_x, chan in enumerate(chanlist):
            print chan
            sr=ddict_pre[chan]['sr']
            
            #get datatrace:
            datatrace_pre = ddict_pre[chan]['trace']
            datatrace_stim = ddict_stim[chan]['trace']
            datatrace_post = ddict_post[chan]['trace']
            
            #cut data trace:
            reccut_pre=datatrace_pre[int(pre_start*sr):int(pre_stop*sr)]
            reccut_stim=datatrace_stim[int(stim_start*sr):int(stim_stop*sr)]
            reccut_post=datatrace_post[int(post_start*sr):int(post_stop*sr)]
            
            #plot spectrograms for 80seconf LFP snippets:
            plt.subplot(4,1,1)
            plot_spec_80szoom(reccut_pre,pre_start,pre_stop,vmin=-100,vmax=-20)
            plt.title(stim_ID)
            plt.subplot(4,1,2)
            plot_spec_80szoom(reccut_stim,stim_start,stim_stop,vmin=-100,vmax=-20) 
            plt.subplot(4,1,3)
            plot_spec_80szoom(reccut_post,post_start,post_stop,vmin=-100,vmax=-20) 
            
            #showing 0.1 Hz pulses:
            plt.subplot(4,1,4)
            rectime_y=np.zeros(3600)
            rectime_x=np.arange(0,3600)
            start=10
            stop=3600
            period=10
            stimtimes=np.arange(start, stop, period)
            for i in rectime_x:
                if i in stimtimes:
                    rectime_y[i]=1
            plt.scatter(rectime_x[stim_start:stim_stop], rectime_y[stim_start:stim_stop], marker='|', color='dodgerblue')
            plt.xlim(stim_start,stim_stop)
            plt.show()

            
            plt.savefig('C:\EAfiles\SPECTROGRAM\saline_bPAC_01Hz/'+stim_ID+'_fft.png')
            plt.close('all')