configfile: '../../config/snakebids.yml'

wildcard_constraints:
    i = '[0-9]+'



rule all:
    input: expand('tpl-upenn_dir-{dir}_level-0_coords.nii.gz', dir=['AP','PD','IO'])

rule create_PD_sink:
    """ this rule is still a WIP and shouldn't be used.. """
    input:
        lbl = 'tpl-upenn_desc-hipptissue_dseg.nii.gz',
    output:
        lbl = 'tpl-upenn_desc-withPDsink_level-0_dseg.nii.gz',
    container: config['singularity']['autotop']
    shell:
        # dilate DG into the SRLM (lbl 2/4/7) and background (lbl 0) to get a PD sink that is outside the laplace domain
        #  create a new label for each region dilated into, so the AP/IO src/sink are unaffected
        #    we can do this easily by just adding 10 to the existing labels (0 -> 10, 2 -> 12, 4->14, 7->17
        'c3d {input} -retain-labels 8 -dilate 8 1x1x1vox -binarize -popas DILATED_DG ' 
        ' {input} -replace 0 10 -retain-labels 2 4 7 10 -binarize -push DILATED_DG -multiply -as DG_SINK '
        ' -scale 10 {input} -add -o {output}' # add 10 to original labels, and save it

rule downsample_coords:
    input: 
        lbl = lambda wildcards: 'tpl-upenn_desc-withPDsink_level-{i}_dseg.nii.gz'.format(i=int(wildcards.i)-1)
    output:
        lbl = 'tpl-upenn_desc-withPDsink_level-{i}_dseg.nii.gz',
    shell:
        'c3d {input} -interpolation NearestNeighbor -resample 50%  {output}'


# run laplace coords without initialization only for lowest resolution
#  this is done by setting constraint on output i  (i.e. {i,2})
rule laplace_coords:
    input:
        lbl = 'tpl-upenn_desc-withPDsink_level-{i}_dseg.nii.gz',
    params:
        gm_labels = lambda wildcards: config['laplace_labels'][wildcards.dir]['gm'],
        src_labels = lambda wildcards: config['laplace_labels'][wildcards.dir]['src'],
        sink_labels = lambda wildcards: config['laplace_labels'][wildcards.dir]['sink'],
        convergence_threshold = 1e-6,
        max_iters = 100000
    output:
        coords = 'tpl-upenn_dir-{dir}_level-{i,2}_coords.nii.gz'
    resources:
        time = 1200
    log: 'log_dir-{dir}_level-{i}_laplace.txt'
    script: '../../workflow/scripts/laplace_coords.py'


rule laplace_init_lores:
    input:
        lbl = 'tpl-upenn_desc-withPDsink_level-{i}_dseg.nii.gz',
        init_coords = 'tpl-upenn_dir-{dir}_level-{i}_desc-upsampled_coords.nii.gz'
    params:
        gm_labels = lambda wildcards: config['laplace_labels'][wildcards.dir]['gm'],
        src_labels = lambda wildcards: config['laplace_labels'][wildcards.dir]['src'],
        sink_labels = lambda wildcards: config['laplace_labels'][wildcards.dir]['sink'],
        convergence_threshold = 1e-8,
        max_iters = 100000
    output:
        coords = 'tpl-upenn_dir-{dir}_level-{i,0|1}_coords.nii.gz'
    resources:
        time = 1200
    log: 'log_dir-{dir}_level-{i}_laplace.txt'
    script: '../../workflow/scripts/laplace_coords_withinit.py'


rule upsample_coords:
    input:
        coords = lambda wildcards: 'tpl-upenn_dir-{dir}_level-{i}_coords.nii.gz'.format(dir=wildcards.dir, i=int(wildcards.i)+1)
    output:
        coords = 'tpl-upenn_dir-{dir}_level-{i}_desc-upsampled_coords.nii.gz'
    shell:
        'c3d {input} -interpolation NearestNeighbor -resample 200%  {output}'
        

