#!/usr/bin/env python
import os
from os import path
import sys
import numpy as np
import nibabel as nib
import argparse
import pathlib

# Function to preprocess the dwi cases
def process_trainingdata(dwib0_arr):
    count = 0
    for b0 in dwib0_arr:
        img = nib.load(b0)
        imgF32 = img.get_fdata().astype(np.float32)
        ''' Intensity based segmentation of MR images is hampered by radio frerquency field
            inhomogeneity causing intensity variation. The intensity range is typically
            scaled between the highest and lowest signal in the Image. Intensity values
            of the same tissue can vary between scans. The pixel value in images must be
            scaled prior to providing the images as input to CNN. The data is projected in to
            a predefined range [0,1] '''
        p = np.percentile(imgF32, 99)
        imgF32_sagittal = imgF32 / p 			          # sagittal view
        imgF32_sagittal[ imgF32_sagittal < 0 ] =  sys.float_info.epsilon	      
        imgF32_sagittal[ imgF32_sagittal > 1 ] = 1 		  
        imgF32_coronal = np.swapaxes(imgF32_sagittal,0,1) # coronal view
        imgF32_axial = np.swapaxes(imgF32_sagittal,0,2)   # Axial view

        # dwi volume data is written to the binary file
        imgF32_sagittal.tofile(sagittal_f_handle)
        imgF32_coronal.tofile(coronal_f_handle)
        imgF32_axial.tofile(axial_f_handle)

        print('Case ' + str(count) + ' done')
        count = count + 1

    # Closing the binary file
    sagittal_f_handle.close()
    axial_f_handle.close()
    coronal_f_handle.close()

# parser module for input arguments
SUFFIX_TXT = "txt"
parser = argparse.ArgumentParser()
parser.add_argument('-i', action='store', dest='b0', type=str,
                        help="txt file containing list of /path/to/b0, one path in each line")
args = parser.parse_args()

try:
    args = parser.parse_args()
    if len(sys.argv) == 1:
        parser.print_help()
        parser.error('too few arguments')
        sys.exit(0)

except SystemExit:
    sys.exit(0)

if args.b0:
    f = pathlib.Path(args.b0)
    if f.exists():
        print ("File exist")
        filename = args.b0
    else:
        print ("File not found")
        sys.exit(1)

    # Input caselist.txt
    if filename.endswith(SUFFIX_TXT):
        with open(filename) as f:
            dwib0_arr = f.read().splitlines()

storage = path.dirname(dwib0_arr[0])
# dwi cases will be written to the below binary files
sagittal_bin_file = storage + '/sagittal-binary-dwi'
coronal_bin_file = storage + '/coronal-binary-dwi'
axial_bin_file = storage + '/axial-binary-dwi'

# The above binary files will be converted to 3D numpy array
sagittal_trainingdata = storage + '/sagittal-traindata-dwi.npy'
coronal_trainingdata = storage + '/coronal-traindata-dwi.npy'
axial_trainingdata = storage + '/axial-traindata-dwi.npy'

# Open the binary file for writing
sagittal_f_handle = open(sagittal_bin_file, 'wb')
coronal_f_handle = open(coronal_bin_file, 'wb')
axial_f_handle = open(axial_bin_file, 'wb')

process_trainingdata(dwib0_arr)
x_dim=len(dwib0_arr)*256
y_dim=256
z_dim=256

# Open the binary file and convert it to 3D numpy array
merge_sagittal = np.memmap(sagittal_bin_file, dtype=np.float32, mode='r+', shape=(x_dim, y_dim, z_dim))
print("Saving sagittal training data to disk")
np.save(sagittal_trainingdata, merge_sagittal)
os.unlink(sagittal_bin_file)

merge_coronal = np.memmap(coronal_bin_file, dtype=np.float32, mode='r+', shape=(x_dim, y_dim, z_dim))
print("Saving coronal training data to disk")
np.save(coronal_trainingdata, merge_coronal)
os.unlink(coronal_bin_file)

merge_axial = np.memmap(axial_bin_file, dtype=np.float32, mode='r+', shape=(x_dim, y_dim, z_dim))
print("Saving axial training data to disk")
np.save(axial_trainingdata, merge_axial)
os.unlink(axial_bin_file)