from IPython import get_ipython
get_ipython().magic(u'matplotlib inline')
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import numpy as np
from scipy import signal


# some program constants
s_band_Hz = [5.0, 50.0]   # where to look for the alpha peak
NFFT = 512  # pitck the length of the fft
fs_Hz = 250.0        # assumed sample rate for the EEG data
f_lim_Hz = [0, 60]   # frequency limits for plotting
t_lim_sec = []       # default empty, it'll show all data
channel_to_plot = 1  # counting the first channel as one

# define which data to load
pname = 'JST'
fname = 'Coba1.txt' 
t_lim_sec = [0, 70]

# load data into numpy array
data = np.loadtxt(pname + fname,
                  delimiter=',',
                  skiprows=5)

# parse the data
data_indices = data[:, 0]  # the first column is the packet index
eeg_data_uV = data[:, channel_to_plot] # ignore the first channel (column 0), so channel 1 is column 1

# check data indices
d_indices = data_indices[2:]-data_indices[1:-1]
n_jump = np.count_nonzero((d_indices != 1) & (d_indices != 255))
print("Discontinuities inthe packet counter: " + str(n_jump))

# filter the data to remove DC
hp_cutoff_Hz = 1.0
print ("highpass filtering at: " + str(hp_cutoff_Hz) + " Hz")
b, a = signal.butter(2, hp_cutoff_Hz/(fs_Hz / 2.0), 'highpass')  # define the filter
f_eeg_data_uV = signal.lfilter(b, a, eeg_data_uV, 0) # apply along the zeroeth dimension

# notch filter the data to remove 60 Hz and 120 Hz interference
notch_freq_Hz = np.array([60.0, 120.0])  # these are the center frequencies
for freq_Hz in np.nditer(notch_freq_Hz):  # loop over each center freq
    bp_stop_Hz = freq_Hz + 3.0*np.array([-1, 1])  # set the stop band
    print ("notch filtering: " + str(bp_stop_Hz[0]) + "-" + str(bp_stop_Hz[1]) + " Hz")
    b, a = signal.butter(3, bp_stop_Hz/(fs_Hz / 2.0), 'bandstop')  # create the filter
    f_eeg_data_uV = signal.lfilter(b, a, f_eeg_data_uV, 0)  # apply along the zeroeth dimension


bp_Hz = np.zeros(0)
if (0):
    #filter to focus on alpha waves
    bp_Hz = np.array([7.0, 12.0])
    print ("bandpass filtering to: " + str(bp_Hz[0]) + "-" + str(bp_Hz[1]) + " Hz")
    b, a = signal.butter(3, bp_Hz/(fs_Hz / 2.0),'bandpass')  # create the filter
    f_eeg_data_uV = signal.lfilter(b, a, f_eeg_data_uV, 0)  # apply along the zeroeth dimension


# %% make time-domain plot

fig = plt.figure(figsize=(10.0, 13.5))  # make new figure, set size in inches

ax1 = plt.subplot(211)
t_sec = np.array(range(0, f_eeg_data_uV.size)) / fs_Hz
plt.plot(t_sec, f_eeg_data_uV)
plt.ylim(-7500, 7500)
plt.ylabel('EEG (uV)')
plt.xlabel('Waktu (s)')
plt.title(fname + " (Channel " + str(channel_to_plot) + ")")
plt.grid()
if (len(t_lim_sec) == 0):
    plt.xlim(t_sec[0], t_sec[-1])
else:
    plt.xlim(t_lim_sec)
    

# add annotation for filtering
if (bp_Hz.size > 0):
    ax1.text(0.025, 0.95,
             "BP = [" + str(bp_Hz[0]) + " " + str(bp_Hz[1]) + "] Hz",
             transform=ax1.transAxes,
             verticalalignment='top',
             horizontalalignment='left',
             backgroundcolor='w')
    



# %% make spectrogram

fig = plt.figure(figsize=(10.0, 10.5))  # make new figure, set size in inches
ax = plt.subplot(212, sharex=ax1)
FFTstep = NFFT/4  # do a new FFT every FFTstep data points
overlap = NFFT - FFTstep  # half-second steps
#overlap = NFFT - int(0.25 * fs_Hz)
spec_PSDperHz, freqs, t = mlab.specgram(np.squeeze(f_eeg_data_uV),
                               NFFT=NFFT,
                               window=mlab.window_hanning,
                               Fs=fs_Hz,
                               noverlap=overlap
                               ) # returns PSD power per Hz
spec_PSDperBin = spec_PSDperHz * fs_Hz / float(NFFT)  #convert to "per bin"
#del spec_PSDperHz  # remove this variable so that I don't mistakenly use it


plt.pcolor(t, freqs, 10*np.log10(spec_PSDperBin))  # dB re: 1 uV
plt.clim(25-5+np.array([-40, 0]))

#plt.ylim([0, fs_Hz/2.0])  # show the full frequency content of the signal
plt.ylim(f_lim_Hz)
plt.xlabel('Time (sec)')
plt.ylabel('Frequency (Hz)')
plt.title(fname + " (Channel " + str(channel_to_plot) + ")")
if (len(t_lim_sec) == 0):
    plt.xlim(t_sec[0], t_sec[-1])
else:
    plt.xlim(t_lim_sec)


