# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
#
# Copyright 2021 The NiPreps Developers <nipreps@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# We support and encourage derived works from this project, please read
# about our expectations at
#
#     https://www.nipreps.org/community/licensing/
#
"""
Functional workflow
===================

.. image :: _static/functional_workflow_source.svg

The functional workflow follows the following steps:

#. Sanitize (revise data types and xforms) input data, read
   associated metadata and discard non-steady state frames.
#. :abbr:`HMC (head-motion correction)` based on ``3dvolreg`` from
   AFNI -- :py:func:`hmc`.
#. Skull-stripping of the time-series (AFNI) --
   :py:func:`fmri_bmsk_workflow`.
#. Calculate mean time-series, and :abbr:`tSNR (temporal SNR)`.
#. Spatial Normalization to MNI (ANTs) -- :py:func:`epi_mni_align`
#. Extraction of IQMs -- :py:func:`compute_iqms`.
#. Individual-reports generation -- :py:func:`individual_reports`.

This workflow is orchestrated by :py:func:`fmri_qc_workflow`.
"""
from mriqc import config
from nipype.interfaces import io as nio
from nipype.interfaces import utility as niu
from nipype.pipeline import engine as pe


def fmri_qc_workflow(name="funcMRIQC"):
    """
    Initialize the (f)MRIQC workflow.

    .. workflow::

        import os.path as op
        from mriqc.workflows.functional import fmri_qc_workflow
        from mriqc.testing import mock_config
        with mock_config():
            wf = fmri_qc_workflow()

    """
    from nipype.algorithms.confounds import TSNR, NonSteadyStateDetector
    from nipype.interfaces.afni import TStat
    from niworkflows.interfaces.header import SanitizeImage

    workflow = pe.Workflow(name=name)

    mem_gb = config.workflow.biggest_file_gb

    dataset = config.workflow.inputs.get("bold", [])
    config.loggers.workflow.info(
        f"""\
Building functional MRIQC workflow for files: {", ".join(dataset)}."""
    )

    # Define workflow, inputs and outputs
    # 0. Get data, put it in RAS orientation
    inputnode = pe.Node(niu.IdentityInterface(fields=["in_file"]), name="inputnode")
    inputnode.iterables = [("in_file", dataset)]

    outputnode = pe.Node(
        niu.IdentityInterface(
            fields=["qc", "mosaic", "out_group", "out_dvars", "out_fd"]
        ),
        name="outputnode",
    )

    non_steady_state_detector = pe.Node(
        NonSteadyStateDetector(), name="non_steady_state_detector"
    )

    sanitize = pe.Node(SanitizeImage(), name="sanitize", mem_gb=mem_gb * 4.0)
    sanitize.inputs.max_32bit = config.execution.float32

    # Workflow --------------------------------------------------------

    # 1. HMC: head motion correct
    hmcwf = hmc()

    # Set HMC settings
    hmcwf.inputs.inputnode.fd_radius = config.workflow.fd_radius

    # 2. Compute mean fmri
    mean = pe.Node(
        TStat(options="-mean", outputtype="NIFTI_GZ"),
        name="mean",
        mem_gb=mem_gb * 1.5,
    )
    skullstrip_epi = fmri_bmsk_workflow()

    # EPI to MNI registration
    ema = epi_mni_align()

    # Compute TSNR using nipype implementation
    tsnr = pe.Node(TSNR(), name="compute_tsnr", mem_gb=mem_gb * 2.5)

    # 7. Compute IQMs
    iqmswf = compute_iqms()
    # Reports
    repwf = individual_reports()

    # fmt: off

    workflow.connect([
        (inputnode, iqmswf, [("in_file", "inputnode.in_file")]),
        (inputnode, sanitize, [("in_file", "in_file")]),
        (inputnode, non_steady_state_detector, [("in_file", "in_file")]),
        (non_steady_state_detector, sanitize, [("n_volumes_to_discard", "n_volumes_to_discard")]),
        (sanitize, hmcwf, [("out_file", "inputnode.in_file")]),
        (mean, skullstrip_epi, [("out_file", "inputnode.in_file")]),
        (hmcwf, mean, [("outputnode.out_file", "in_file")]),
        (hmcwf, tsnr, [("outputnode.out_file", "in_file")]),
        (mean, ema, [("out_file", "inputnode.epi_mean")]),
        (skullstrip_epi, ema, [("outputnode.out_file", "inputnode.epi_mask")]),
        (sanitize, iqmswf, [("out_file", "inputnode.in_ras")]),
        (mean, iqmswf, [("out_file", "inputnode.epi_mean")]),
        (hmcwf, iqmswf, [("outputnode.out_file", "inputnode.hmc_epi"),
                         ("outputnode.out_fd", "inputnode.hmc_fd")]),
        (skullstrip_epi, iqmswf, [("outputnode.out_file", "inputnode.brainmask")]),
        (tsnr, iqmswf, [("tsnr_file", "inputnode.in_tsnr")]),
        (sanitize, repwf, [("out_file", "inputnode.in_ras")]),
        (mean, repwf, [("out_file", "inputnode.epi_mean")]),
        (tsnr, repwf, [("stddev_file", "inputnode.in_stddev")]),
        (skullstrip_epi, repwf, [("outputnode.out_file", "inputnode.brainmask")]),
        (hmcwf, repwf, [("outputnode.out_fd", "inputnode.hmc_fd"),
                        ("outputnode.out_file", "inputnode.hmc_epi")]),
        (ema, repwf, [("outputnode.epi_parc", "inputnode.epi_parc"),
                      ("outputnode.report", "inputnode.mni_report")]),
        (non_steady_state_detector, iqmswf, [("n_volumes_to_discard", "inputnode.exclude_index")]),
        (iqmswf, repwf, [("outputnode.out_file", "inputnode.in_iqms"),
                         ("outputnode.out_dvars", "inputnode.in_dvars"),
                         ("outputnode.outliers", "inputnode.outliers")]),
        (hmcwf, outputnode, [("outputnode.out_fd", "out_fd")]),
    ])
    # fmt: on

    if config.workflow.fft_spikes_detector:
        # fmt: off
        workflow.connect([
            (iqmswf, repwf, [("outputnode.out_spikes", "inputnode.in_spikes"),
                             ("outputnode.out_fft", "inputnode.in_fft")]),
        ])
        # fmt: on

    if config.workflow.ica:
        from niworkflows.interfaces.reportlets.segmentation import MELODICRPT

        melodic = pe.Node(
            MELODICRPT(
                no_bet=True,
                no_mask=True,
                no_mm=True,
                compress_report=False,
                generate_report=True,
            ),
            name="ICA",
            mem_gb=max(mem_gb * 5, 8),
        )
        # fmt: off
        workflow.connect([
            (sanitize, melodic, [("out_file", "in_files")]),
            (skullstrip_epi, melodic, [("outputnode.out_file", "report_mask")]),
            (melodic, repwf, [("out_report", "inputnode.ica_report")])
        ])
        # fmt: on

    # Upload metrics
    if not config.execution.no_sub:
        from mriqc.interfaces.webapi import UploadIQMs

        upldwf = pe.Node(UploadIQMs(), name="UploadMetrics")
        upldwf.inputs.url = config.execution.webapi_url
        upldwf.inputs.strict = config.execution.upload_strict
        if config.execution.webapi_port:
            upldwf.inputs.port = config.execution.webapi_port

        # fmt: off
        workflow.connect([
            (iqmswf, upldwf, [("outputnode.out_file", "in_iqms")]),
        ])
        # fmt: on

    return workflow


