"""

This code is to test NN model and visualize output

"""

import numpy as np

import sys

import time

import matplotlib.pyplot as plt

from keras.models import Model, load_model

from keras.layers import Input, Activation, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose, UpSampling2D, ZeroPadding2D, BatchNormalization

from keras.optimizers import Adam, SGD

from keras.callbacks import ModelCheckpoint

from keras import backend as K

import tensorflow as tf

from data import load_train_data, load_test_data

from utils import *

K.set_image_data_format('channels_last')  # Tensorflow dimension ordering

data_path  = sys.argv[1] + "/"

model_path = data_path + "models/"

# dir for storing results that contains

rst_path = data_path + "test-records/"

if not os.path.exists(rst_path):

    os.makedirs(rst_path)

model_to_test = sys.argv[2]

cur_fold = sys.argv[3]

plane = sys.argv[4]

im_z = int(sys.argv[5])

im_y = int(sys.argv[6])

im_x = int(sys.argv[7])

high_range = float(sys.argv[8])

low_range = float(sys.argv[9])

margin = int(sys.argv[10])

vis = sys.argv[11]

# prediction of trained model

pred_path = os.path.join(rst_path, "pred-%s/"%cur_fold)

if not os.path.exists(pred_path):

    os.makedirs(pred_path)

"""

Dice Ceofficient and Cost functions for training

"""

smooth = 1.

def dice_coef(y_true, y_pred):

    y_true_f = K.flatten(y_true)

    y_pred_f = K.flatten(y_pred)

    intersection = K.sum(y_true_f * y_pred_f)

    return (2.0 * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):

    return  -dice_coef(y_true, y_pred)

def test(model_to_test, current_fold, plane, rst_dir, vis):

    print "-"*50

    print "loading model ", model_to_test

    print "-"*50

    model = load_model(model_path + model_to_test + '.h5', custom_objects={'dice_coef_loss': dice_coef_loss, 'dice_coef':dice_coef})

    volume_list = open(testing_set_filename(current_fold), 'r').read().splitlines()

    total = len(volume_list)

    dsc = np.zeros((total, 2))

    # iterate all test cases

    for i in range(total):

        s = volume_list[i].split(' ')

        image = np.load(s[1])

        label = np.load(s[2])

        case_num = s[1].split("00")[1].split(".")[0]

        print "testing case: ", case_num

        image_ = np.transpose(image, (2, 0, 1))

        label_ = np.transpose(label, (2, 0, 1))

        # standardize test data

        image_[image_ < low_range] = low_range

        image_[image_ > high_range] = high_range

        image_ = (image_ - low_range) / float(high_range - low_range)

        # for creating final prediction visualization

        pred = np.zeros_like(image_)

        for sli in range(label_.shape[0]):

            try:

                # crop each slice according to smallest bounding box of each slice

                width = label_[sli].shape[0]

                height = label_[sli].shape[1]

                arr = np.nonzero(label_[sli])

                if len(arr[0]) == 0:

                    continue

                minA = min(arr[0])

                maxA = max(arr[0])

                minB = min(arr[1])

                maxB = max(arr[1])

                minAdiff = margin

                maxAdiff = margin

                minBdiff = margin

                maxBdiff = margin

                cropped = image_[sli, max(minA - minAdiff, 0): min(maxA + maxAdiff + 1, width), \

                        max(minB - minBdiff, 0): min(maxB + maxBdiff + 1, height)]

                cropped_mask = label_[sli, max(minA - minAdiff, 0): min(maxA + maxAdiff + 1, width), \

                        max(minB - minBdiff, 0): min(maxB + maxBdiff + 1, height)]

                image_padded_ = pad_2d(cropped, plane, 0, im_x, im_y, im_z)

                mask_padded_ = pad_2d(cropped_mask, plane, 0, im_x, im_y, im_z)

                image_padded_prep = preprocess_front(preprocess(image_padded_))

                out_ori = (model.predict(image_padded_prep) > 0.5).astype(np.uint8)

                out = out_ori[:,0:cropped.shape[0], 0:cropped.shape[1],:].reshape(cropped.shape)

                pred[sli, max(minA - minAdiff, 0): min(maxA + maxAdiff + 1, width), max(minB - minBdiff, 0): min(maxB + maxBdiff+ 1, height)] = out

                pred_vis = pred[sli, max(minA - minAdiff, 0): min(maxA + maxAdiff + 1, width), max(minB - minBdiff, 0): min(maxB + maxBdiff+ 1, height)]

                if vis == "true":

                    fig = plt.figure()

                    ax = fig.add_subplot(1, 3, 1)

                    ax.set_title("input test image")

                    ax.imshow(cropped, cmap=plt.cm.gray)

                    ax = fig.add_subplot(1, 3, 2)

                    ax.set_title("prediction")

                    ax.imshow(pred_vis, cmap=plt.cm.gray)

                    ax = fig.add_subplot(1, 3, 3)

                    ax.set_title("ground truth")

                    ax.imshow(cropped_mask, cmap=plt.cm.gray)

                    # plt.suptitle("slice %s"%sli)

                    fig.canvas.set_window_title("slice %s"%sli)

                    plt.axis('off')

                    plt.show()

            except KeyboardInterrupt:

                print 'KeyboardInterrupt caught'

                raise ValueError("terminate because of keyboard interruption")

        # ------------ write out for visualization ---------------

        np.save(pred_path + case_num + ".npy", pred) # prediction made by the trained model

        # compute DSC

        cur_dsc, _, _, _ = DSC_computation(label_, pred)

        print cur_dsc

        dsc[i][0] = case_num

        dsc[i][1] = cur_dsc

    dsc_mean = np.mean(dsc[:,1])

    dsc_std = np.std(dsc[:,1])

    # record test dsc mean and standard deviation for each fold in the one file

    fd = open(rst_path + 'test_stats.csv','a+')

    fd.write("%s,%s,%s,%s\n"%(cur_fold, model_to_test, dsc_mean, dsc_std))

    fd.close()

    print "---------------------------------"

    print "mean: ", dsc_mean

    print "std: ", dsc_std

    # record test result case by case

    np.savetxt(rst_path + model_to_test + ".csv", dsc, fmt = "%i, %.5f", delimiter=",", header="case_num,DSC")

if __name__ == "__main__":

    start_time = time.time()

    test(model_to_test, cur_fold, plane, rst_path, vis)

    print "-----------test done, total time used: %s ------------"% (time.time() - start_time)