from file_tree import FileTree
from fsl_pipe import Pipeline, In, Out, Ref

# Load some libraries that allow us to run FSL tools
from fsl import wrappers
from subprocess import run
from os import getenv

# Load the file-tree describing the data directory
tree = FileTree.read("data.tree").update_glob("T1w", link=[("group", "subject")])

# Create the pipeline
pipe = Pipeline()


# Add recipes to the pipeline
# Filled by user
@pipe
def brain_extract(T1w: In, T1w_brain_first_attempt: Out, T1w_cut: Out, T1w_brain: Out):
    wrappers.bet(T1w, T1w_brain_first_attempt)
    run(["standard_space_roi", T1w_brain_first_attempt, T1w_cut, "-roiNONE", "-altinput", T1w, "-ssref", f"{getenv('FSLDIR')}/data/standard/MNI152_T1_2mm_brain"])
    wrappers.bet(T1w_cut, T1w_brain, fracintensity=0.4)

@pipe
def segmentation(T1w_brain: In, basename: Ref("segment/basename"), gm_pve: Out("segment/gm_pve")):
    wrappers.fast(T1w_brain, out=basename)

@pipe
def linear_registration(gm_pve: In("segment/gm_pve"), linear_reg: Out):
    run(["flirt", "-in", gm_pve, "-ref", f"{getenv('FSLDIR')}/data/standard/tissuepriors/avg152T1_gray", "-omat", linear_reg])

@pipe
def nonlinear_registration(gm_pve: In("segment/gm_pve"), linear_reg: In, nonlinear_reg: Out, gm_pve_in_standard: Out):
    run(["fnirt", f"--in={gm_pve}", f"--ref={getenv('FSLDIR')}/data/standard/tissuepriors/avg152T1_gray", f"--aff={linear_reg}", f"--cout={nonlinear_reg}", "--config=bad-fnirt-config.cnf", f"--iout={gm_pve_in_standard}"])


# Run the pipeline command line interface
if __name__ == "__main__":
    pipe.cli(tree)