from numpy.random import seed

seed(1017)

from tensorflow import set_random_seed

set_random_seed(1017)

import os

from glob import glob

from collections import OrderedDict

import mne

from mne.io import RawArray

from mne import read_evokeds, read_source_spaces, compute_covariance

from mne import channels, find_events, concatenate_raws

from mne import pick_types, viz, io, Epochs, create_info

from mne import pick_channels, concatenate_epochs

from mne.datasets import sample

from mne.simulation import simulate_sparse_stc, simulate_raw

from mne.channels import read_montage

from mne.time_frequency import tfr_morlet

import numpy as np

from numpy import genfromtxt

import pandas as pd

pd.options.display.precision = 4

pd.options.display.max_columns = None

import matplotlib.pyplot as plt

plt.rcParams["figure.figsize"] = (12,12)

import keras

from keras import regularizers

from keras.callbacks import TensorBoard

from keras.models import Sequential, Model

from keras.layers import Dense, Dropout, Activation, Input

from keras.layers import Flatten, Conv2D, MaxPooling2D, LSTM

from keras.layers import BatchNormalization, Conv3D, MaxPooling3D

from sklearn.utils import class_weight

from sklearn.model_selection import train_test_split

class Feats:

  def __init__(self, num_classes=2, class_weights=[1,1], input_shape=[16,], 

               new_times=1, model_type='1', 

               x_train=1, y_train=1, x_test=1, y_test=1, x_val=1, y_val=1):

    self.num_classes = num_classes

    self.class_weights = class_weights

    self.input_shape = input_shape

    self.new_times = new_times

    self.model_type = model_type

    self.x_train = x_train

    self.y_train = y_train

    self.x_test = x_test

    self.y_test = y_test

    self.x_val = x_val

    self.y_val = y_val

def LoadBVData(sub,session,data_dir,exp):

  #for isub,sub in enumerate(subs):       

  print('Loading data for subject number: ' + sub)

  fname = data_dir + exp + '/' + sub + '_' + exp + '_' + session + '.vhdr'

  raw,sfreq = loadBV(fname,plot_sensors=False,plot_raw=False,

          plot_raw_psd=False,stim_channel=True)

  return raw

def loadBV(filename, plot_sensors=True, plot_raw=True,

  plot_raw_psd=True, stim_channel=False, ):

  """Load in recorder data files."""

  #load .vhdr files from brain vision recorder

  raw = io.read_raw_brainvision(filename,

                          montage='standard_1020',

                          eog=('HEOG', 'VEOG'),

                          preload=True,stim_channel=stim_channel)

  #set sampling rate

  sfreq = raw.info['sfreq']

  print('Sampling Rate = ' + str(sfreq))

  #load channel locations

  print('Loading Channel Locations')

  if plot_sensors:

    raw.plot_sensors(show_names='True')

  ##Plot raw data

  if plot_raw:

    raw.plot(n_channels=16, block=True)

   #plot raw psd

  if plot_raw_psd:

    raw.plot_psd(fmin=.1, fmax=100 )

  return raw, sfreq

def LoadMuseData(subs, nsesh, data_dir, load_verbose=False, sfreq=256.):

  nsubs = len(subs)

  raw = []

  print('Loading Data')

  for isub,sub in enumerate(subs):

    print('Subject number ' + str(isub+1) + '/' + str(nsubs))

    for isesh in range(nsesh):

      print(' Session number ' + str(isesh+1) + '/' + str(nsesh))

      raw.append(muse_load_data(data_dir, sfreq=sfreq ,subject_nb=sub,

                    session_nb=isesh+1,verbose=load_verbose))

  raw = concatenate_raws(raw)

  return raw

#from eeg-notebooks load_data

