%matplotlib inline
from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
import nibabel as nib

def plot_slice(fname, z_idx=5):

    # Load the image and collect the data
    # and orientation information
    img = nib.load(fname)
    data = img.get_data()
    aff = img.get_affine()

    # Find the center of the brain matrix
    ctr = np.dot(np.linalg.inv(aff), [0, 0, 0, 1])[:3]

    # Plot the data
    vmin, vmax = (0, 1) if data.dtype == np.int16 else (30, 150)
    plt.imshow(np.rot90(data[:, :, ctr[2] + z_idx]),
           cmap="gray", vmin=vmin, vmax=vmax)
    plt.gca().set_axis_off()

%%bash
bet data/T1 data/T1_brain -m
fslmaths data/T1 -s 2 data/T1_smooth
fslmaths data/T1_smooth -mas data/T1_brain_mask data/T1_smooth_mask

f = plt.figure(figsize=(12, 4))
for i, img in enumerate(["T1", "T1_smooth", "T1_smooth_mask"]):
    f.add_subplot(1, 3, i + 1)
    plot_slice("data/%s.nii.gz" % img)
    plt.title(img)

!make clean

from nipype.interfaces import fsl

skullstrip = fsl.BET(in_file="data/T1.nii.gz", mask=True,
                     out_file="data/T1_brain.nii.gz")
res = skullstrip.run()

smooth = fsl.IsotropicSmooth(in_file="data/T1.nii.gz", fwhm=4,
                             out_file="data/T1_smooth.nii.gz")
smooth.run()

mask = fsl.ApplyMask(in_file="data/T1_smooth.nii.gz",
                     mask_file="data/T1_brain_mask.nii.gz",
                     out_file="data/T1_smooth_mask.nii.gz")
mask.run()

f = plt.figure(figsize=(12, 4))
for i, img in enumerate(["T1", "T1_smooth", "T1_smooth_mask"]):
    f.add_subplot(1, 3, i + 1)
    plot_slice("data/%s.nii.gz" % img)
    plt.title(img)

!make clean

skullstrip = fsl.BET(in_file="data/T1.nii.gz", mask=True)
bet_result = skullstrip.run()

smooth = fsl.IsotropicSmooth(in_file="data/T1.nii.gz", fwhm=4)
smooth_result = smooth.run()

mask = fsl.ApplyMask(in_file=smooth_result.outputs.out_file,
                     mask_file=bet_result.outputs.mask_file)
mask.run()

f = plt.figure(figsize=(12, 4))
for i, img in enumerate(["data/T1", "T1_smooth", "T1_smooth_masked"]):
    f.add_subplot(1, 3, i + 1)
    plot_slice("%s.nii.gz" % img)
    plt.title(img)

from nipype import Node, Workflow

# For reasons that will later become clear, it's important to
# pass filenames to Nodes as absolute paths
from os.path import abspath
in_file = abspath("data/T1.nii.gz")

skullstrip = Node(fsl.BET(in_file=in_file, mask=True), name="skullstrip")

smooth = Node(fsl.IsotropicSmooth(in_file=in_file, fwhm=4), name="smooth")

mask = Node(fsl.ApplyMask(), name="mask")

wf = Workflow(name="smoothflow")  # Workflows need names too

print(wf.connect.__doc__)

# First the "simple", but more restricted method
wf.connect(skullstrip, "mask_file", mask, "mask_file")

# Now the more complicated method. Note this way you can define several connections at once,
# and you can even define several connnections between two nodes in one smalller step
wf.connect([(smooth, mask, [("out_file", "in_file")])])

wf.write_graph("workflow_graph.dot")
from IPython.display import Image
Image(filename="workflow_graph.dot.png")

!make clean

wf.run()

wf.base_dir = "working_dir"
wf.run()

f = plt.figure(figsize=(12, 4))
for i, img in enumerate(["data/T1.nii.gz",
                         "working_dir/smoothflow/smooth/T1_smooth.nii.gz",
                         "working_dir/smoothflow/mask/T1_smooth_masked.nii.gz"]):
    f.add_subplot(1, 3, i + 1)
    plot_slice(img)

from nipype.workflows.fmri.fsl import create_susan_smooth

susan = create_susan_smooth(separate_masks=False)

susan.write_graph("susan_workflow.dot")
Image(filename="susan_workflow.dot.png")

print("Inputs:")
print(susan.inputs.inputnode)
print("Outputs:")
print(susan.outputs.outputnode)

print(susan.inputs.smooth)

from nipype import Function
extract_func = lambda l: l[0]
list_extract = Node(Function(input_names=["l"], output_names=["out"],
                             function=extract_func), name="list_extract")

wf2 = Workflow(name="susanflow", base_dir="working_dir")
skullstrip2 = Node(fsl.BET(in_file=in_file, mask=True), name="skullstrip")
mask2 = Node(fsl.ApplyMask(), name="mask")


wf2.connect([
    (skullstrip2, mask2, [("mask_file", "mask_file")]),
    (skullstrip2, susan, [("mask_file", "inputnode.mask_file")]),
    (susan, list_extract, [("outputnode.smoothed_files", "l")]),
    (list_extract, mask2, [("out", "in_file")])
    ])
susan.inputs.inputnode.in_files = abspath("data/T1.nii.gz")
susan.inputs.inputnode.fwhm = 5

wf2.write_graph("./full_susan_flow.dot")
Image(filename="full_susan_flow.dot.png")

wf2.write_graph("./full_susan_workflow_toplevel.dot", "orig")
Image(filename="full_susan_workflow_toplevel.dot.png")

wf2.run()

f = plt.figure(figsize=(12, 4))
for i, img in enumerate(["data/T1.nii.gz",
                         "working_dir/susanflow/mask/T1_smooth_masked.nii.gz"]):
    f.add_subplot(1, 2, i + 1)
    plot_slice(img)

wf.run()

wf.inputs.smooth.fwhm = 1
wf.run()

wf.run("MultiProc")

wf.run("MultiProc", dict(n_procs=2))

!make clean