import os, time, math
import nibabel as nib, numpy as np
import tensorflow as tf
import glob
from image_utils import *

def deeplearningseg(model_path, test_dir, atlas_dir):           
     
   with tf.Session() as sess:
       
        sess.run(tf.global_variables_initializer())

        # Import the computation graph and restore the variable values
        saver = tf.train.import_meta_graph('{0}.meta'.format(model_path))
        saver.restore(sess, '{0}'.format(model_path))
        
        # Process each subject subdirectory
        table_time = []


        if os.path.exists('{0}/subjnames.txt'.format(test_dir)):
            os.system('rm {0}/*.txt'.format(test_dir))
        os.system('touch {0}/subjnames.txt'.format(test_dir))
        for data in sorted(os.listdir(test_dir)):
            
            print(data)
            
            data_dir = os.path.join(test_dir, data)
            
            if not os.path.isdir(data_dir):
                print('  {0} is not a valid directory, Skip'.format(data_dir))
                continue
            
            file = open('{0}/subjnames.txt'.format(test_dir),'a')
            file.write('{0}\n'.format(data))
            file.close()
            
            if os.path.exists('{0}/PHsegmentation_ED.gipl'.format(data_dir)):
                os.system('rm {0}/*.gipl'.format(data_dir))
            if os.path.exists('{0}/lvsa_.nii.gz'.format(data_dir)):
                os.system('rm {0}/lvsa_*.nii.gz'.format(data_dir))
                os.system('rm {0}/seg_*.nii.gz'.format(data_dir))
            
            originalnii = glob.glob('{0}/*.nii'.format(data_dir))     
            if not originalnii:
                print('  original nifit image does not exist, use lvsa.nii.gz')
                originalnii = glob.glob('{0}/*.nii.gz'.format(data_dir))  
                imagePreprocessing(originalnii[0], data_dir, atlas_dir) 
            else:
                print('  start image preprocessing ...')
                imagePreprocessing(originalnii[0], data_dir, atlas_dir)
            
            # Process ED and ES time frames
            image_ED_name = '{0}/lvsa_{1}.nii.gz'.format(data_dir, 'ED')
            image_ES_name = '{0}/lvsa_{1}.nii.gz'.format(data_dir, 'ES')
   
            if not os.path.exists(image_ED_name) or not os.path.exists(image_ES_name):
                print(' Image {0} or {1} does not exist. Skip.'.format(image_ED_name, image_ES_name))
                continue
                               
            if os.path.exists('{0}/{1}'.format(data_dir, 'dofs')) or \
               os.path.exists('{0}/{1}'.format(data_dir, 'segs')) or \
               os.path.exists('{0}/{1}'.format(data_dir, 'tmps')) or \
               os.path.exists('{0}/{1}'.format(data_dir, 'sizes')) or \
               os.path.exists('{0}/{1}'.format(data_dir, 'motion')) or \
               os.path.exists('{0}/{1}'.format(data_dir, 'vtks')):
                    
                os.system('rm -rf {0}/{1}'.format(data_dir, 'dofs'))
                os.system('rm -rf {0}/{1}'.format(data_dir, 'segs'))
                os.system('rm -rf {0}/{1}'.format(data_dir, 'tmps'))
                os.system('rm -rf {0}/{1}'.format(data_dir, 'sizes'))
                os.system('rm -rf {0}/{1}'.format(data_dir, 'motion'))
                os.system('rm -rf {0}/{1}'.format(data_dir, 'vtks'))
                
                os.mkdir('{0}/{1}'.format(data_dir, 'dofs'))
                os.mkdir('{0}/{1}'.format(data_dir, 'segs'))
                os.mkdir('{0}/{1}'.format(data_dir, 'tmps'))
                os.mkdir('{0}/{1}'.format(data_dir, 'sizes'))
                os.mkdir('{0}/{1}'.format(data_dir, 'motion'))
                os.mkdir('{0}/{1}'.format(data_dir, 'vtks'))
                
            else: 
                
                os.mkdir('{0}/{1}'.format(data_dir, 'dofs'))
                os.mkdir('{0}/{1}'.format(data_dir, 'segs'))
                os.mkdir('{0}/{1}'.format(data_dir, 'tmps'))
                os.mkdir('{0}/{1}'.format(data_dir, 'sizes'))
                os.mkdir('{0}/{1}'.format(data_dir, 'motion'))
                os.mkdir('{0}/{1}'.format(data_dir, 'vtks'))

            for fr in ['ED', 'ES']:
       
                image_name = '{0}/lvsa_{1}.nii.gz'.format(data_dir, fr)

                # Read the image
                print('  Reading {} ...'.format(image_name))
                nim = nib.load(image_name)
                image = nim.get_data()

                imageOrg = np.squeeze(image, axis=-1).astype(np.int16)
                tmp = imageOrg

                X, Y, Z = image.shape[:3]
                
                print('  Segmenting {0} frame ...'.format(fr))
              
                # print('  Segmenting {0} frame {1} ...'.format(fr, slice))
                start_seg_time = time.time()
                
                for slice in range(Z):
                    
                    image = imageOrg[:,:,slice]
                    
                    if image.ndim == 2:
                        image = np.expand_dims(image, axis=2)
                        
                        # Intensity rescaling
                        image = rescale_intensity(image, (1, 99))
                        # Pad the image size to be a factor of 16 so that the downsample and upsample procedures
                        # in the network will result in the same image size at each resolution level.
                        X2, Y2 = int(math.ceil(X / 16.0)) * 16, int(math.ceil(Y / 16.0)) * 16
                        x_pre, y_pre = int((X2 - X) / 2), int((Y2 - Y) / 2)
                        x_post, y_post = (X2 - X) - x_pre, (Y2 - Y) - y_pre
                        image = np.pad(image, ((x_pre, x_post), (y_pre, y_post), (0, 0)), 'constant')
                        
                        # Transpose the shape to NXYC
                        image = np.transpose(image, axes=(2, 0, 1)).astype(np.float32)
                        image = np.expand_dims(image, axis=-1)
                        
                        # Evaluate the networ
                        prob, pred = sess.run(['probE:0', 'predR:0'], feed_dict={'image:0': image, 'training:0': False})
                        
                        # Transpose and crop the segmentation to recover the original size
                        pred = np.transpose(pred, axes=(1, 2, 0))
                        
                        pred = pred[x_pre:x_pre + X, y_pre:y_pre + Y]
                        pred = np.squeeze(pred, axis=-1).astype(np.int16)
                        tmp[:,:,slice] = pred
                    
                seg_time = time.time() - start_seg_time
                print('  Segmentation time = {:3f}s'.format(seg_time))
                table_time += [seg_time]

                pred = tmp
        
                nim2 = nib.Nifti1Image(pred, nim.affine)
                nim2.header['pixdim'] = nim.header['pixdim']
                nib.save(nim2, '{0}/segs/seg_lvsa_{1}.nii.gz'.format(data_dir, fr))

        print('Average segmentation time = {:.3f}s per frame'.format(np.mean(table_time)))