# ==============================================================================

# Copyright (C) 2020 Vladimir Juras, Ravinder Regatte and Cem M. Deniz

#

# This file is part of 2019_IWOAI_Challenge

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published

# by the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License

# along with this program.  If not, see <https://www.gnu.org/licenses/>.

# ==============================================================================

import tensorflow as tf

import tf_utilities as tfut

import tf_layers as tflay

import models

import nibabel as nib

import numpy as np

import re

import time

import os

from functools import partial

from pathlib import Path

import h5py

from scipy import ndimage

from scipy.spatial import distance

import glob

from sklearn.model_selection import StratifiedKFold

from sys import platform

from sklearn.preprocessing import label_binarize

tf.app.flags.DEFINE_string('model_path', "./InferenceModel", 'Name of output folder.')

tf.app.flags.DEFINE_string('data_folder', './data', 'Data Folder')

tf.app.flags.DEFINE_integer('cv', -1, 'which fold to run')

tf.app.flags.DEFINE_integer('feature', 16, 'which fold to run')

tf.app.flags.DEFINE_string('model', '4atrous248', 'Model name.')

tf.app.flags.DEFINE_integer('reso', 384, 'Image size.')

tf.app.flags.DEFINE_integer('slices', 160, 'Number Of Slices')

tf.app.flags.DEFINE_integer('seed', 1234, 'Graph-level random seed.')

tf.app.flags.DEFINE_boolean('resnet', False, 'Whether to use resnet shortcut.')

tf.app.flags.DEFINE_integer('noImages', -1, 'how many images to infer')

FLAGS = tf.app.flags.FLAGS

num_classes = 7

num_CV =1

num_channels = 1

def main(argv=None):

    print('OUT:: ',FLAGS.feature,FLAGS.seed, FLAGS.resnet,FLAGS.model,FLAGS.reso)

    tf.set_random_seed(FLAGS.seed)

    np.random.seed(FLAGS.seed)

    batch_x = tf.placeholder(tf.float32, shape=(None, FLAGS.reso, FLAGS.reso, FLAGS.slices, 1))

    batch_y = tf.placeholder(tf.float32, shape=(None, FLAGS.reso, FLAGS.reso, FLAGS.slices, num_classes))

    keep_prob = tf.placeholder(tf.float32, shape=[])

    class_weights = tf.placeholder(tf.float32, shape=(num_classes))

    # choose the model

    inference_raw = {'4unet': partial(models.inference_unet4),# the original architecture and use 4 layers 

                      # replace the convolution operations between down-convolution and up-convolution layers 

                      # by atrous convolution

                     '4atrous248': partial(models.inference_atrous4, n_class=num_classes, dilation_rates=[2,4,8])}[FLAGS.model]

    inference = partial(inference_raw, resnet=FLAGS.resnet)

    score = inference(batch_x, features_root=FLAGS.feature, keep_prob=keep_prob, n_class=num_classes)

    logits = tf.nn.softmax(score)

    # load dataset from folder

    dataFolder = FLAGS.data_folder + '/test'

    pathNifti = Path(dataFolder)

    X = []  # create an empty list

    for fileList in list(pathNifti.glob('**/*.im')):

        X.append(fileList)

    X = sorted(X)

    if FLAGS.noImages ==-1:

        noOfFiles = len(X)

    else:

        noOfFiles = FLAGS.noImages

    list_X = list( X[i] for i in range(noOfFiles) )

    n_samples = len(list_X)

    X_test, train_info = tfut.loadData_list_h5_image(list_X,num_channels)

    X_test = tfut.zeroMeanUnitVariance(X_test)

    X_test = X_test[...,np.newaxis]

    cv=1

    sample_size = X_test.shape[0]

    output_path = FLAGS.model_path

    # find the model to read

    fdr = Path(output_path)

    cpktFile = sorted(fdr.glob(('**/*.meta')))

    read_file = str((cpktFile[-1]))

    with tf.Session() as sess:

        sess.run(tf.global_variables_initializer())

        saver = tf.train.Saver(max_to_keep=0)

        saver.restore(sess, read_file[:-5])

        print('Model restored from file: %s' % read_file[:-5])

        start = time.clock()

        y_out = np.zeros((FLAGS.reso, FLAGS.reso, FLAGS.slices, num_classes))

        for xi in range(sample_size):

            prob = sess.run(logits,

                                feed_dict={batch_x: X_test[xi:xi+1],

                                            keep_prob:1})

            y_out=prob

            winOut = np.zeros(y_out.shape)

            winOut[y_out[...,1]>0.5,...,1] =1

            winOut[y_out[...,2]>0.5,...,2] =2

            winOut[y_out[...,3]>0.5,...,3] =3

            winOut[y_out[...,4]>0.5,...,4] =4

            winOut[y_out[...,5]>0.5,...,5] =5

            winOut[y_out[...,6]>0.1,...,6] =6

            # place to keep only largest connected volume    

            if 1:

                for iii in range(1,7):

                    inn = winOut[...,iii]

                    all_labels, num_features = ndimage.label(inn)

                    print('Label #:',iii, num_features,'Number of Connected Volumes')

                    if num_features > 1:

                        volume = ndimage.sum(inn, all_labels, index=range(num_features+1))

                        print("Volume:", volume)

                        cem = all_labels == np.argmax(volume)

                        winOut[...,iii] = winOut[...,iii] * cem

            winOut =  np.sum(winOut,axis=4)

            winOut.astype(int)

            seg = np.pad(np.squeeze(winOut),((1,1),(1,1),(1,1)),'edge')

            #For the classes dimensions, the order for 4 classes are as the following: 

            #0 = femoral cartilage, 1 = tibial cartilage, 2 = patellar cartilage, and 3 = meniscus.

            #save as numpy array

            saveNumpy = np.zeros((384,384,160,4))

            saveNumpy[seg==1,...,0] = 1

            saveNumpy[seg==2,...,1] = 1

            saveNumpy[seg==3,...,1] = 1

            saveNumpy[seg==4,...,2] = 1

            saveNumpy[seg==5,...,3] = 1

            saveNumpy[seg==6,...,3] = 1

            savename= str(X[xi])

            fdr = Path('./InferenceResults/%s.npy' % (savename[-15:-3]))

            np.save(fdr, saveNumpy.astype(int), allow_pickle = False)

if __name__ == '__main__':

    tf.app.run()