# add annotation for FFT Parameters
ax.text(0.025, 0.95,
        "NFFT = " + str(NFFT) + "\nfs = " + str(int(fs_Hz)) + " Hz",
        transform=ax.transAxes,
        verticalalignment='top',
        horizontalalignment='left',
        backgroundcolor='w')

plt.tight_layout()

#%% find spectra that are in our time span
foo_spec = spec_PSDperBin
if (t_lim_sec[2-1] != 0):
    foo_spec = foo_spec[:, ((t >= t_lim_sec[0]) & (t <= t_lim_sec[1]))]
    
    # add markers on the figure
    yl = ax.get_ylim()
    plt.plot(t_lim_sec[0]*np.array([1, 1]), yl, 'k--', linewidth=2);
    plt.plot(t_lim_sec[1]*np.array([1, 1]), yl, 'k--', linewidth=2);

# get the mean spectra and convert from PSD to uVrms
mean_spectra_PSDperBin = np.mean(foo_spec, 1);
mean_spectra_uVrmsPerSqrtBin = np.sqrt(mean_spectra_PSDperBin)

# plot the mean spectra
fig = plt.figure(figsize=(10.0, 10.5))  # make new figure, set size in inches
ax = plt.subplot(212, sharex=ax1)
plt.plot(freqs, mean_spectra_uVrmsPerSqrtBin, '.-')
plt.xlim(f_lim_Hz)
plt.ylim([0, 200])
plt.xlabel('Frequency (Hz)')
plt.ylabel('RMS Amplitude\n(uV per sqrt(Bin))')

# add generic annotation about the FFT processing
ax.text(1.0-0.025, 0.95,
        "NFFT = " + str(NFFT) + "\nfs = " + str(int(fs_Hz)) + " Hz",
        transform=ax.transAxes,
        verticalalignment='top',
        horizontalalignment='right',
        backgroundcolor='w')        

plt.tight_layout()

#%% PSD

foo_spec = spec_PSDperBin
ind = ((t > t_lim_sec[0]) & (t < t_lim_sec[1]))
  
#get the mean spectrum in that time
spectrum_PSDperHz = np.mean(spec_PSDperHz[:,ind],1)  #time is horizontal in the 2D array?

#plot
fig = plt.figure(figsize=(10.0, 10.5))  # make new figure, set size in inches
ax = plt.subplot(212, sharex=ax1)
plt.plot(freqs, 10*np.log10(spectrum_PSDperHz))  # dB re: 1 uV
plt.xlim([0, fs_Hz/2.0])  # show the full frequency content of the signal
plt.xlim(f_lim_Hz)
#plt.xlim([0, 70])
plt.ylim([0.0, 50.0])
#plt.plot(38.0*np.array([1, 1]),ax.get_ylim(),'k--',linewidth=2)
#plt.plot(40.0*np.array([1, 1]),ax.get_ylim(),'k--',linewidth=2)
#plt.plot(42.0*np.array([1, 1]),ax.get_ylim(),'k--',linewidth=2)
plt.xlabel('Frequency (Hz)')
plt.ylabel('PSD')
plt.title(fname + " (Channel " + str(channel_to_plot) + ")")
plt.grid()

ax.text(1-0.025, 0.95,
    "NFFT = " + str(NFFT) + "\nfs = " + str(int(fs_Hz)) + " Hz",
    transform=ax.transAxes,
    verticalalignment='top',
    horizontalalignment='right',
    backgroundcolor='w')

plt.tight_layout()

# %% Plot the amplitude of alpha vs time
#reduce size
full_spec_PSDperBin = spec_PSDperBin
full_t_spec = t


# compute alpha vs time
bool_inds = (freqs > s_band_Hz[0]) & (freqs < s_band_Hz[1])
alpha_max_uVperSqrtBin = np.sqrt(np.amax(full_spec_PSDperBin[bool_inds,:],0))
alpha_sum_uVrms = np.sqrt(np.sum(full_spec_PSDperBin[bool_inds, :],0))


# make figure window
fig = plt.figure(figsize=(10.0, 15.5))  # make new figure, set size in inches

# make spectrogram (again)


# make time-domain plot of alpha amplitude
ax = plt.subplot(312, sharex=ax1)
#plt.plot(full_t_spec, alpha_sum_uVrms, '.-',
#         full_t_spec, alpha_max_uVperSqrtBin, '.-')
plt.plot(full_t_spec, alpha_max_uVperSqrtBin, '.-')
plt.ylim([0, 150])
plt.xlim([t_sec[0], t_sec[-1]])
if (t_lim_sec[2-1] != 0):
    plt.xlim(t_lim_sec)
plt.xlabel('Waktu (s)')
plt.ylabel('Amplitud (uVrms)')
plt.axhline(y=70.0, color='r', linestyle='--')
plt.text(30.0, 70.0, 'Threshold', fontsize=15, va='center', ha='center', backgroundcolor='w')
#plt.plot(38.0*np.array([1, 1]),ax.get_lim(),'k--',linewidth=2)
# plt.legend(('Sum In-Band', 'Max In-Band'), loc=2, fontsize='medium')
plt.legend(['Max In-Band'], loc=2, fontsize='medium')


plt.tight_layout()


# In[32]:

spec_PSDperBin[1,1]