# -*- DISCLAIMER: this file contains code derived from Nipype (https://github.com/nipy/nipype/blob/master/LICENSE)  -*-

from nipype.interfaces.base import (
    traits,
    BaseInterface,
    BaseInterfaceInputSpec,
    TraitedSpec,
    File,
    Bunch,
    isdefined,
)
from swane.config.config_enums import BLOCK_DESIGN


# -*- DISCLAIMER: this class extends a Nipype class (nipype.interfaces.base.BaseInterfaceInputSpec)  -*-
class FMRIGenSpecInputSpec(BaseInterfaceInputSpec):
    TR = traits.Float(mandatory=True, desc="Repetition Time")
    nvols = traits.Int(mandatory=True, desc="Number of EPI runs")
    task_duration = traits.Int(mandatory=True, desc="Task duration")
    rest_duration = traits.Int(mandatory=True, desc="Rest duration")
    block_design = traits.Enum(BLOCK_DESIGN, usedefault=True)
    task_a_name = traits.String(mandatory=False, desc="Task A name")
    task_b_name = traits.String(mandatory=False, desc="Task A name")
    hpcutoff = traits.Int(mandatory=False, desc="Cutoff for highpass filtering")
    tempMean = File(exists=True, mandatory=True, desc="mean functional image")


# -*- DISCLAIMER: this class extends a Nipype class (nipype.interfaces.base.TraitedSpec)  -*-
class FMRIGenSpecOutputSpec(TraitedSpec):
    hpcutoff = traits.Int(desc="Cutoff for highpass filtering")
    hp_sigma_vol = traits.Float(desc="Sigma volume for highpass filtering")
    evs_run = traits.Any(desc="task events")
    hpstring = traits.String(desc="op_string for highpass filtering")
    task_a_name = traits.String(desc="Task name")
    task_b_name = traits.String(desc="Task name")
    contrasts = traits.List(desc="T contrast array")


# -*- DISCLAIMER: this class extends a Nipype class (nipype.interfaces.base.BaseInterface)  -*-
class FMRIGenSpec(BaseInterface):
    """
    Formats the fMRI parameters for FSL feat nodes.

    """

    input_spec = FMRIGenSpecInputSpec
    output_spec = FMRIGenSpecOutputSpec

    def _run_interface(self, runtime):

        if isdefined(self.inputs.hpcutoff):
            self.hpcutoff = self.inputs.hpcutoff
        elif self.inputs.block_design == BLOCK_DESIGN.RARA:
            self.hpcutoff = self.inputs.task_duration + self.inputs.rest_duration
        else:
            self.hpcutoff = (self.inputs.task_duration + self.inputs.rest_duration) * 2

        self.hp_sigma_vol = self.hpcutoff / (2 * self.inputs.TR)

        self.hpstring = "-bptf %f -1 -add %s" % (
            self.hp_sigma_vol,
            self.inputs.tempMean,
        )

        if not isdefined(self.inputs.task_a_name):
            self.inputs.task_a_name = "TaskA"
        if not isdefined(self.inputs.task_b_name):
            self.inputs.task_b_name = "TaskB"

        if self.inputs.block_design == BLOCK_DESIGN.RARA:
            self.contrasts = [
                [
                    "%s_versus_Rest" % self.inputs.task_a_name,
                    "T",
                    [self.inputs.task_a_name],
                    [1],
                ]
            ]

            self.evs_run = Bunch(
                conditions=[self.inputs.task_a_name],
                onsets=[
                    list(
                        range(
                            self.inputs.rest_duration,
                            int(self.inputs.TR * self.inputs.nvols),
                            (self.inputs.task_duration + self.inputs.rest_duration),
                        )
                    )
                ],
                durations=[[self.inputs.task_duration]],
            )
        else:
            cont1 = [
                "%s_versus_%s" % (self.inputs.task_a_name, self.inputs.task_b_name),
                "T",
                [self.inputs.task_a_name, self.inputs.task_b_name],
                [1, -1],
            ]
            cont2 = [
                "%s_versus_%s" % (self.inputs.task_b_name, self.inputs.task_a_name),
                "T",
                [self.inputs.task_a_name, self.inputs.task_b_name],
                [-1, 1],
            ]
            self.contrasts = [cont1, cont2]

            self.evs_run = Bunch(
                conditions=[self.inputs.task_a_name, self.inputs.task_b_name],
                onsets=[
                    list(
                        range(
                            self.inputs.rest_duration,
                            int(self.inputs.TR * self.inputs.nvols),
                            (
                                self.inputs.task_duration * 2
                                + self.inputs.rest_duration * 2
                            ),
                        )
                    ),
                    list(
                        range(
                            (self.inputs.rest_duration * 2 + self.inputs.task_duration),
                            int(self.inputs.TR * self.inputs.nvols),
                            (
                                self.inputs.task_duration * 2
                                + self.inputs.rest_duration * 2
                            ),
                        )
                    ),
                ],
                durations=[[self.inputs.task_duration], [self.inputs.task_duration]],
            )

        return runtime

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["hpcutoff"] = self.hpcutoff
        outputs["hp_sigma_vol"] = self.hp_sigma_vol
        outputs["hpstring"] = self.hpstring
        outputs["evs_run"] = self.evs_run
        outputs["contrasts"] = self.contrasts
        outputs["task_a_name"] = self.inputs.task_a_name
        outputs["task_b_name"] = self.inputs.task_b_name
        return outputs
