import scipy.io
import numpy as np
import random
import matplotlib.pyplot as plt

from sklearn import preprocessing
from keras.regularizers import l2
from keras.layers import Dense, Activation, Dropout, Flatten
from keras.utils.np_utils import to_categorical
from keras.optimizers import SGD, Nadam
from keras.models import Sequential
from keras.datasets import mnist
from keras.callbacks import ModelCheckpoint, EarlyStopping
from keras.models import model_from_json
from keras.callbacks.callbacks import Callback

# use the identical train/test data as in Scientific Reports paper
# link: https://www.nature.com/articles/s41598-020-61588-w
data = scipy.io.loadmat('../../data/nsclc/Data_all.mat')
train_x = data['train_x'].astype('float32')
train_c = data['clinical_train_x'].astype('float32')
train_y = data['train_y'].astype(int)
test_x  = data['test_x'].astype('float32')
test_c  = data['clinical_test_x'].astype('float32')
test_y  = data['test_y'].astype(int)

# scaler_x = preprocessing.StandardScaler().fit(train_x.astype(np.float32))
# train_x = scaler_x.transform(train_x.astype(np.float32))
# test_x = scaler_x.transform(test_x.astype(np.float32))
# scaler_c = preprocessing.StandardScaler().fit(train_c.astype(np.float32))
# train_c = scaler_c.transform(train_c.astype(np.float32))
# test_c = scaler_c.transform(test_c.astype(np.float32))

# change this to train/test_c for another branch of bimodal subnetwork
X_train = train_x
X_test  = test_x
n_classes = 2
train_y = train_y -1
test_y = test_y -1
y_train = to_categorical(train_y,n_classes)
y_test = to_categorical(test_y,n_classes)

# subnetwork model (follows the identical settings in the paper)
model = Sequential()
model.add(Dense(40,input_dim=np.shape(X_train)[1],W_regularizer=l2(0.001),name='pre_x_input'))
model.add(Activation('relu'))
model.add(Dense(40,W_regularizer=l2(0.001),name='pre_x_h1'))
model.add(Activation('relu'))
model.add(Dense(40,W_regularizer=l2(0.001),name='pre_x_h2'))
model.add(Activation('relu'))
model.add(Dense(40,W_regularizer=l2(0.001),name='pre_x_h3'))
model.add(Activation('relu'))
model.add(Dense(n_classes,name='pre_x_out'))
model.add(Activation('softmax'))
model.summary()

nadam = Nadam(lr=0.006,beta_1=0.9,beta_2=0.999,epsilon=1e-08,schedule_decay=0.004)
ES = EarlyStopping(monitor='val_loss',patience=30,verbose=0,mode='min')
checkpointer = ModelCheckpoint(
    filepath="../../model/bimodal/nsclc_weights_x.hdf5",
    verbose=1,
    save_best_only=True,
    monitor="val_loss",
    mode="min")
model.compile(
    loss='categorical_crossentropy',
    optimizer=nadam,
    metrics=['accuracy']
)

loss_test = []
class TestCallback(Callback):
    def __init__(self, test_data):
        self.test_data = test_data

    def on_epoch_end(self, epoch, logs={}):
        x, y = self.test_data
        loss, acc = self.model.evaluate(x, y, verbose=0)
        loss_test.append(loss)
        #print('\nTesting loss: {}, acc: {}\n'.format(loss, acc))

batch_size = 20
n_epochs = 100
hist = model.fit(X_train,
                 y_train,
                 batch_size=batch_size,
                 nb_epoch=n_epochs,
				 validation_split=0.25,
				 callbacks=[ES,checkpointer,TestCallback((test_x,y_test))]
                )

loss,accuracy = model.evaluate(X_test,y_test)
print('loss:', loss)
print('accuracy:', accuracy)

with open("test.json", "w") as f:
    f.write(model.to_json())
with open("test.json", "r") as f:
    model = model_from_json(f.read())
model.load_weights("../../model/bimodal/nsclc_weights_x.hdf5")

print("drawing the training process...")
TrainERR = hist.history['loss']
ValidERR = hist.history['val_loss']
plt.figure(1)
plt.plot(TrainERR, 'b', label='TrainERR')
# plt.plot(ValidERR, 'r', label='ValidERR')
plt.plot(loss_test, 'g', label='TestERR')
plt.legend()
plt.xlabel('epoch')
plt.ylabel('error')
plt.grid(True)
plt.show()