def muse_load_data(data_dir, subject_nb=1, session_nb=1, sfreq=256.,

                   ch_ind=[0, 1, 2, 3], stim_ind=5, replace_ch_names=None,

                   verbose=1):

    """Load CSV files from the /data directory into a Raw object.

    Args:

        data_dir (str): directory inside /data that contains the

            CSV files to load, e.g., 'auditory/P300'

    Keyword Args:

        subject_nb (int or str): subject number. If 'all', load all

            subjects.

        session_nb (int or str): session number. If 'all', load all

            sessions.

        sfreq (float): EEG sampling frequency

        ch_ind (list): indices of the EEG channels to keep

        stim_ind (int): index of the stim channel

        replace_ch_names (dict or None): dictionary containing a mapping to

            rename channels. Useful when an external electrode was used.

    Returns:

        (mne.io.array.array.RawArray): loaded EEG

    """

    if subject_nb == 'all':

        subject_nb = '*'

    if session_nb == 'all':

        session_nb = '*'

    data_path = os.path.join(

            'eeg-notebooks_v0.1/data', data_dir,

            'subject{}/session{}/*.csv'.format(subject_nb, session_nb))

    fnames = glob(data_path)

    return load_muse_csv_as_raw(fnames,

                                sfreq=sfreq,

                                ch_ind=ch_ind,

                                stim_ind=stim_ind,

                                replace_ch_names=replace_ch_names,

                                verbose=verbose)

#from eeg-notebooks

def load_muse_csv_as_raw(filename, sfreq=256., ch_ind=[0, 1, 2, 3],

                         stim_ind=5, replace_ch_names=None, verbose=1):

    """Load CSV files into a Raw object.

    Args:

        filename (str or list): path or paths to CSV files to load

    Keyword Args:

        subject_nb (int or str): subject number. If 'all', load all

            subjects.

        session_nb (int or str): session number. If 'all', load all

            sessions.

        sfreq (float): EEG sampling frequency

        ch_ind (list): indices of the EEG channels to keep

        stim_ind (int): index of the stim channel

        replace_ch_names (dict or None): dictionary containing a mapping to

            rename channels. Useful when an external electrode was used.

    Returns:

        (mne.io.array.array.RawArray): loaded EEG

    """

    n_channel = len(ch_ind)

    raw = []

    for fname in filename:

        # read the file

        data = pd.read_csv(fname, index_col=0)

        # name of each channels

        ch_names = list(data.columns)[0:n_channel] + ['Stim']

        if replace_ch_names is not None:

            ch_names = [c if c not in replace_ch_names.keys()

                        else replace_ch_names[c] for c in ch_names]

        # type of each channels

        ch_types = ['eeg'] * n_channel + ['stim']

        montage = read_montage('standard_1005')

        # get data and exclude Aux channel

        data = data.values[:, ch_ind + [stim_ind]].T

        # convert in Volts (from uVolts)

        data[:-1] *= 1e-6

        # create MNE object

        info = create_info(ch_names=ch_names, ch_types=ch_types,

                           sfreq=sfreq, montage=montage, verbose=verbose)

        raw.append(RawArray(data=data, info=info, verbose=verbose))

    # concatenate all raw objects

    if len(raw) > 0:

      raws = concatenate_raws(raw, verbose=verbose)

    else:

      print('No files for subject with filename ' + str(filename))

      raws = raw

    return raws

