"""
This code is to
1. Create train & test input to Network as numpy arrays
2. Load the train & test numpy arrays
"""

import numpy as np
from utils import *

# data type to save as np array
npdtype = np.float32


def create_train_data(current_fold, plane):
    """
    Crop each slice by its ground truth bounding box,
    then pad zeros to form uniform dimension,
    rescale pixel intensities to [0,1]
    """
    # get the list of image and label number of current_fold
    imlb_list = open(training_set_filename(current_fold), 'r').read().splitlines()
    current_fold = current_fold
    training_image_set = np.zeros((len(imlb_list)), dtype = np.int)

    for i in range(len(imlb_list)):
        s = imlb_list[i].split(' ')
        training_image_set[i] = int(s[0])

    slice_list = open(list_training[plane], 'r').read().splitlines()
    slices = len( slice_list)
    image_ID = np.zeros(( slices), dtype = np.int)
    slice_ID = np.zeros(( slices), dtype = np.int)
    image_filename = ['' for l in range( slices)]
    label_filename = ['' for l in range( slices)]
    pixels = np.zeros((slices), dtype = np.int)

    for l in range(slices):
        s =  slice_list[l].split(' ')
        image_ID[l] = s[0]
        slice_ID[l] = s[1]
        image_filename[l] = s[2]
        label_filename[l] = s[3]
        pixels[l] = int(s[organ_ID * 5])

    create_slice_list = []
    create_label_list = []

    for l in range(slices):
        if image_ID[l] in training_image_set and pixels[l] >= 100:
            create_slice_list.append(image_filename[l])
            create_label_list.append(label_filename[l])
    if len(create_slice_list)!= len(create_label_list):
        raise ValueError('slice number does not equal label number!')

    total = len(create_slice_list)

    img_rows = XMAX
    img_cols = YMAX

    imgs = np.ndarray((total, img_rows, img_cols), dtype = npdtype)
    imgs_mask = np.ndarray((total, img_rows, img_cols), dtype = npdtype)

    print('-'*30)
    print('  Creating training data...')
    print('-'*30)

    for i in range(len(create_slice_list)):
        cur_im = np.load(create_slice_list[i])
        cur_mask = np.load(create_label_list[i])

        cur_im = (cur_im - low_range) / float(high_range - low_range)
        arr = np.nonzero(cur_mask)

        width = cur_mask.shape[0]
        height = cur_mask.shape[1]

        minA = min(arr[0])
        maxA = max(arr[0])
        minB = min(arr[1])
        maxB = max(arr[1])

        # with margin
        cropped_im = cur_im[max(minA - margin, 0): min(maxA + margin + 1, width), \
                                    max(minB - margin, 0): min(maxB + margin + 1, height)]
        cropped_mask = cur_mask[max(minA - margin, 0): min(maxA + margin + 1, width), \
                                    max(minB - margin, 0): min(maxB + margin + 1, height)]

        imgs[i] = pad_2d(cropped_im, plane, 0, XMAX, YMAX, ZMAX)
        imgs_mask[i] = pad_2d(cropped_mask, plane, 0, XMAX, YMAX, ZMAX)

        if i % 100 == 0:
            print('Done: {0}/{1} slices'.format(i, total))

    np.save('imgs_train_%s_%s.npy'%(current_fold, plane), imgs)
    np.save('masks_train_%s_%s.npy'%(current_fold, plane), imgs_mask)
    print('Training data created for fold %s, plane %s'%(current_fold, plane))


def load_train_data(current_fold, plane):
    imgs_train = np.load('imgs_train_%s_%s.npy'%(current_fold, plane))
    mask_train = np.load('masks_train_%s_%s.npy'%(current_fold, plane))
    return imgs_train, mask_train


def load_test_data(current_fold, plane):
    imgs_test = np.load('imgs_test_%s_%s.npy'%(current_fold, plane))
    mask_test = np.load('masks_test_$s_%s.npy'%(current_fold, plane))
    return imgs_test, mask_test


if __name__ == '__main__':

    data_path = sys.argv[1]
    current_fold = int(sys.argv[2])
    plane = sys.argv[3]

    # dim of each case (after padding zeors to max gt bounding box)
    ZMAX = int(sys.argv[4])
    YMAX = int(sys.argv[5])
    XMAX = int(sys.argv[6])

    margin = int(sys.argv[7])
    organ_ID = int(sys.argv[8])
    low_range = int(sys.argv[9])
    high_range = int(sys.argv[10])

    create_train_data(current_fold, plane)