#--- snakebids boilerplate

import snakebids
from snakebids import bids, filter_list

configfile: 'config/snakebids.yml'

#writes inputs_config.yml and updates config dict
snakebids.generate_inputs_config(config)

#this adds constraints to the bids naming
wildcard_constraints:  **snakebids.get_wildcard_constraints(config)

 
#--- rest of workflow below

#additional constraints for wildcards not defined from inputs
wildcard_constraints:
    desc='[a-zA-Z0-9]+',
    fwhm='[0-9]+',
    confounds='[0-9]+'


rule all:
    input: 
        # using the zip lists to expand over all scans, note use of the zip option in expand:
        denoised = expand(
                        expand(
                            bids(root='results',datatype='func',desc='{{desc}}',fwhm='{{fwhm}}',confounds='{{confounds_idx}}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
                             zip, **config['input_zip_lists']['preproc_bold']),
                        fwhm=config['fwhm'],confounds_idx=range(1,len(config['confounds'])+1),desc=['denoised','AROMAdenoised'])

  
rule smooth:
    input:
        nii = config['input_path']['preproc_bold'],
        json = re.sub('.nii.gz','.json',config['input_path']['preproc_bold'])
    params:
        fwhm = lambda wildcards: float(wildcards.fwhm)
    output:
        nii = bids(root='results',datatype='func',desc='smoothed',fwhm='{fwhm}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
        json = bids(root='results',datatype='func',desc='smoothed',fwhm='{fwhm}',suffix='bold.json',**config['input_wildcards']['preproc_bold'])
    group: 'subj'
    script: 'scripts/smooth.py'




rule denoise:
    input: 
        nii = bids(root='results',datatype='func',desc='smoothed',fwhm='{fwhm}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
        json = bids(root='results',datatype='func',desc='smoothed',fwhm='{fwhm}',suffix='bold.json',**config['input_wildcards']['preproc_bold']),
        confounds_tsv = config['input_path']['confounds'],
        mask_nii = config['input_path']['preproc_mask']
    params:
        confounds_to_use = lambda wildcards: config['confounds'][int(wildcards.confounds_idx)-1]['regressors'],
        confounds_name = lambda wildcards: config['confounds'][int(wildcards.confounds_idx)-1]['name'],
        standardize = True,
        detrend = True,
        low_pass = False,
        high_pass = False,
    output: 
        nii = bids(root='results',datatype='func',desc='denoised',fwhm='{fwhm}',confounds='{confounds_idx}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
        json= bids(root='results',datatype='func',desc='denoised',fwhm='{fwhm}',confounds='{confounds_idx}',suffix='bold.json',**config['input_wildcards']['preproc_bold'])
    group: 'subj'
    script: 'scripts/denoise.py'

                

rule aroma_nonaggr:
    input: 
        nii = bids(root='results',datatype='func',desc='smoothed',fwhm='{fwhm}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
        mixing = config['input_path']['mixing'],
        noiseICs = config['input_path']['noiseICs'],
        mask_nii = config['input_path']['preproc_mask']
    params:
        container = config['singularity']['fsl']
    output: 
        nii = bids(root='results',datatype='func',desc='AROMAnonaggr',fwhm='{fwhm}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
    container: config['singularity']['fsl']
    group: 'subj'
    shell:
        'fsl_regfilt -i {input.nii} -f `cat {input.noiseICs}` -d {input.mixing} -o {output.nii} -m {input.mask_nii}' 

#this regresses out confounds after aroma 
rule aroma_aggr:
    input: 
        nii = bids(root='results',datatype='func',desc='AROMAnonaggr',fwhm='{fwhm}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
        json = bids(root='results',datatype='func',desc='smoothed',fwhm='{fwhm}',suffix='bold.json',**config['input_wildcards']['preproc_bold']),
        confounds_tsv = config['input_path']['confounds'],
        mask_nii = config['input_path']['preproc_mask']
    params:
        confounds_to_use = lambda wildcards: config['confounds'][int(wildcards.confounds_idx)-1]['regressors'],
        confounds_name = lambda wildcards: config['confounds'][int(wildcards.confounds_idx)-1]['name'],
        standardize = True,
        detrend = True,
        low_pass = False,
        high_pass = False,
    output: 
        nii = bids(root='results',datatype='func',desc='AROMAdenoised',fwhm='{fwhm}',confounds='{confounds_idx}',suffix='bold.nii.gz',**config['input_wildcards']['preproc_bold']),
        json= bids(root='results',datatype='func',desc='AROMAdenoised',fwhm='{fwhm}',confounds='{confounds_idx}',suffix='bold.json',**config['input_wildcards']['preproc_bold'])
    group: 'subj'
    script: 'scripts/denoise.py'

    