def SimulateRaw(amp1 = 50, amp2 = 100, freq = 1., batch=1):

  """Create simulated raw data and events of two kinds

  Keyword Args:

      amp1 (float): amplitude of first condition effect

      amp2 (float): ampltiude of second condition effect, 

          null hypothesis amp1=amp2

      freq (float): Frequency of simulated signal 1. for ERP 10. for alpha

      batch (int): number of groups of 255 trials in each condition

  Returns: 

      raw: simulated EEG MNE raw object with two event types

      event_id: dict of the two events for input to PreProcess()

  """

  data_path = sample.data_path()

  raw_fname = data_path + '/MEG/sample/sample_audvis_raw.fif'

  trans_fname = data_path + '/MEG/sample/sample_audvis_raw-trans.fif'

  src_fname = data_path + '/subjects/sample/bem/sample-oct-6-src.fif'

  bem_fname = (data_path + 

        '/subjects/sample/bem/sample-5120-5120-5120-bem-sol.fif')

  raw_single = mne.io.read_raw_fif(raw_fname,preload=True)

  raw_single.set_eeg_reference(projection=True)

  raw_single = raw_single.crop(0., 255.)

  raw_single = raw_single.copy().pick_types(meg=False, eeg=True, eog=True, stim=True)

  #concatenate 4 raws together to make 1000 trials

  raw = []

  for i in range(batch):

    raw.append(raw_single)

  raw = concatenate_raws(raw)

  epoch_duration = 1.

  def data_fun(amp, freq):

    """Create function to create fake signal"""

    def data_fun_inner(times):

      """Create fake signal with no noise"""

      n_samp = len(times)

      window = np.zeros(n_samp)

      start, stop = [int(ii * float(n_samp) / 2)

        for ii in (0, 1)]

      window[start:stop] = np.hamming(stop - start)

      data = amp * 1e-9 * np.sin(2. * np.pi * freq * times)

      data *= window

      return data

    return data_fun_inner

  times = raw.times[:int(raw.info['sfreq'] * epoch_duration)]

  src = read_source_spaces(src_fname)

  stc_zero = simulate_sparse_stc(src, n_dipoles=1, times=times,

              data_fun=data_fun(amp1,freq), random_state=0)

  stc_one = simulate_sparse_stc(src, n_dipoles=1, times=times,

              data_fun=data_fun(amp2,freq), random_state=0)

  raw_sim_zero = simulate_raw(raw, stc_zero, trans_fname, src, bem_fname, 

            cov='simple', blink=True, n_jobs=1, verbose=True)

  raw_sim_one = simulate_raw(raw, stc_one, trans_fname, src, bem_fname, 

            cov='simple', blink=True, n_jobs=1, verbose=True)

  stim_pick = raw_sim_one.info['ch_names'].index('STI 014')

  raw_sim_one._data[stim_pick][np.where(raw_sim_one._data[stim_pick]==1)] = 2

  raw = concatenate_raws([raw_sim_zero, raw_sim_one])

  event_id = {'CondZero': 1,'CondOne': 2}

  return raw, event_id

def mastoidReref(raw):

  ref_idx = pick_channels(raw.info['ch_names'],['M2'])

  eeg_idx = pick_types(raw.info,eeg=True)

  raw._data[eeg_idx,:] =  raw._data[eeg_idx,:]  -  raw._data[ref_idx,:] * .5 ;

  return raw

def GrattonEmcpRaw(raw):

  raw_eeg = raw.copy().pick_types(eeg=True)[:][0]

  raw_eog = raw.copy().pick_types(eog=True)[:][0]

  b = np.linalg.solve(np.dot(raw_eog,raw_eog.T), np.dot(raw_eog,raw_eeg.T))

  eeg_corrected = (raw_eeg.T - np.dot(raw_eog.T,b)).T

  raw_new = raw.copy()

  raw_new._data[pick_types(raw.info,eeg=True),:] = eeg_corrected

  return raw_new

def GrattonEmcpEpochs(epochs):

  '''

  # Correct EEG data for EOG artifacts with regression

  # INPUT - MNE epochs object (with eeg and eog channels)

  # OUTPUT - MNE epochs object (with eeg corrected)

  # After: Gratton,Coles,Donchin, 1983

  # -compute the ERP in each condition

  # -subtract ERP from each trial

  # -subtract baseline (mean over all epoch)

  # -predict eye channel remainder from eeg remainder

  # -use coefficients to subtract eog from eeg

  '''

  event_names = ['A_error','B_error']

  i = 0

  for key, value in sorted(epochs.event_id.items(), key=lambda x: (x[1], x[0])):

    event_names[i] = key

    i += 1

  #select the correct channels and data

  eeg_chans = pick_types(epochs.info, eeg=True, eog=False)

  eog_chans = pick_types(epochs.info, eeg=False, eog=True)

  original_data = epochs._data

  #subtract the average over trials from each trial

  rem = {}

  for event in event_names:

    data = epochs[event]._data

    avg = np.mean(epochs[event]._data,axis=0)

    rem[event] = data-avg

  #concatenate trials together of different types

  ## then put them all back together in X (regression on all at once)

  allrem = np.concatenate([rem[event] for event in event_names])

  #separate eog and eeg

  X = allrem[:,eeg_chans,:]

  Y = allrem[:,eog_chans,:]

  #subtract mean over time from every trial/channel

  X = (X.T - np.mean(X,2).T).T

  Y = (Y.T - np.mean(Y,2).T).T

  #move electrodes first

  X = np.moveaxis(X,0,1)

  Y = np.moveaxis(Y,0,1)

  #make 2d and compute regression

  X = np.reshape(X,(X.shape[0],np.prod(X.shape[1:])))

  Y = np.reshape(Y,(Y.shape[0],np.prod(Y.shape[1:])))

  b = np.linalg.solve(np.dot(Y,Y.T), np.dot(Y,X.T))

  #get original data and electrodes first for matrix math

  raw_eeg = np.moveaxis(original_data[:,eeg_chans,:],0,1)

  raw_eog = np.moveaxis(original_data[:,eog_chans,:],0,1)

  #subtract weighted eye channels from eeg channels

  eeg_corrected = (raw_eeg.T - np.dot(raw_eog.T,b)).T

  #move back to match epochs

  eeg_corrected = np.moveaxis(eeg_corrected,0,1)

  #copy original epochs and replace with corrected data

  epochs_new = epochs.copy()

  epochs_new._data[:,eeg_chans,:] = eeg_corrected

  return epochs_new

