from __future__ import print_function

from keras.layers import Dense, Activation, LSTM

from keras.models import Sequential

from sklearn.preprocessing import MinMaxScaler

from sklearn.metrics import mean_squared_error

import time

import math

import numpy as np

import os.path as op

import mne

import babymeg

from matplotlib import pyplot as plt

# prepare the dataset for the prediction of outer layer chanel to the inner layer channel

def prepare_dataset(oData, iData):

    dataX, dataY = [], []

    for i in range(len(oData[0])):

        a = []

        for j in range(len(oData)):

            a.append(oData[j,i])

        dataX.append(a)

        dataY.append(iData[i])

    return np.array(dataX), np.array(dataY)

#read data from a fif file with mne-python

datapath   = '/neuro/users/limin/Limin_Script/machine learning/data/'

datafile   = '170810_104407_5106943_empty_room_raw.fif'

info = mne.io.read_info(op.join(datapath,datafile), verbose=False)

print('Keys in info dictionary:\n', info.keys())

[ind_inner, labels_inner] = babymeg.channels.LM_SelLayerChan(info,"inner")

[ind_outer, labels_outer] = babymeg.channels.LM_SelLayerChan(info,"outer")

nInner = ind_inner.__len__()

nOuter = ind_outer.__len__()

raw = mne.io.read_raw_fif(op.join(datapath, datafile), preload=True)

tmin, tmax = 0, 2

raw.crop(tmin, tmax).load_data()

InnerData = raw._data[ind_inner]

OuterData = raw._data[ind_outer]

#raw.plot(block=True)

InnerData = InnerData.astype('float32')

OuterData = OuterData.astype('float32')

datalen = len(InnerData[0])

# normalize the dataset

scaler1 = MinMaxScaler(feature_range=(0, 1))

TempData = np.reshape(InnerData[0],(datalen,1))

TempData = scaler1.fit_transform(TempData)

scaler2 = MinMaxScaler(feature_range=(0, 1))

OuterData = scaler2.fit_transform(OuterData)

# split into train and test sets

train_size = int(datalen * 0.67)

test_size = datalen  - train_size

train_inner, test_inner = TempData[0:train_size,0], TempData[train_size:datalen,0]

train_outer, test_outer = OuterData[:, 0:train_size], OuterData[:, train_size:datalen]

print('Start machine learning')

start_time = time.time()

# prepare the dataset

trainX, trainY = prepare_dataset(train_outer[:,:], train_inner)

testX, testY = prepare_dataset(test_outer[:,:], test_inner)

# reshape input to be [samples, time steps, features]

trainX = np.reshape(trainX, (trainX.shape[0], trainX.shape[1]))

trainY = np.reshape(trainY, (trainY.shape[0], 1))

testX = np.reshape(testX, (testX.shape[0], testX.shape[1]))

testY = np.reshape(testY, (testY.shape[0], 1))

ninput = len(train_outer)

model = Sequential()

model.add(Dense(80, input_dim=ninput, activation='relu'))

model.add(Dense(40, input_dim= 80, activation='relu'))

model.add(Dense(40, activation='relu'))

model.add(Dense(1, activation='linear'))

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

for i in range(50):

    model.fit(trainX, trainY,  epochs=100, verbose=2, shuffle=False)

    model.reset_states()

# make predictions

trainPredict = model.predict(trainX)

model.reset_states()

testPredict = model.predict(testX)

model.reset_states()

# invert predictions

trainPredict = scaler1.inverse_transform(trainPredict)

trainY = scaler1.inverse_transform(trainY)

testPredict = scaler1.inverse_transform(testPredict)

testY = scaler1.inverse_transform(testY)

# calculate root mean squared error

trainScore = math.sqrt(mean_squared_error(trainY, trainPredict))

print('Train Score: %.2f RMSE' % (trainScore))

testScore = math.sqrt(mean_squared_error(testY, testPredict))

print('Test Score: %.2f RMSE' % (testScore))

elapsed_time = time.time() - start_time

print(elapsed_time)

time.strftime("%H:%M:%S", time.gmtime(elapsed_time))

# save model

modelfile = op.join(datapath,'noise_model.h5')

model.save(modelfile)

PredictPlot = np.empty_like(InnerData[0])

PredictPlot[:] = np.nan

PredictPlot[0:len(trainPredict)] = trainPredict[:,0]

PredictPlot[len(trainPredict):len(InnerData[0])+1] = testPredict[:,0]

OrgData = scaler1.inverse_transform(TempData)

OrgData = np.reshape(OrgData, (OrgData.shape[0]))

# plot baseline and predictions

fig1, ax1 = plt.subplots()

ax1.plot(OrgData)

ax1.plot(PredictPlot)

fig2, ax2 = plt.subplots()

ax2.plot(np.subtract(OrgData,PredictPlot))

plt.show()

print('All is done!')