"""

load_file takes the file name and reads in the data. It returns an

array containing only REM sleep stage

"""

gc.collect()

data = np.load(filename)

a = data['stages']

d = data['DATA']

srate = data['srate']

j=0

it = np.nditer(a, flags=['f_index'])

while not it.finished:

if( it[0]==5 ):

eeg = pull_epoch(d, srate, it.index)

if(j==0):

samples = eeg

else:

samples = np.hstack((samples,eeg))

j=j+1

it.iternext()

elength = srate * 30.0

startpos = epoch_number * elength

endpos = startpos + elength

eegdata = data[0:,int(startpos):int(endpos)]

plt.close()

plt.rc('font', family='Arial')

fig, ax = plt.subplots(nrows=9, ncols=2, sharex='col', sharey='row')

fig.suptitle(title + " - REM stage", fontsize=20)

plt.subplots_adjust(wspace = .05,hspace = 0.4 )

ny_nfft=1024

i=0

plt.tick_params(axis='both', labelsize=8)

while i<9:

Pxx, freq, bins, im = ax[i,0].specgram(bsl[i],NFFT=ny_nfft,Fs=rate)

ax[i,0].set_yticks(np.arange(0, 50, 10))

ax[i,0].set_ylim([0, 40])

if(i==8):

ax[i,0].set_xlabel("Time, seconds", fontsize=10)

ax[i,0].set_ylabel("Freq, Hz", fontsize=8)

ax[i,0].set_title('Baseline, channel:'+str(i+1), fontsize=10)

for label in (ax[i,0].get_xticklabels() + ax[i,0].get_yticklabels()):

label.set_fontname('Arial')

label.set_fontsize(8)

i=i+1

i=0

while i<9:

Pxx, freq, bins, im = ax[i,1].specgram(rec[i],NFFT=ny_nfft,Fs=rate)

ax[i,0].set_yticks(np.arange(0, 50, 10))

ax[i,1].set_ylim([0, 40])

#ax[i,1].set_xlim([0, 10000]) #13000])

if(i==8):

ax[i,1].set_xlabel("Time, seconds", fontsize=10)

#ax[i,1].set_ylabel("Freq, Hz")

ax[i,1].set_title('Recovery, channel:'+str(i+1), fontsize=10)

for label in (ax[i,0].get_xticklabels() + ax[i,0].get_yticklabels()):

label.set_fontname('Arial')

label.set_fontsize(8)

i=i+1

plt.show()