def PreProcess(raw, event_id, plot_psd=False, filter_data=True,

               filter_range=(1,30), plot_events=False, epoch_time=(-.2,1),

               baseline=(-.2,0), rej_thresh_uV=200, rereference=False, 

               emcp_raw=False, emcp_epochs=False, epoch_decim=1, plot_electrodes=False,

               plot_erp=False):

  sfreq = raw.info['sfreq']

  #create new output freq for after epoch or wavelet decim

  nsfreq = sfreq/epoch_decim

  tmin=epoch_time[0]

  tmax=epoch_time[1]

  if filter_range[1] > nsfreq:

    filter_range[1] = nsfreq/2.5  #lower than 2 to avoid aliasing from decim??

  #pull event names in order of trigger number

  event_names = ['A_error','B_error']

  i = 0

  for key, value in sorted(event_id.items(), key=lambda x: (x[1], x[0])):

    event_names[i] = key

    i += 1

  #Filtering

  if rereference:

    print('Rerefering to average mastoid')

    raw = mastoidReref(raw)

  if filter_data:

    print('Filtering Data Between ' + str(filter_range[0]) + 

            ' and ' + str(filter_range[1]) + ' Hz.')

    raw.filter(filter_range[0],filter_range[1],

               method='iir', verbose='WARNING' )

  if plot_psd:

    raw.plot_psd(fmin=filter_range[0], fmax=nsfreq/2 )

  #Eye Correction

  if emcp_raw:

    print('Raw Eye Movement Correction')

    raw = GrattonEmcpRaw(raw)

  #Epoching

  events = find_events(raw,shortest_event=1)

  color = {1: 'red', 2: 'black'}

  #artifact rejection

  rej_thresh = rej_thresh_uV*1e-6

  #plot event timing

  if plot_events:

    viz.plot_events(events, sfreq, raw.first_samp, color=color,

                        event_id=event_id)

  #Construct events - Main function from MNE

  epochs = Epochs(raw, events=events, event_id=event_id,

                  tmin=tmin, tmax=tmax, baseline=baseline,

                  preload=True,reject={'eeg':rej_thresh},

                  verbose=False, decim=epoch_decim)

  print('Remaining Trials: ' + str(len(epochs)))

  #Gratton eye movement correction procedure on epochs

  if emcp_epochs:

    print('Epochs Eye Movement Correct')

    epochs = GrattonEmcpEpochs(epochs)

  ## plot ERP at each electrode

  evoked_dict = {event_names[0]:epochs[event_names[0]].average(),

                              event_names[1]:epochs[event_names[1]].average()}

  # butterfly plot

  if plot_electrodes:

    picks = pick_types(evoked_dict[event_names[0]].info, meg=False, eeg=True, eog=False)

    fig_zero = evoked_dict[event_names[0]].plot(spatial_colors=True,picks=picks)

    fig_zero = evoked_dict[event_names[1]].plot(spatial_colors=True,picks=picks)

  # plot ERP in each condition on same plot

  if plot_erp:

    #find the electrode most miximal on the head (highest in z)

    picks = np.argmax([evoked_dict[event_names[0]].info['chs'][i]['loc'][2] 

              for i in range(len(evoked_dict[event_names[0]].info['chs']))])

    colors = {event_names[0]:"Red",event_names[1]:"Blue"}

    viz.plot_compare_evokeds(evoked_dict,colors=colors,

                            picks=picks,split_legend=True)

  return epochs

