# this will compute TF spectrum from 1st awakening
# if it breaks, it's because that subject didn't ever fall asleep. Specify to skip them on line 205 (e.g., if (f != 41))

from scipy.stats import zscore
from scipy.signal import welch
from scipy.stats import linregress
import yasa
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as signal
import seaborn as sns
import pandas as pd
sns.set(font_scale=1.2)
import mne
import glob
from scipy.special import erf
import datetime
import csv
import os

PATH = '/Users/ajsimon/Dropbox (Personal)/Data/Overnight PSG/data/preprocessed/'

#create list of PSG files
FILES = glob.glob(PATH + '*.fif')
FILES.sort()

#create list of hypnogram files  
HYPNO = glob.glob(PATH + '*.csv')
HYPNO.sort()

# this is for indexing the last '/' in the path so that we can pull the filename easily for each participant
sep = '/'
def find(PATH, sep):
    return [i for i, ltr in enumerate(PATH) if ltr == sep]

pathsep = list(find(PATH,sep))
pathsep = pathsep[len(pathsep)-1]

sf = 100.

#format output header  -- spectral power for NREM stages
SleepHealth = np.zeros((len(FILES)+1,12), dtype=np.int)
SleepHealth = SleepHealth.astype('U20')

SleepHealth[0,0] = 'Subject'
SleepHealth[0,1] = 'Sleep Latency'
SleepHealth[0,2] = 'Total WASO min'
SleepHealth[0,3] = 'Total num WASO instances'
SleepHealth[0,4] = 'Avg len WASO'
SleepHealth[0,5] = 'Max len WASO'
SleepHealth[0,6] = 'Total num WASO instances 15 sec cutoff'
SleepHealth[0,7] = 'Avg len WASO 15 sec cutoff'
SleepHealth[0,8] = 'Max len WASO 15 sec cutoff'
SleepHealth[0,9] = 'Theta 1 min WASO slope'
SleepHealth[0,10] = 'Delta 1 min WASO slope'
SleepHealth[0,11] = 'Alpha 1 min WASO slope'

#find continuous regions of WASO (start and stop indexes) in sleep data
def contiguous_regions(condition):
    idx = []
    i = 0
    while i < len(condition):
        x1 = i + condition[i:].argmax()
        try:
            x2 = x1 + condition[x1:].argmin()
        except:
            x2 = x1 + 1
        if x1 == x2:
            if condition[x1] == True:
                x2 = len(condition)
            else:
                break
        idx.append( [x1,x2] )
        i = x2
    return idx


#create a variable that will contain a theta TS (all elecs averaged) for each sub
GlobalTheta = np.zeros((len(FILES),30),dtype=float)
GlobalAlpha = np.zeros((len(FILES),30),dtype=float)
GlobalDelta = np.zeros((len(FILES),30),dtype=float)

