include { DENOISING_MPPCA } from './modules/nf-neuro/denoising/mppca/main.nf'

include { PREPROC_SINGLEEDDY } from './modules/local/preproc/singleeddy/main.nf'

include { UTILS_EXTRACTB0 } from './modules/nf-neuro/utils/extractb0/main.nf'

include { NNUNET } from './subworkflows/local/nnunet/'

include { MOUSE_N4 } from './modules/local/mouse/n4/main.nf'

include { IMAGE_RESAMPLE as RESAMPLE_DWI} from './modules/nf-neuro/image/resample/main.nf'

include { IMAGE_RESAMPLE as RESAMPLE_MASK} from './modules/nf-neuro/image/resample/main.nf'

include { IMAGE_CONVERT } from './modules/nf-neuro/image/convert/main.nf'

include { MOUSE_REGISTRATION } from './modules/local/mouse/register/main.nf'

include { RECONST_DTIMETRICS } from './modules/nf-neuro/reconst/dtimetrics/main.nf'

include { RECONST_FRF } from './modules/nf-neuro/reconst/frf/main.nf'

include { RECONST_FODF } from './modules/nf-neuro/reconst/fodf/main.nf'

include { RECONST_QBALL } from './modules/nf-neuro/reconst/qball/main.nf'

include { TRACKING_MASK } from './modules/local/tracking/mask/main.nf'

include { TRACKING_LOCALTRACKING } from './modules/nf-neuro/tracking/localtracking/main.nf'

include { MOUSE_EXTRACTMASKS } from './modules/local/mouse/extractmasks/main.nf'

include { MOUSE_VOLUMEROISTATS } from './modules/local/mouse/volumeroistats/main.nf'

include { STATS_METRICSINROI as STATS_AMBA } from './modules/nf-neuro/stats/metricsinroi/main'

include { STATS_METRICSINROI as STATS_AMBA_LR } from './modules/nf-neuro/stats/metricsinroi/main'

include { MOUSE_CONVERTJSON as CONVERTJSON_AMBA } from './modules/local/mouse/convertjson/main.nf'

include { MOUSE_CONVERTJSON as CONVERTJSON_AMBA_LR } from './modules/local/mouse/convertjson/main.nf'

include { MOUSE_COMBINESTATS as COMBINESTATS_AMBA } from './modules/local/mouse/combinestats/main.nf'

include { MOUSE_COMBINESTATS as COMBINESTATS_AMBA_LR } from './modules/local/mouse/combinestats/main.nf'

include { MOUSE_COMBINESTATS as COMBINESTATS_MERGED} from './modules/local/mouse/combinestats/main.nf'

include { MULTIQC } from "./modules/nf-core/multiqc/main"

include { PRE_QC } from './modules/local/mouse/preqc/main.nf'

workflow get_data {

main:

if ( !params.input ) {

log.info "You must provide an input directory containing all images using:"

log.info ""

log.info " --input=/path/to/[input] Input directory containing your subjects"

log.info ""

log.info " [input]"

log.info " ├-- S1"

log.info " | ├-- *dwi.nii.gz"

log.info " | ├-- *dwi.bval"

log.info " | └-- *dwi.bvec"

log.info " └-- S2"

log.info " ├-- *dwi.nii.gz"

log.info " ├-- *dwi.bval"

log.info " └-- *dwi.bvec"

log.info ""

log.info ""

error "Please resubmit your command with the previous file structure."

}

input = file(params.input)

// ** Loading all files. ** //

dwi_channel = Channel.fromFilePairs("$input/**/*dwi.{nii.gz,bval,bvec}", size: 3, flat: true)

{ it.parent.name }

.map{ sid, bvals, bvecs, dwi -> [ [id: sid], dwi, bvals, bvecs ] } // Reordering the inputs.

mask_channel = Channel.fromPath("$input/**/*mask.nii.gz")

.map { mask_file -> def sid = mask_file.parent.name

[[id: sid], mask_file] }

template_channel = Channel.fromPath("$projectDir/assets/reference_rgb_mqc.png")

lut_channel = Channel.of([

amba : file("$projectDir/assets/LUT_AMBA.json"),

amba_lr: file("$projectDir/assets/LUT_AMBA-LR.json")

])

emit:

dwi = dwi_channel

mask = mask_channel

template_rgb = template_channel

lut = lut_channel

}