def FeatureEngineer(epochs, model_type='NN',

                    frequency_domain=False,

                    normalization=False, electrode_median=False,

                    wavelet_decim=1, flims=(3,30), include_phase=False,

                    f_bins=20, wave_cycles=3, 

                    wavelet_electrodes = [11,12,13,14,15],

                    spect_baseline=[-1,-.5],

                    test_split = 0.2, val_split = 0.2,

                    random_seed=1017, watermark = False):

  """

  Takes epochs object as 

  input and settings, 

  outputs  feats(training, test and val data option to use frequency or time domain)

  TODO: take tfr? or autoencoder encoded object?

  FeatureEngineer(epochs, model_type='NN',

                    frequency_domain=False,

                    normalization=False, electrode_median=False,

                    wavelet_decim=1, flims=(3,30), include_phase=False,

                    f_bins=20, wave_cycles=3, 

                    wavelet_electrodes = [11,12,13,14,15],

                    spect_baseline=[-1,-.5],

                    test_split = 0.2, val_split = 0.2,

                    random_seed=1017, watermark = False):

  """

  np.random.seed(random_seed)

  #pull event names in order of trigger number

  epochs.event_id = {'cond0':1, 'cond1':2}

  event_names = ['cond0','cond1']

  i = 0

  for key, value in sorted(epochs.event_id.items(),

                           key=lambda item: (item[1],item[0])):

    event_names[i] = key

    i += 1

  #Create feats object for output

  feats = Feats()

  feats.num_classes = len(epochs.event_id)

  feats.model_type = model_type

  if frequency_domain:

    print('Constructing Frequency Domain Features')

    #list of frequencies to output

    f_low = flims[0]

    f_high = flims[1]

    frequencies =  np.linspace(f_low, f_high, f_bins, endpoint=True)

    #option to select all electrodes for fft

    if wavelet_electrodes == 'all':

      wavelet_electrodes = pick_types(epochs.info,eeg=True,eog=False)

    #type of output from wavelet analysis

    if include_phase:

      tfr_output_type = 'complex'

    else:

      tfr_output_type = 'power'

    tfr_dict = {}

    for event in event_names:

      print('Computing Morlet Wavelets on ' + event)

      tfr_temp = tfr_morlet(epochs[event], freqs=frequencies,

                            n_cycles=wave_cycles, return_itc=False,

                            picks=wavelet_electrodes, average=False,

                            decim=wavelet_decim, output=tfr_output_type)

      # Apply spectral baseline and find stim onset time

      tfr_temp = tfr_temp.apply_baseline(spect_baseline,mode='mean')

      stim_onset = np.argmax(tfr_temp.times>0)

      # Reshape power output and save to tfr dict

      power_out_temp = np.moveaxis(tfr_temp.data[:,:,:,stim_onset:],1,3)

      power_out_temp = np.moveaxis(power_out_temp,1,2)

      print(event + ' trials: ' + str(len(power_out_temp)))

      tfr_dict[event] = power_out_temp

    #reshape times (sloppy but just use the last temp tfr)

    feats.new_times = tfr_temp.times[stim_onset:]

    for event in event_names:

      print(event + ' Time Points: ' + str(len(feats.new_times)))

      print(event + ' Frequencies: ' + str(len(tfr_temp.freqs)))

    #Construct X and Y

    for ievent,event in enumerate(event_names):

      if ievent == 0:

        X = tfr_dict[event]

        Y_class = np.zeros(len(tfr_dict[event]))

      else:

        X = np.append(X,tfr_dict[event],0)

        Y_class = np.append(Y_class,np.ones(len(tfr_dict[event]))*ievent,0)

    #concatenate real and imaginary data

    if include_phase:

      print('Concatenating the real and imaginary components')

      X = np.append(np.real(X),np.imag(X),2)

    #compute median over electrodes to decrease features

    if electrode_median:

      print('Computing Median over electrodes')

      X = np.expand_dims(np.median(X,axis=len(X.shape)-1),2)

    #reshape for various models

    if model_type == 'NN' or model_type == 'LSTM':

      X = np.reshape(X, (X.shape[0], X.shape[1], np.prod(X.shape[2:])))

    if model_type == 'CNN3D':

      X = np.expand_dims(X,4)

    if model_type == 'AUTO' or model_type == 'AUTODeep':

      print('Auto model reshape')

      X = np.reshape(X, (X.shape[0],np.prod(X.shape[1:])))

  if not frequency_domain:

    print('Constructing Time Domain Features')

    #if using muse aux port as eeg must label it as such

    eeg_chans = pick_types(epochs.info,eeg=True,eog=False)

    #put channels last, remove eye and stim

    X = np.moveaxis(epochs._data[:,eeg_chans,:],1,2);

    #take post baseline only

    stim_onset = np.argmax(epochs.times>0)

    feats.new_times = epochs.times[stim_onset:]

    X = X[:,stim_onset:,:]

    #convert markers to class

    #requires markers to be 1 and 2 in data file?

    #This probably is not robust to other marker numbers

    Y_class = epochs.events[:,2]-1  #subtract 1 to make 0 and 1

    #median over electrodes to reduce features

    if electrode_median:

      print('Computing Median over electrodes')

      X = np.expand_dims(np.median(X,axis=len(X.shape)-1),2)

    ## Model Reshapes:

    # reshape for CNN

    if model_type == 'CNN':

      print('Size X before reshape for CNN: ' + str(X.shape))

      X = np.expand_dims(X,3 )

      print('Size X before reshape for CNN: ' + str(X.shape))

    # reshape for CNN3D

    if model_type == 'CNN3D':

      print('Size X before reshape for CNN3D: ' + str(X.shape))

      X = np.expand_dims(np.expand_dims(X,3),4)

      print('Size X before reshape for CNN3D: ' + str(X.shape))

    #reshape for autoencoder

    if model_type == 'AUTO' or model_type == 'AUTODeep':

      print('Size X before reshape for Auto: ' + str(X.shape))

      X = np.reshape(X, (X.shape[0], np.prod(X.shape[1:])))

      print('Size X after reshape for Auto: ' + str(X.shape))

  #Normalize X - TODO: need to save mean and std for future test + val

  if normalization:

    print('Normalizing X')

    X = (X - np.mean(X)) / np.std(X)

  # convert class vectors to one hot Y and recast X

  Y = keras.utils.to_categorical(Y_class,feats.num_classes)

  X = X.astype('float32')

  # add watermark for testing models

  if watermark:

    X[Y[:,0]==0,0:2,] = 0

    X[Y[:,0]==1,0:2,] = 1

  # Compute model input shape

  feats.input_shape = X.shape[1:]

  # Split training test and validation data

  val_prop = val_split / (1-test_split)

  (feats.x_train,

    feats.x_test,

    feats.y_train,

    feats.y_test) = train_test_split(X, Y,

                                     test_size=test_split,

                                     random_state=random_seed)

  (feats.x_train,

   feats.x_val,

   feats.y_train,

   feats.y_val) = train_test_split(feats.x_train, feats.y_train,

                                   test_size=val_prop,

                                   random_state=random_seed)

  #compute class weights for uneven classes

  y_ints = [y.argmax() for y in feats.y_train]

  feats.class_weights = class_weight.compute_class_weight('balanced',

                                                 np.unique(y_ints),

                                                 y_ints)

  #Print some outputs

  print('Combined X Shape: ' + str(X.shape))

  print('Combined Y Shape: ' + str(Y_class.shape))

  print('Y Example (should be 1s & 0s): ' + str(Y_class[0:10]))

  print('X Range: ' + str(np.min(X)) + ':' + str(np.max(X)))

  print('Input Shape: ' + str(feats.input_shape))

  print('x_train shape:', feats.x_train.shape)

  print(feats.x_train.shape[0], 'train samples')

  print(feats.x_test.shape[0], 'test samples')

  print(feats.x_val.shape[0], 'validation samples')

  print('Class Weights: ' + str(feats.class_weights))

  return feats

