#!/usr/bin/env python

import os

from os import path

import sys

import numpy as np

import nibabel as nib

import argparse

import pathlib

def process_trainingdata(mask_arr):

    count = 0

    for b0_mask in mask_arr:

        img = nib.load(b0_mask)

        imgU8_sagittal = img.get_fdata().astype(np.uint8) # sagittal view

        imgU8_sagittal[ imgU8_sagittal < 0 ] = 0        

        imgU8_sagittal[ imgU8_sagittal > 1 ] = 1        

        imgU8_coronal = np.swapaxes(imgU8_sagittal,0,1) # coronal view

        imgU8_axial = np.swapaxes(imgU8_sagittal,0,2)   # Axial view

        # dwi mask volume data is written to the binary file

        imgU8_sagittal.tofile(sagittal_f_handle)

        imgU8_coronal.tofile(coronal_f_handle)

        imgU8_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='mask', type=str,

                        help="txt file containing list of /path/to/mask, 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.mask:

    f = pathlib.Path(args.mask)

    if f.exists():

        print ("File exist")

        filename = args.mask

    else:

        print ("File not found")

        sys.exit(1)

    # Input caselist.txt

    if filename.endswith(SUFFIX_TXT):

        with open(filename) as f:

            mask_arr = f.read().splitlines()

storage = path.dirname(mask_arr[0])

# dwi cases mask will be written to the below binary files

sagittal_bin_file = storage + '/sagittal-binary-mask'

coronal_bin_file = storage + '/coronal-binary-mask'

axial_bin_file = storage + '/axial-binary-mask'

# The above binary files will be converted to 3D numpy array

sagittal_trainingdata = storage + '/sagittal-traindata-mask.npy'

coronal_trainingdata = storage + '/coronal-traindata-mask.npy'

axial_trainingdata = storage + '/axial-traindata-mask.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(mask_arr)

x_dim=len(mask_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.uint8, mode='r+', shape=(x_dim, y_dim, z_dim))

print("Saving sagittal training data mask to disk")

np.save(sagittal_trainingdata, merge_sagittal)

os.unlink(sagittal_bin_file)

merge_coronal = np.memmap(coronal_bin_file, dtype=np.uint8, mode='r+', shape=(x_dim, y_dim, z_dim))

print("Saving coronal training data mask to disk")

np.save(coronal_trainingdata, merge_coronal)

os.unlink(coronal_bin_file)

merge_axial = np.memmap(axial_bin_file, dtype=np.uint8, mode='r+', shape=(x_dim, y_dim, z_dim))

print("Saving axial training data mask to disk")

np.save(axial_trainingdata, merge_axial)

os.unlink(axial_bin_file)