def compute_iqms(name="ComputeIQMs"):
    """
    Initialize the workflow that actually computes the IQMs.

    .. workflow::

        from mriqc.workflows.functional import compute_iqms
        from mriqc.testing import mock_config
        with mock_config():
            wf = compute_iqms()

    """
    from nipype.algorithms.confounds import ComputeDVARS
    from nipype.interfaces.afni import OutlierCount, QualityIndex
    from niworkflows.interfaces.bids import ReadSidecarJSON

    from mriqc.interfaces import FunctionalQC, IQMFileSink
    from mriqc.interfaces.reports import AddProvenance
    from mriqc.interfaces.transitional import GCOR
    from mriqc.workflows.utils import _tofloat, get_fwhmx

    mem_gb = config.workflow.biggest_file_gb

    workflow = pe.Workflow(name=name)
    inputnode = pe.Node(
        niu.IdentityInterface(
            fields=[
                "in_file",
                "in_ras",
                "epi_mean",
                "brainmask",
                "hmc_epi",
                "hmc_fd",
                "fd_thres",
                "in_tsnr",
                "metadata",
                "exclude_index",
            ]
        ),
        name="inputnode",
    )
    outputnode = pe.Node(
        niu.IdentityInterface(
            fields=[
                "out_file",
                "out_dvars",
                "outliers",
                "out_spikes",
                "out_fft",
            ]
        ),
        name="outputnode",
    )

    # Set FD threshold
    inputnode.inputs.fd_thres = config.workflow.fd_thres

    # Compute DVARS
    dvnode = pe.Node(
        ComputeDVARS(save_plot=False, save_all=True),
        name="ComputeDVARS",
        mem_gb=mem_gb * 3,
    )

    # AFNI quality measures
    fwhm_interface = get_fwhmx()
    fwhm = pe.Node(fwhm_interface, name="smoothness")
    # fwhm.inputs.acf = True  # add when AFNI >= 16
    outliers = pe.Node(
        OutlierCount(fraction=True, out_file="outliers.out"),
        name="outliers",
        mem_gb=mem_gb * 2.5,
    )

    quality = pe.Node(
        QualityIndex(automask=True),
        out_file="quality.out",
        name="quality",
        mem_gb=mem_gb * 3,
    )

    gcor = pe.Node(GCOR(), name="gcor", mem_gb=mem_gb * 2)

    measures = pe.Node(FunctionalQC(), name="measures", mem_gb=mem_gb * 3)

    # fmt: off
    workflow.connect([
        (inputnode, dvnode, [("hmc_epi", "in_file"),
                             ("brainmask", "in_mask")]),
        (inputnode, measures, [("epi_mean", "in_epi"),
                               ("brainmask", "in_mask"),
                               ("hmc_epi", "in_hmc"),
                               ("hmc_fd", "in_fd"),
                               ("fd_thres", "fd_thres"),
                               ("in_tsnr", "in_tsnr")]),