## initialize for loop
for f, file in enumerate(FILES):

    if (f != 41):
        PSG_fnum = FILES[f][np.s_[pathsep+1:pathsep+6]]
        Hyp_fnum= HYPNO[f][np.s_[pathsep+1:pathsep+6]]

            # make sure that the participant numbers between the PSG file and the hypnogram file for iteration 'f' are the same
        if PSG_fnum == Hyp_fnum:

            eeg = mne.io.read_raw_fif(FILES[f], preload=True)
            data = eeg.get_data() 
            data = data*1000000
            
            channels = eeg.ch_names
            
            sf = 100.
            times = np.arange(data.size) / sf
            
            hypnogram = []
            hypnogram = np.loadtxt(fname = HYPNO[f],dtype = 'str',delimiter = ',')  
            hypnogram = hypnogram.astype('U3')
            
            for r in range(len(hypnogram)):
                if hypnogram[r] == '- 1':
                    hypnogram[r] = '-1'
                elif hypnogram[r] == '2 1':
                    hypnogram[r] = '2'
                elif hypnogram[r] == '-':
                    hypnogram[r] = '-1'
                    
            
            where_sleep = np.isin(hypnogram, ['1','2','3','4'])  # True if sample is in any sleep stage, False otherwise
            sleep_ix = np.where(where_sleep)    #get a vector of sleep indeces so that I can split the night into before and after initially falling asleep
            sleep_ix = np.transpose(sleep_ix)
            begin_sleep = sleep_ix[0][0]
            end_sleep = sleep_ix[len(sleep_ix)-1][0]
            
            during_sleep_hypno = hypnogram[begin_sleep:end_sleep]
            
            during_sleep_data = data[:,begin_sleep:end_sleep]
            
            where_during_sleep_wakes = np.isin(during_sleep_hypno, ['0']) 
            
         #   data_pre_sleep = before_sleep_data[:, where_pre_sleep_wakes]
            data_during_sleep = during_sleep_data[:, where_during_sleep_wakes]
        
            wake_bouts = contiguous_regions(where_during_sleep_wakes)

            wake_len = np.zeros((len(wake_bouts),1), dtype=np.float)
            
            ix_awa = 0
            for w in range(len(wake_bouts)):
                wake_len[w] = (wake_bouts[w][1] - wake_bouts[w][0])   #WASO lengths in min
                
                if wake_len[w] > 1500:   #if greater than 15 sec, count as awakening
                    ix_awa = ix_awa + 1
                
            #for removing wakes less than 3 sec long
            wake_len_15sec = np.nonzero(wake_len > 1500)   #index WASOs greater than 3 sec -- Sleep Fragmentation, Mezick 2013
            wake_len_15sec = wake_len_15sec[0][:]
            wake_len_15sec = wake_len[wake_len_15sec]
            
            
            SleepHealth[f+1,0] = Hyp_fnum
            SleepHealth[f+1,1] = begin_sleep/(sf*60)
            SleepHealth[f+1,2] = len(data_during_sleep[1])/(60*sf)
            SleepHealth[f+1,3] = len(wake_bouts)
            SleepHealth[f+1,4] = np.mean(wake_len)
            SleepHealth[f+1,5] = np.max(wake_len)
            
            SleepHealth[f+1,6] = ix_awa
            SleepHealth[f+1,7] = np.mean(wake_len_15sec)
            SleepHealth[f+1,8] = np.max(wake_len_15sec)
            
            #for computing TF metrics on only first min on WASOs
            wake_len_1min = np.nonzero(wake_len > 6000)   
            wake_len_1min = wake_len_1min[0][:]
            
            wake_bouts = np.asarray(wake_bouts)
            
            wake_starts = wake_bouts[wake_len_1min,0]
            
            #create vector bools that is the length of data, default falses anything where they woke up for >1min will be 'True' for 1st min of WASO
            where_1min_WASOs = np.zeros((1,np.shape(during_sleep_data)[1]),dtype=bool)
            
            ix = 0
            for wk in range(len(where_1min_WASOs[0])):
                if wk == wake_starts[ix]:
                    if ix == 0:
                        where_1min_WASOs[0][wk:wk+6000] = True
                        ix = ix+1
                    
            where_1min_WASOs = np.squeeze(where_1min_WASOs)
            data_during_sleep = during_sleep_data[:, where_1min_WASOs]
            
            frame = 2
            itwin = int(frame*sf)

            DSlen = np.ceil(len(data_during_sleep[1])/itwin)

            thetaTS = np.zeros((6,int(DSlen)),dtype=float)
            alphaTS = np.zeros((6,int(DSlen)),dtype=float)
            deltaTS = np.zeros((6,int(DSlen)),dtype=float)

            for elec in range(data_during_sleep.shape[0]):
                ix = 0
                elec_thetaTS = np.zeros(int(DSlen),dtype=float)
                elec_alphaTS = np.zeros(int(DSlen),dtype=float)
                elec_deltaTS = np.zeros(int(DSlen),dtype=float)
                
                for wind in range(0,data_during_sleep.shape[1],itwin):

                    tempdata = data_during_sleep[elec,wind:wind+itwin]

                    win = int(2 * sf)  # Window size is set to 2 seconds
                    freqs, psd = welch(tempdata, sf, nperseg=win, average='median')  # Works with single or multi-channel data, freq bins are 0.25 Hz

                    theta = np.average(psd[8:15])
                    delta = np.average(psd[3:9])
                    alpha = np.average(psd[16:25])

                    thetaTS[elec,ix] = theta    #this one will store an electrode x time matrix that all be saved independently for each subject. This way, we have either a subject x electrode x time matrix, or keep the elec x time matrix but make 6 rows for each subject
                    alphaTS[elec,ix] = alpha
                    deltaTS[elec,ix] = delta
                    
                    elec_thetaTS[ix] = theta   #this one is for computing slope
                    elec_deltaTS[ix] = delta   #this one is for computing slope
                    elec_alphaTS[ix] = alpha   #this one is for computing slope

                    ix = ix+1
                    
            if thetaTS[1].shape[0] > 0:
                
                GlobalTheta[f,:] = thetaTS.mean(axis=0)
                GlobalAlpha[f,:] = alphaTS.mean(axis=0)
                GlobalDelta[f,:] = deltaTS.mean(axis=0)

                ## compute the slope of theta change across the first min of the first WASO that exceeds 1 min
                xx = np.zeros(GlobalTheta[f,:].shape[0],dtype = float)
                for x in range(GlobalTheta[f,:].shape[0]):
                    xx[x] = int(x)

                slope, intercept, r_value, p_value, std_err = linregress([xx],[GlobalTheta[f,:]])

                SleepHealth[f+1,9] = slope
                
                
                 ## compute the slope of delta change across the first min of the first WASO that exceeds 1 min
                xx = np.zeros(GlobalDelta[f,:].shape[0],dtype = float)
                for x in range(GlobalDelta[f,:].shape[0]):
                    xx[x] = int(x)

                slope, intercept, r_value, p_value, std_err = linregress([xx],[GlobalDelta[f,:]])

                SleepHealth[f+1,10] = slope
                
                 ## compute the slope of alpha change across the first min of the first WASO that exceeds 1 min
                xx = np.zeros(GlobalAlpha[f,:].shape[0],dtype = float)
                for x in range(GlobalAlpha[f,:].shape[0]):
                    xx[x] = int(x)

                slope, intercept, r_value, p_value, std_err = linregress([xx],[GlobalAlpha[f,:]])

                SleepHealth[f+1,11] = slope
                
                
 
 
WASO_outfile = open('/Users/ajsimon/Dropbox (Personal)/Data/Overnight PSG/data/' + 'WASO stats 15 sec threshold.csv','w')
with WASO_outfile:
    writer = csv.writer(WASO_outfile,delimiter=',')
    writer.writerows(SleepHealth)
    
    
theta_outfile = open('/Users/ajsimon/Dropbox (Personal)/Data/Overnight PSG/data/' + 'first min WASO theta time series.csv','w')
with theta_outfile:
    writer = csv.writer(theta_outfile,delimiter=',')
    writer.writerows(GlobalTheta)
    
alpha_outfile = open('/Users/ajsimon/Dropbox (Personal)/Data/Overnight PSG/data/' + 'first min WASO alpha time series.csv','w')
with alpha_outfile:
    writer = csv.writer(alpha_outfile,delimiter=',')
    writer.writerows(GlobalAlpha)
    
delta_outfile = open('/Users/ajsimon/Dropbox (Personal)/Data/Overnight PSG/data/' + 'first min WASO delta time series.csv','w')
with delta_outfile:
    writer = csv.writer(delta_outfile,delimiter=',')
    writer.writerows(GlobalDelta)