# -*- coding: utf-8 -*-
"""
Created on Wed Nov 02 21:35:59 2016

@author: seeker105
"""
import os.path
import sys
from ConvNet import LeNet
import json
import SimpleITK as sitk
import pylab
from skimage import color
from sklearn.utils import shuffle
from scipy.ndimage.interpolation import rotate
from keras.optimizers import SGD
from keras.utils import np_utils
from keras.callbacks import LearningRateScheduler, ModelCheckpoint
import numpy as np
from Brain_pipeline import Pipeline
import Brain_pipeline
import Metrics
from glob import glob
import model_test

''' Script to drive loading, training, testing and saving the brain MRI
    First, we load all the images, and process them through the Pipeline,
    and get the pre-processed images as output.
    Then, we train the model from ConvNet or load the weights into it.
    We divide the training and test data using train_test_split.
    '''
def show_segmented_image(orig_img, pred_img):
    '''
    Show the prediction over the original image
    INPUT:
        1)orig_img: the test image, which was used as input
        2)pred_img: the prediction output
    OUTPUT:
        segmented image rendering
    '''
    #define the colours of the labels
    red = [10, 0, 0] #label 1
    yellow = [10, 10, 0] #label 2
    green = [0, 10, 0]  #label 3
    blue = [0, 0, 10] #label 4
    #convert original image to rgb
    gray_im = color.gray2rgb(orig_img)
    #color the tumor voxels
    gray_im[pred_img == 1] = red
    gray_im[pred_img == 2] = yellow
    gray_im[pred_img == 3] = green
    gray_im[pred_img == 4] = blue
    pylab.imshow(gray_im)


def step_decay(epochs):
    init_rate = 0.003
    fin_rate = 0.00003
    total_epochs = 24
    print 'ep: {}'.format(epochs)
    if epochs<25:
        lrate = init_rate - (init_rate - fin_rate)/total_epochs * float(epochs)
    else: lrate = 0.00003
    print 'lrate: {}'.format(model.optimizer.lr.get_value())
    return lrate

pth_train = 'D:/New folder/BRATS2015_Training/train_slices/'
pth_test = 'D:/New folder/BRATS2015_Training/test_slices/'
x = Pipeline(pth_train, pth_test)   #pass the images through the preprocessing steps

#build the model
model = LeNet.build_Pereira(33, 33, 4, 5)   

#callback
change_lr = LearningRateScheduler(step_decay)

#initialize the optimizer and model
opt = SGD(lr = 0.003, momentum=0.9, decay= 0, nesterov = True)
model.compile(loss = 'categorical_crossentropy', optimizer=opt, metrics = ['accuracy'])


#load training patches
X_patches, Y_labels, mu, sigma = x.training_patches([180000, 67500, 67500, 67500, 67500])
tmp = rotate(X_patches, 90, (2, 3))
tmp = np.append(tmp, rotate(X_patches, -90, (2, 3)), axis=0)
tmp = np.append(tmp, rotate(X_patches, 180, (2, 3)), axis=0)
X_patches = np.append(X_patches, tmp, axis=0)
Y_labels = np.hstack(Y_labels)
for i in xrange(2):
    Y_labels = np.append(Y_labels, Y_labels, axis=0)
# Labels should be in categorical array form 1x5
Y_labels = np_utils.to_categorical(Y_labels, 5)
X_patches, Y_labels = shuffle(X_patches, Y_labels, random_state=0)

#save model after each epoch
os.mkdir(r'D:\New folder\Pereira_model_checkpoints')
checkpointer = ModelCheckpoint(filepath='D:/New folder/Pereira_model_checkpoints/weights.{epoch:02d}-{val_loss:.2f}.keras2.hdf5',monitor = 'val_loss', verbose=1)
#fit model and shuffle training data
hist = model.fit(X_patches[:200000], Y_labels[:200000], nb_epoch=25, batch_size=128, verbose=1, validation_split=0.1, callbacks = [change_lr, checkpointer])
 
#save model
sv_pth = 'D:/New Folder/Pereira_model_checkpoints/model_weights'
m = '{}.json'.format(sv_pth)
w = '{}.hdf5'.format(sv_pth)
model.save_weights(w)
json_strng = model.to_json()
with open(m, 'w') as f:
    json.dump(json_strng, f)
    

#test all the test image slices
test_im = x.test_im.swapaxes(0,1)
gt = test_im[4]
test_im = test_im[:4].swapaxes(0, 1)
predicted_images, params = model_test.test_slices(test_im[158:159], gt[158:159], model, mu, sigma)

'''test_pths = zip(*x.pathnames_test)
#show a segmented slice
tst = test_pths[0]#random.choice(test_pths)
test_arr = [sitk.GetArrayFromImage(sitk.ReadImage(i)) for i in tst]
final_pth = os.path.dirname(os.path.dirname(tst[0])) +  '/' + os.path.splitext(os.path.splitext(os.path.basename(tst[0]))[0])[0] + '_processed_predicted_70.mha'  
slice_arr = [test_arr[j][70] for j in xrange(4)]
patches = Brain_pipeline.test_patches(slice_arr)
pred = model.predict_classes(patches)
pred = Brain_pipeline.reconstruct_labels(pred)
show_segmented_image(test_arr[0][70], pred)
sitk.WriteImage(sitk.GetImageFromArray(np.array(pred.astype(float))), final_pth)


#evaluate metrics
DSC_arr = [] #stores DSC
DSC_core_arr = [] #stores list of core DSCs
PPV_arr = []
acc_arr = []

#use for getting orignal brain image and prediction label slices
# use for:
    #overlay images
    #segmentation vs orig label. it's in test_paths
    #with/without nyul
    #4 sequences after nyul. for original ones, redefine paths
    #ok. now we gotta see metrics brother
pred_pth = []
t1c_pth = []
pred_arr = []

for i in xrange(len(test_pths)):
    tst = test_pths[i]
    test_arr = [sitk.GetArrayFromImage(sitk.ReadImage(j)) for j in tst]
    #take slices
    slice_arr = [test_arr[j][70] for j in xrange(4)]
    #read original slice label
    orig = test_arr[4][70]
    patches = Brain_pipeline.test_patches(slice_arr)
    pred = model.predict_classes(patches)
    pred = Brain_pipeline.reconstruct_labels(pred)
    acc_arr.append(Metrics.accuracy(pred, orig))
    DSC_arr.append(Metrics.DSC(pred, orig, 2))
    DSC_core_arr.append(Metrics.DSC_core_tumor(pred, orig))
    PPV_arr.append(Metrics.PPV(pred, orig))
    print 'acc: {}'.format(acc_arr[i])
    print 'DSC: {}'.format(DSC_arr[i])
    print 'DSC_core: {}'.format(DSC_core_arr[i])
    print 'PPV : {}'.format(PPV_arr[i])
    sys.stdout.flush()
    final_pth = os.path.dirname(tst[4]) +  '/' + os.path.splitext(os.path.basename(tst[0]))[0] + '_predicted_70.mha'  
    pred_pth.append(final_pth)
    pred_arr.append(pred)
    t1c_pth.append([flp for flp in glob(os.path.dirname(tst[2]) + '/*.mha') if 'n4' not in flp])
               '''