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

from dipy.data import fetch_sherbrooke_3shell
from matplotlib import pyplot

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()

input_dir = 'sherbrooke_3shell'
# Get input data.
fetch_sherbrooke_3shell()
fdwi = os.path.join(input_dir, 'HARDI193.nii.gz')
fbval = os.path.join(input_dir, 'HARDI193.bval')
fbvec = os.path.join(input_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
pyplot.figure('Showing the datasets')
pyplot.subplot(1, 2, 1).set_axis_off()
pyplot.imshow(data[:, :, axial_middle, 0].T, cmap='gray', origin='lower')
pyplot.subplot(1, 2, 2).set_axis_off()
pyplot.imshow(data[:, :, axial_middle, 10].T, cmap='gray', origin='lower')
pyplot.savefig('data.png', bbox_inches='tight')
shutil.move('data.png', args.output)