def CreateModel(feats,units=[16,8,4,8,16], dropout=.25,

                batch_norm=True, filt_size=3, pool_size=2):

  print('Creating ' +  feats.model_type + ' Model')

  print('Input shape: ' + str(feats.input_shape))

  nunits = len(units)

  ##---LSTM - Many to two, sequence of time to classes

  #Units must be at least two

  if feats.model_type == 'LSTM':

    if nunits < 2:

      print('Warning: Need at least two layers for LSTM')

    model = Sequential()

    model.add(LSTM(input_shape=(None, feats.input_shape[1]),

                   units=units[0], return_sequences=True))

    if batch_norm:

      model.add(BatchNormalization())

    model.add(Activation('relu'))

    if dropout:

      model.add(Dropout(dropout))

    if len(units) > 2:

      for unit in units[1:-1]:

        model.add(LSTM(units=unit,return_sequences=True))

        if batch_norm:

          model.add(BatchNormalization())

        model.add(Activation('relu'))

        if dropout:

          model.add(Dropout(dropout))

    model.add(LSTM(units=units[-1],return_sequences=False))

    if batch_norm:

      model.add(BatchNormalization())

    model.add(Activation('relu'))

    if dropout:

      model.add(Dropout(dropout))

    model.add(Dense(units=feats.num_classes))

    model.add(Activation("softmax"))

  ##---DenseFeedforward Network

  #Makes a hidden layer for each item in units

  if feats.model_type == 'NN':

    model = Sequential()

    model.add(Flatten(input_shape=feats.input_shape))

    for unit in units:

      model.add(Dense(unit))

      if batch_norm:

        model.add(BatchNormalization())

      model.add(Activation('relu'))

      if dropout:

        model.add(Dropout(dropout))

    model.add(Dense(feats.num_classes, activation='softmax'))

  ##----Convolutional Network

  if feats.model_type == 'CNN':

    if nunits < 2:

      print('Warning: Need at least two layers for CNN')

    model = Sequential()

    model.add(Conv2D(units[0], filt_size,

              input_shape=feats.input_shape, padding='same'))

    model.add(Activation('relu'))

    model.add(MaxPooling2D(pool_size=pool_size, padding='same'))

    if nunits > 2:

      for unit in units[1:-1]:

        model.add(Conv2D(unit, filt_size, padding='same'))

        model.add(Activation('relu'))

        model.add(MaxPooling2D(pool_size=pool_size, padding='same'))

    model.add(Flatten())

    model.add(Dense(units[-1]))

    model.add(Activation('relu'))

    model.add(Dense(feats.num_classes))

    model.add(Activation('softmax'))

  ##----Convolutional Network

  if feats.model_type == 'CNN3D':

    if nunits < 2:

      print('Warning: Need at least two layers for CNN')

    model = Sequential()

    model.add(Conv3D(units[0], filt_size,

                     input_shape=feats.input_shape, padding='same'))

    model.add(Activation('relu'))

    model.add(MaxPooling3D(pool_size=pool_size, padding='same'))

    if nunits > 2:

      for unit in units[1:-1]:

        model.add(Conv3D(unit, filt_size, padding='same'))

        model.add(Activation('relu'))

        model.add(MaxPooling3D(pool_size=pool_size, padding='same'))

    model.add(Flatten())

    model.add(Dense(units[-1]))

    model.add(Activation('relu'))

    model.add(Dense(feats.num_classes))

    model.add(Activation('softmax'))

  ## Autoencoder

  #takes the first item in units for hidden layer size

  if feats.model_type == 'AUTO':

    encoding_dim = units[0]

    input_data = Input(shape=(feats.input_shape[0],))

    #,activity_regularizer=regularizers.l1(10e-5)

    encoded = Dense(encoding_dim, activation='relu')(input_data)

    decoded = Dense(feats.input_shape[0], activation='sigmoid')(encoded)

    model = Model(input_data, decoded)

    encoder = Model(input_data,encoded)

    encoded_input = Input(shape=(encoding_dim,))

    decoder_layer = model.layers[-1]

    decoder = Model(encoded_input, decoder_layer(encoded_input))

  #takes an odd number of layers > 1

  #e.g. units = [64,32,16,32,64]

  if feats.model_type == 'AUTODeep':

    if nunits % 2 == 0:

      print('Warning: Please enter odd number of layers into units')

    half = nunits/2

    midi = int(np.floor(half))

    input_data = Input(shape=(feats.input_shape[0],))

    encoded = Dense(units[0], activation='relu')(input_data)

    #encoder decreases

    if nunits >= 3:

        for unit in units[1:midi]:

          encoded = Dense(unit, activation='relu')(encoded)

    #latent space

    decoded = Dense(units[midi], activation='relu')(encoded)

    #decoder increses

    if nunits >= 3:

      for unit in units[midi+1:-1]:

        decoded = Dense(unit, activation='relu')(decoded)

    decoded = Dense(units[-1], activation='relu')(decoded)

    decoded = Dense(feats.input_shape[0], activation='sigmoid')(decoded)

    model = Model(input_data, decoded)

    encoder = Model(input_data,encoded)

    encoded_input = Input(shape=(units[midi],))

  if feats.model_type == 'AUTO' or feats.model_type == 'AUTODeep':

    opt = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999,

                                epsilon=None, decay=0.0, amsgrad=False)

    model.compile(optimizer=opt, loss='mean_squared_error')

  if ((feats.model_type == 'CNN') or

      (feats.model_type == 'CNN3D') or

      (feats.model_type == 'LSTM') or

      (feats.model_type == 'NN')):

    # initiate adam optimizer

    opt = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999,

                                epsilon=None, decay=0.0, amsgrad=False)

    # Let's train the model using RMSprop

    model.compile(loss='binary_crossentropy',

                  optimizer=opt,

                  metrics=['accuracy'])

    encoder = []

  model.summary()

  return model, encoder