workflow {

log.warn('During the first execution, sf-mouse may take some time to download the necessary modules.')

log.info("Uses GPU: $params.use_gpu")

// Define channel for multiqc files

ch_multiqc_files = Channel.empty()

ch_multiqc_config = Channel.fromPath("$projectDir/assets/multiqc_config.yml", checkIfExists: true)

// ** Now call your input workflow to fetch your files ** //

data = get_data()

if ( params.use_fodf_for_tracking && ! params.run_tracking ) {

error "The parameter use_fodf_for_tracking cannot be enabled if run_tracking is disabled."

}

ch_dwi_bvalbvec = data.dwi

.multiMap { meta, dwi, bval, bvec ->

dwi: [ meta, dwi ]

bvs_files: [ meta, bval, bvec ]

bval: [meta, bval]

bvec: [meta, bvec]

}

ch_ref_rgb = data.template_rgb

ch_lut = data.lut

if ( params.run_preqc ) {

PRE_QC(ch_dwi_bvalbvec.dwi.join(ch_dwi_bvalbvec.bvs_files).combine(ch_ref_rgb))

ch_multiqc_files = ch_multiqc_files.mix(PRE_QC.out.rgb_mqc)

ch_multiqc_files = ch_multiqc_files.mix(PRE_QC.out.sampling_mqc)

if (params.use_preqc){

log.warn('Using the output from the preqc module is highly experimental. Please be careful.')

ch_after_preqc = PRE_QC.out.dwi

bvs_after_preqc = PRE_QC.out.bvs

}

else {

ch_after_preqc = Channel.empty()

bvs_after_preqc = Channel.empty()

}

}

else {

ch_after_preqc = ch_dwi_bvalbvec.dwi

bvs_after_preqc = ch_dwi_bvalbvec.bvs_files

}

if (params.run_denoising){

ch_mppca = ch_after_preqc

.map{ it + [[]] } // This add one empty list to the channel, since we do not have a mask.

DENOISING_MPPCA( ch_mppca )

ch_after_denoising = DENOISING_MPPCA.out.image

}

else {

ch_after_denoising = ch_after_preqc

}

ch_eddy = ch_after_denoising.join(bvs_after_preqc)

if (params.run_eddy){

PREPROC_SINGLEEDDY(ch_eddy)

ch_after_eddy = PREPROC_SINGLEEDDY.out.dwi_corrected.join(

PREPROC_SINGLEEDDY.out.bval_corrected).join(

PREPROC_SINGLEEDDY.out.bvec_corrected)

}

else {

ch_after_eddy = ch_eddy

}

UTILS_EXTRACTB0(ch_after_eddy)

ch_nnunet = ch_after_eddy.join(UTILS_EXTRACTB0.out.b0)

.join(data.mask, by: 0, remainder: true)

.map { meta, dwi, bval, bvec, b0, mask ->

[meta, dwi, bval, b0, mask ?: [ ]]} // Use empty list if mask is null

NNUNET(ch_nnunet)

if ( params.run_n4 ) {

ch_N4 = ch_after_eddy

.map{ meta, dwi, _bval, _bvec ->

tuple(meta, dwi)}

.join(UTILS_EXTRACTB0.out.b0)

.join(NNUNET.out.mask)

MOUSE_N4(ch_N4)

ch_after_n4 = MOUSE_N4.out.dwi_n4

}

else {

ch_after_n4 = ch_after_eddy

.map{ meta, dwi, _bval, _bvec -> tuple(meta, dwi)}

}

if ( params.run_resampling ) {

RESAMPLE_DWI(ch_after_n4.map{ meta, dwi -> [meta, dwi, []] }) // Add an empty list for the optional reference image

RESAMPLE_MASK(NNUNET.out.mask.map{ meta, mask -> [meta, mask, []] })

IMAGE_CONVERT(RESAMPLE_MASK.out.image)

dwi_after_resample = RESAMPLE_DWI.out.image

mask_after_resample = IMAGE_CONVERT.out.image

}

else {

dwi_after_resample = ch_after_n4

IMAGE_CONVERT(NNUNET.out.mask)

mask_after_resample = IMAGE_CONVERT.out.image

}

ch_for_mouse_registration = dwi_after_resample

.join(ch_after_eddy.map{ [it[0], it[2], it[3]] })

.join(mask_after_resample)

MOUSE_REGISTRATION(ch_for_mouse_registration)

ch_multiqc_files = ch_multiqc_files.mix(MOUSE_REGISTRATION.out.mqc)

ch_for_reconst = dwi_after_resample

.join(ch_after_eddy.map{ [it[0], it[2], it[3]] })

.join(mask_after_resample)

RECONST_DTIMETRICS(ch_for_reconst)

ch_multiqc_files = ch_multiqc_files.mix(RECONST_DTIMETRICS.out.mqc)

/* FODF */

RECONST_FRF(ch_for_reconst.map{ it + [[], [], []]})

ch_for_reconst_fodf = ch_for_reconst

.join(RECONST_DTIMETRICS.out.fa)

.join(RECONST_DTIMETRICS.out.md)

.join(RECONST_FRF.out.frf)

.map{ it + [[], []]}

RECONST_FODF(ch_for_reconst_fodf)

/* QBALL */

RECONST_QBALL(ch_for_reconst)

if ( params.use_fodf_for_tracking ) {

reconst_sh = RECONST_FODF.out.fodf

}

else {

reconst_sh = RECONST_QBALL.out.qball

}

if ( params.run_tracking ) {

TRACKING_MASK(mask_after_resample

.join(MOUSE_REGISTRATION.out.ANO))

ch_multiqc_files = ch_multiqc_files.mix(TRACKING_MASK.out.mqc)

TRACKING_LOCALTRACKING(TRACKING_MASK.out.tracking_mask

.join(reconst_sh)

.join(TRACKING_MASK.out.seeding_mask))

ch_multiqc_files = ch_multiqc_files.mix(TRACKING_LOCALTRACKING.out.mqc)

}

MOUSE_EXTRACTMASKS(MOUSE_REGISTRATION.out.ANO)

ch_metrics = RECONST_DTIMETRICS.out.md

.join(RECONST_DTIMETRICS.out.fa)

.join(RECONST_DTIMETRICS.out.rd)

.join(RECONST_DTIMETRICS.out.ad)

.map{ meta, fa, md, ad, rd ->

tuple(meta, [ fa, md, ad, rd ])}

ch_for_stats = ch_metrics

.combine(MOUSE_EXTRACTMASKS.out.masks_dir, by: 0)

MOUSE_VOLUMEROISTATS(ch_for_stats)

STATS_AMBA(ch_metrics.join(MOUSE_REGISTRATION.out.ANO.combine(ch_lut.map{ it.amba })))

STATS_AMBA_LR(ch_metrics.join(MOUSE_REGISTRATION.out.ANO_LR.combine(ch_lut.map{ it.amba_lr })))

CONVERTJSON_AMBA(STATS_AMBA.out.stats_json)

CONVERTJSON_AMBA_LR(STATS_AMBA_LR.out.stats_json)

all_stats_amba = CONVERTJSON_AMBA.out.stats_reorganized

.map{ _meta, json -> json}

.collect()

all_stats_lr = CONVERTJSON_AMBA_LR.out.stats_reorganized

.map{ _meta, json -> json}

.collect()

all_stats_merged = MOUSE_VOLUMEROISTATS.out.stats

.map{ _meta, json -> json}

.collect()

COMBINESTATS_AMBA(all_stats_amba)

COMBINESTATS_AMBA_LR(all_stats_lr)

COMBINESTATS_MERGED(all_stats_merged)

ch_multiqc_files = ch_multiqc_files

.groupTuple()

.map { meta, files_list ->

def files = files_list.flatten().findAll { it != null }

return tuple(meta, files)

}

MULTIQC(ch_multiqc_files, [], ch_multiqc_config.toList(), [], channel.fromPath("${projectDir}/assets/sf-mouse-light-logo.png").toList(), [], [])

}