#!/usr/bin/env python
import argparse
import os
import nibabel
import shutil

from dipy.data import fetch_sherbrooke_3shell

parser = argparse.ArgumentParser()
parser.add_argument('--input', dest='input', help='Input dataset')
parser.add_argument('--output', dest='output', help='Output dataset')

args = parser.parse_args()

def move_directory_files(source_dir, destination_dir, copy=False, remove_source_dir=False):
    source_directory = os.path.abspath(source_dir)
    destination_directory = os.path.abspath(destination_dir)
    if not os.path.isdir(destination_directory):
        os.makedirs(destination_directory)
    for dir_entry in os.listdir(source_directory):
        source_entry = os.path.join(source_directory, dir_entry)
        if copy:
            shutil.copy(source_entry, destination_directory)
        else:
            shutil.move(source_entry, destination_directory)
    if remove_source_dir:
        os.rmdir(source_directory)
source_dir = os.path.abspath(os.path.join('home', 'greg', '.dipy', 'sherbrooke_3shell'))
dest_dir = 'input_dir'
if not os.path.exists(source_dir):
    # Get input data.
    fetch_sherbrooke_3shell()
# Move inputs to working directory.
move_directory_files(source_dir, dest_dir, copy=True)
fdwi = os.path.join(dest_dir, 'HARDI193.nii.gz')
fbval = os.path.join(dest_dir, 'HARDI193.bval')
fbvec = os.path.join(dest_dir, 'HARDI193.bvec')
# Load the dMRI datasets.
img = nibabel.load(fdwi)
data = img.get_data()
# data is a 4D array where the first 3 dimensions are the i, j,
# k voxel coordinates and the last dimension is the number of
# non-weighted (S0s) and diffusion-weighted volumes.
# Visualize the results using matplotlib.
axial_middle = data.shape[2] // 2
plt.figure('Showing the datasets')
plt.subplot(1, 2, 1).set_axis_off()
plt.imshow(data[:, :, axial_middle, 0].T, cmap='gray', origin='lower')
plt.subplot(1, 2, 2).set_axis_off()
plt.imshow(data[:, :, axial_middle, 10].T, cmap='gray', origin='lower')
plt.show()
plt.savefig(args.output, bbox_inches='tight')