def TrainTestVal(model, feats, batch_size=2, 

                train_epochs=20, show_plots=True):

  print('Training Model:')

  # Train Model

  if feats.model_type == 'AUTO' or feats.model_type == 'AUTODeep':

    print('Training autoencoder:')

    history = model.fit(feats.x_train, feats.x_train,

                        batch_size = batch_size,

                        epochs=train_epochs,

                        validation_data=(feats.x_val,feats.x_val),

                        shuffle=True,

                        verbose=True,

                        class_weight=feats.class_weights

                       )

    # list all data in history

    print(history.history.keys())

    if show_plots:

      # summarize history for loss

      plt.semilogy(history.history['loss'])

      plt.semilogy(history.history['val_loss'])

      plt.title('model loss')

      plt.ylabel('loss')

      plt.xlabel('epoch')

      plt.legend(['train', 'val'], loc='upper left')

      plt.show()

  else:

    history = model.fit(feats.x_train, feats.y_train,

              batch_size=batch_size,

              epochs=train_epochs,

              validation_data=(feats.x_val, feats.y_val),

              shuffle=True,

              verbose=True,

              class_weight=feats.class_weights

              )

    # list all data in history

    print(history.history.keys())

    if show_plots:

      # summarize history for accuracy

      plt.plot(history.history['acc'])

      plt.plot(history.history['val_acc'])

      plt.title('model accuracy')

      plt.ylabel('accuracy')

      plt.xlabel('epoch')

      plt.legend(['train', 'val'], loc='upper left')

      plt.show()

      # summarize history for loss

      plt.semilogy(history.history['loss'])

      plt.semilogy(history.history['val_loss'])

      plt.title('model loss')

      plt.ylabel('loss')

      plt.xlabel('epoch')

      plt.legend(['train', 'val'], loc='upper left')

      plt.show()

    # Test on left out Test data

    score, acc = model.evaluate(feats.x_test, feats.y_test,

                                batch_size=batch_size)

    print(model.metrics_names)

    print('Test loss:', score)

    print('Test accuracy:', acc)

    # Build a dictionary of data to return

    data = {}

    data['score'] = score

    data['acc'] = acc

    return model, data