#!/usr/bin/env python

from __future__ import division, print_function

import os, glob

import argparse

from PIL import Image

import patient

def generate_unet_data(patient_dirs, x_dir, y_dir, file_format="{:04d}"):

    BW_8BIT = 'L'

    i = 0

    for patient_dir in patient_dirs:

        p = patient.PatientData(patient_dir)

        for image, endo in zip(p.images, p.endocardium_masks):

            x_outfile = os.path.join(x_dir, file_format.format(i) + ".jpg")

            Image.fromarray(image, BW_8BIT).save(x_outfile)

            y_outfile = os.path.join(y_dir, file_format.format(i) + ".png")

            Image.fromarray(endo, BW_8BIT).save(y_outfile)

            i += 1

def main(args):

    glob_search = os.path.join(args.indir, "patient*")

    patient_dirs = glob.glob(glob_search)

    print("Found patient directories:")

    for patient_dir in patient_dirs:

        print(patient_dir)

    split_index = int((1-args.split/100) * len(patient_dirs))

    train_dirs = patient_dirs[:split_index]

    test_dirs = patient_dirs[split_index:]

    print("First {} patients used as training set, remaining as test.".format(split_index))

    x_train = os.path.join(args.outdir, "train/img/0")

    y_train = os.path.join(args.outdir, "train/gt/0")

    os.makedirs(x_train)

    os.makedirs(y_train)

    x_test = os.path.join(args.outdir, "test/img/0")

    y_test = os.path.join(args.outdir, "test/gt/0")

    os.makedirs(x_test)

    os.makedirs(y_test)

    generate_unet_data(train_dirs, x_train, y_train)

    generate_unet_data(test_dirs, x_test, y_test)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="Generate data for U-Net from RV MRI dicom images.")

    parser.add_argument('indir', default='.', help="TrainingSet/ directory")

    parser.add_argument('-o', '--outdir', default='.', help="Directory to write output data")

    parser.add_argument('-s', '--split', default=20, type=int, help='Percentage of patients used for test set')

    args = parser.parse_args()

    main(args)