{#
Macros for common ANTs tasks:
* Registration, both rigid and SyN
* Application of transformations
#}

{# Return a step calling antsRegistrationSyN.sh #}
{%- macro register_outputs(fixed, moving, prefix, suffixes, transform_type) -%}
{%- for suffix in suffixes %}
- {{ prefix }}{{ suffix }} 
{%- endfor %}
{%- endmacro -%}

{%- macro rigid_registration_stage(fixed, moving) -%}
--initial-moving-transform '[{{ fixed }},{{ moving }},1]' {#
#}--transform Rigid[0.1] {#
#}--metric 'MI[{{ fixed }},{{ moving }},1,32,Regular,0.25]' {#
#}--convergence [1000x500x250x100,1e-6,10] {#
#}--shrink-factors 8x4x2x1 {#
#}--smoothing-sigmas 3x2x1x0vox
{%- endmacro -%}

{%- macro affine_registration_stage(fixed, moving) -%}
--transform Affine[0.1] {#
#}--metric 'MI[{{ fixed }},{{ moving }},1,32,Regular,0.25]' {#
#}--convergence [1000x500x250x100,1e-6,10] {#
#}--shrink-factors 8x4x2x1 {#
#}--smoothing-sigmas 3x2x1x0vox
{%- endmacro -%}

{%- macro syn_registration_stage(fixed, moving) -%}
--transform SyN[0.1,3,0] {#
#}--metric 'CC[{{ fixed }},{{ moving }},1,4]' {#
#}--convergence [100x70x50x20,1e-6,10] {#
#}--shrink-factors 8x4x2x1 {#
#}--smoothing-sigmas 3x2x1x0vox
{%- endmacro -%}

{# FIXME: slots, use ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS #}
{%- macro register_commands(fixed, moving, prefix, suffixes, transform_type) -%}
- antsRegistration {#
    #}--dimensionality 3 --float 0 {#
    #}--interpolation Linear --winsorize-image-intensities [0.005,0.995] {#
    #}--use-histogram-matching 0 {#
    #}--output '[{{ prefix }},{{ prefix }}Warped.nii.gz]' {#
    #}{% 
      if transform_type in ["r","s"] 
        %}{{ rigid_registration_stage(fixed, moving) }} {%
      endif %}{%
      if transform_type == "s"
        %}{{ affine_registration_stage(fixed, moving) }} {{ syn_registration_stage(fixed, moving) }} {%
      endif %}
{%- endmacro -%}

{%- macro register(fixed, moving, prefix, suffixes, transform_type) -%}
outputs:
  {{ register_outputs(fixed, moving, prefix, suffixes, transform_type)|indent(2) }}
inputs:
  - {{ fixed }} 
  - {{ moving }}
commands:
  {{ register_commands(fixed, moving, prefix, suffixes, transform_type) }}
{%- endmacro -%}

{# Return a step performing a rigid registration using antsRegistrationSyN.sh #}
{%- macro rigid_register_outputs(fixed, moving, prefix) -%}
{{ register_outputs(
	fixed, moving, prefix, ["0GenericAffine.mat", "Warped.nii.gz"], "r") }}
{%- endmacro -%}

{%- macro rigid_register_commands(fixed, moving, prefix) -%}
{{ register_commands(
	fixed, moving, prefix, ["0GenericAffine.mat", "Warped.nii.gz"], "r") }}
{%- endmacro -%}

{%- macro rigid_register(fixed, moving, prefix) -%}
{{ register(
	fixed, moving, prefix, ["0GenericAffine.mat", "Warped.nii.gz"], "r") }}
{%- endmacro -%}

{# Return a step performing a SyN registration using antsRegistrationSyN.sh #}
{%- macro syn_register_outputs(fixed, moving, prefix) -%}
{{ register_outputs(
	fixed, moving, prefix, [
    "0GenericAffine.mat", "1Warp.nii.gz", "1InverseWarp.nii.gz", "Warped.nii.gz"
  ], "s") }}
{%- endmacro -%}

{%- macro syn_register_commands(fixed, moving, prefix) -%}
{{ register_commands(
	fixed, moving, prefix, [
    "0GenericAffine.mat", "1Warp.nii.gz", "1InverseWarp.nii.gz", "Warped.nii.gz"
  ], "s") }}
{%- endmacro -%}


{%- macro syn_register(fixed, moving, prefix) -%}
{{ register(
  fixed, moving, prefix, [
    "0GenericAffine.mat", "1Warp.nii.gz", "1InverseWarp.nii.gz", "Warped.nii.gz"
  ], "s") }}
{%- endmacro -%}

{# 
Return the transform files used by a SyN tranformation, in reverse order of
application, separated by an optional separator.

For example:

antsApplyTransforms -i source.nii -r ref.nii -o dest.nii {{ syn_transforms("in_source", "-t ") }}
#}
{%- macro syn_transforms(root, separator="") -%}
  {% for suffix in ["1Warp.nii.gz", "0GenericAffine.mat"] -%}
    {{ separator }}{{ root }}{{ suffix }}{{ " " if not loop.last }}
  {%- endfor %}
{%- endmacro -%}

{# Return a step appling transforms #}
{% macro apply_transforms_commands(
    source, reference, destination, transforms, 
    interpolation="BSpline", input_image_type="scalar") -%}
- antsApplyTransforms {#
  #}-i '{{ source }}' -r '{{ reference }}' -o '{{ destination }}' {#
  #}-n {{ interpolation }} -e {{ input_image_type }} {#
  #}{% for transform in transforms 
    %}-t {% 
      if transform is string 
        %}'{{ transform }}'{% 
      else 
        %}'[{{ transform|join(",") }}]'{% 
      endif %} {% endfor %}
{%- endmacro -%}

{% macro apply_transforms(
    source, reference, destination, transforms, 
    interpolation="BSpline", input_image_type="scalar") -%}
outputs: {{ destination }}
inputs:
  - {{ source|json }}
  {% for transform in transforms %}
    {% if transform is string %}
  - {{ transform|json }}
    {% else %}
  - {{ transform|first|json }}
    {% endif %}
  {% endfor %} 
commands:
  {{ apply_transforms_commands(
        source, reference, destination, transforms, 
        interpolation, input_image_type)|indent(2) }}
{%- endmacro -%}

{# Return a step appling a SyN transform #}
{% macro apply_syn_transforms(source, reference, destination, prefix) -%}
{{ apply_transforms(
    source, reference, destination, 
    [prefix+"1Warp.nii.gz", prefix+"0GenericAffine.mat"])}}
{%- endmacro -%}
