from pathlib import Path
from typing import List

from nipype import config

from clinica.pipelines.dwi.preprocessing.engine import DWIPreprocessingPipeline

cfg = dict(execution={"parameterize_dirs": False})
config.update_config(cfg)


class DwiPreprocessingUsingT1(DWIPreprocessingPipeline):
    """DWI Preprocessing using T1 image for susceptibility distortion step.

    Ideas for improvement:
        - Replace prepare_reference_b0 function by a first run of FSL eddy
        - Replace B0-T1w registration by FA-T1w registration
        - Replace os.system(cmd) by call of Nipype interface
        - Use promising sdcflows workflows and/or dMRIprep

    Warnings:
        - Do not use this pipeline if you have fieldmap data in your dataset.

    Note:
        Some reading regarding the reproducibility of FSL eddy command:
        https://www.jiscmail.ac.uk/cgi-bin/webadmin?A2=fsl;1ccf038f.1608

    Returns:
        A clinica pipeline object containing the DWIPreprocessingUsingT1 pipeline.
    """

    @staticmethod
    def get_processed_images(
        caps_directory: Path, subjects: List[str], sessions: List[str]
    ) -> List[str]:
        from clinica.utils.filemanip import extract_image_ids
        from clinica.utils.input_files import DWI_PREPROC_NII
        from clinica.utils.inputs import clinica_file_reader

        image_ids: List[str] = []
        if caps_directory.is_dir():
            preproc_files, _ = clinica_file_reader(
                subjects, sessions, caps_directory, DWI_PREPROC_NII, False
            )
            image_ids = extract_image_ids(preproc_files)
        return image_ids

    def _check_pipeline_parameters(self) -> None:
        """Check pipeline parameters."""
        super()._check_pipeline_parameters()
        self.parameters.setdefault("delete_cache", False)
        self.parameters.setdefault("random_seed", None)
        self.parameters.setdefault("double_precision", True)

    def get_input_fields(self) -> List[str]:
        """Specify the list of possible inputs of this pipeline.

        Returns
        -------
        List[str] :
            The list of inputs for the DwiPreprocessingUsingT1 pipeline namely:
                - "t1w_filename" : Path of the T1w image in BIDS format.
                - "dwi_filename": Path of the diffusion weighted image in BIDS format.
                - "dwi_json_filename" : Path of the DWI JSON file in BIDS format and containing
                  'TotalReadoutTime' and 'PhaseEncodingDirection' metadata (see BIDS specifications).
                - "b_vectors_filename" : Path of the b-vectors file in BIDS format.
                - "b_values_filename" : Path of the b-values file in BIDS format.
        """
        return [
            "t1w_filename",
            "dwi_filename",
            "dwi_json_filename",
            "b_vectors_filename",
            "b_values_filename",
        ]

    def get_output_fields(self) -> List[str]:
        """Specify the list of possible outputs of this pipeline.

        Returns:
            List[str]: The list of outputs generated by the DwiPreprocessingUsingT1 pipeline namely:
                * preproc_dwi: Path of the preprocessed DWI
                * preproc_bvec: Path of the preprocessed bvec
                * preproc_bval: Path of the preprocessed bval
                * b0_mask: Path of the b0 brainmask
        """
        return ["preproc_dwi", "preproc_bvec", "preproc_bval", "b0_mask"]

    def filter_qc(self) -> tuple[list[str], list[str]]:
        from clinica.pipelines.dwi.preprocessing.utils import check_dwi_volume
        from clinica.pipelines.dwi.utils import DWIDataset
        from clinica.utils.bids import BIDSFileName
        from clinica.utils.input_files import (
            DWI_BVAL,
            DWI_BVEC,
            DWI_NII,
        )
        from clinica.utils.inputs import clinica_list_of_files_reader
        from clinica.utils.stream import cprint

        subjects = []
        sessions = []
        list_bids_files = clinica_list_of_files_reader(
            self.subjects,
            self.sessions,
            self.bids_directory,
            [DWI_NII, DWI_BVEC, DWI_BVAL],
            raise_exception=True,
        )
        for dwi_image_file, b_vectors_file, b_values_file in zip(*list_bids_files):
            dwi_image_file_bids = BIDSFileName.from_name(dwi_image_file)
            try:
                check_dwi_volume(
                    DWIDataset(
                        dwi=dwi_image_file,
                        b_values=b_values_file,
                        b_vectors=b_vectors_file,
                    )
                )
                subjects.append(f"sub-{dwi_image_file_bids.subject}")
                sessions.append(f"ses-{dwi_image_file_bids.session}")
            except (IOError, ValueError) as e:
                cprint(
                    f"Ignoring DWI scan for subject {dwi_image_file_bids.subject} "
                    f"and session {dwi_image_file_bids.session} for the following reason: {e}",
                    lvl="warning",
                )

        return subjects, sessions

    def _build_input_node(self):
        """Build and connect an input node to the pipeline."""
        import nipype.interfaces.utility as nutil
        import nipype.pipeline.engine as npe

        from clinica.utils.filemanip import save_participants_sessions
        from clinica.utils.input_files import (
            DWI_BVAL,
            DWI_BVEC,
            DWI_JSON,
            DWI_NII,
            T1W_NII,
        )
        from clinica.utils.inputs import clinica_list_of_files_reader
        from clinica.utils.stream import cprint
        from clinica.utils.ux import print_images_to_process

        list_bids_files = clinica_list_of_files_reader(
            self.subjects,
            self.sessions,
            self.bids_directory,
            [T1W_NII, DWI_NII, DWI_JSON, DWI_BVEC, DWI_BVAL],
            raise_exception=True,
        )

        save_participants_sessions(
            self.subjects, self.sessions, self.base_dir / self.name
        )

        if len(self.subjects):
            print_images_to_process(self.subjects, self.sessions)
            cprint(
                f"List available in {self.base_dir / self.name / 'participants.tsv'}"
            )
            cprint(
                "Computational time will depend of the number of volumes in your DWI dataset and the use of CUDA."
            )

        read_node = npe.Node(
            name="ReadingFiles",
            iterables=[
                (x, y) for x, y in zip(self.get_input_fields(), list_bids_files)
            ],
            synchronize=True,
            interface=nutil.IdentityInterface(fields=self.get_input_fields()),
        )
        self.connect(
            [
                (read_node, self.input_node, [(x, x) for x in self.get_input_fields()]),
            ]
        )

    def _build_output_node(self):
        """Build and connect an output node to the pipeline."""
        import nipype.interfaces.io as nio
        import nipype.interfaces.utility as nutil
        import nipype.pipeline.engine as npe

        from clinica.utils.nipype import container_from_filename, fix_join

        from .tasks import rename_into_caps_task

        container_path = npe.Node(
            nutil.Function(
                input_names=["bids_or_caps_filename"],
                output_names=["container"],
                function=container_from_filename,
            ),
            name="container_path",
        )
        files_to_write_in_caps = [
            "dwi_preproc_filename",
            "b_values_preproc_filename",
            "b_vectors_preproc_filename",
            "b0_brain_mask_filename",
        ]

        rename_into_caps = npe.Node(
            nutil.Function(
                input_names=["dwi_filename"] + files_to_write_in_caps,
                output_names=[f"{x}_caps" for x in files_to_write_in_caps],
                function=rename_into_caps_task,
            ),
            name="rename_into_caps",
        )

        write_results = npe.Node(name="write_results", interface=nio.DataSink())
        write_results.inputs.base_directory = str(self.caps_directory)
        write_results.inputs.parameterization = False

        self.connect(
            [
                (
                    self.input_node,
                    container_path,
                    [("dwi_filename", "bids_or_caps_filename")],
                ),
                (self.input_node, rename_into_caps, [("dwi_filename", "dwi_filename")]),
                (
                    self.output_node,
                    rename_into_caps,
                    [
                        ("preproc_dwi", "dwi_preproc_filename"),
                        ("preproc_bval", "b_values_preproc_filename"),
                        ("preproc_bvec", "b_vectors_preproc_filename"),
                        ("b0_mask", "b0_brain_mask_filename"),
                    ],
                ),
                (
                    container_path,
                    write_results,
                    [(("container", fix_join, "dwi"), "container")],
                ),
                (
                    rename_into_caps,
                    write_results,
                    [
                        ("dwi_preproc_filename_caps", "preprocessing.@preproc_dwi"),
                        (
                            "b_values_preproc_filename_caps",
                            "preprocessing.@preproc_bval",
                        ),
                        (
                            "b_vectors_preproc_filename_caps",
                            "preprocessing.@preproc_bvec",
                        ),
                        ("b0_brain_mask_filename_caps", "preprocessing.@b0_mask"),
                    ],
                ),
            ]
        )

    def _build_core_nodes(self):
        """Build and connect the core nodes of the pipeline."""
        import nipype.interfaces.fsl as fsl
        import nipype.interfaces.mrtrix3 as mrtrix3
        import nipype.interfaces.utility as nutil
        import nipype.pipeline.engine as npe

        from clinica.pipelines.dwi.preprocessing.tasks import compute_average_b0_task
        from clinica.pipelines.dwi.preprocessing.workflows import eddy_fsl_pipeline

        from .tasks import prepare_reference_b0_task
        from .utils import init_input_node, print_end_pipeline
        from .workflows import epi_pipeline

        init_node = npe.Node(
            interface=nutil.Function(
                input_names=self.get_input_fields(),
                output_names=[
                    "image_id",
                    "t1w_filename",
                    "dwi_filename",
                    "b_vectors_filename",
                    "b_values_filename",
                    "total_readout_time",
                    "phase_encoding_direction",
                ],
                function=init_input_node,
            ),
            name="0-InitNode",
        )

        # Prepare b0 image for further corrections
        prepare_b0 = npe.Node(
            name="PrepareB0",
            interface=nutil.Function(
                input_names=[
                    "dwi_filename",
                    "b_values_filename",
                    "b_vectors_filename",
                    "b_value_threshold",
                    "working_directory",
                ],
                output_names=[
                    "out_reference_b0",
                    "out_b0_dwi_merge",
                    "out_updated_bval",
                    "out_updated_bvec",
                ],
                function=prepare_reference_b0_task,
            ),
        )
        prepare_b0.inputs.b_value_threshold = self.parameters["low_bval"]
        prepare_b0.inputs.working_directory = str(self.base_dir)

        # Head-motion correction + Eddy-currents correction
        eddy_fsl = eddy_fsl_pipeline(
            base_dir=self.base_dir,
            use_cuda=self.parameters["use_cuda"],
            initrand=self.parameters["initrand"],
            compute_mask=True,
            image_id=True,
        )
        # Susceptibility distortion correction using T1w image
        sdc = epi_pipeline(
            str(self.base_dir),
            self.parameters["delete_cache"],
            name="SusceptibilityDistortionCorrection",
            ants_random_seed=self.parameters["random_seed"],
            use_double_precision=self.parameters["double_precision"],
        )

        # Remove bias correction from (Jeurissen et al., 2014)
        bias = npe.Node(
            mrtrix3.DWIBiasCorrect(
                use_ants=True,
                out_file="Jacobian_image_maths_thresh_merged_biascorr.nii.gz",
            ),
            name="RemoveBias",
        )
        # Compute b0 mask on corrected avg b0
        compute_avg_b0 = npe.Node(
            nutil.Function(
                input_names=["dwi_filename", "b_value_filename"],
                output_names=["out_b0_average"],
                function=compute_average_b0_task,
            ),
            name="ComputeB0Average",
        )
        compute_avg_b0.inputs.b_value_threshold = self.parameters["low_bval"]

        # Compute brain mask from reference b0
        mask_avg_b0 = npe.Node(fsl.BET(mask=True, robust=True), name="MaskB0")

        print_end_message = npe.Node(
            interface=nutil.Function(
                input_names=["image_id", "final_file"], function=print_end_pipeline
            ),
            name="WriteEndMessage",
        )

        self.connect(
            [
                (self.input_node, init_node, [(x, x) for x in self.get_input_fields()]),
                (
                    init_node,
                    prepare_b0,
                    [
                        (x, x)
                        for x in (
                            "dwi_filename",
                            "b_values_filename",
                            "b_vectors_filename",
                        )
                    ],
                ),
                (
                    init_node,
                    eddy_fsl,
                    [
                        (x, f"inputnode.{x}")
                        for x in (
                            "total_readout_time",
                            "phase_encoding_direction",
                            "image_id",
                        )
                    ],
                ),
                (
                    prepare_b0,
                    eddy_fsl,
                    [
                        ("out_b0_dwi_merge", "inputnode.dwi_filename"),
                        ("out_updated_bval", "inputnode.b_values_filename"),
                        ("out_updated_bvec", "inputnode.b_vectors_filename"),
                        ("out_reference_b0", "inputnode.reference_b0"),
                    ],
                ),
                # Magnetic susceptibility correction
                (init_node, sdc, [("t1w_filename", "inputnode.t1_filename")]),
                (
                    eddy_fsl,
                    sdc,
                    [("outputnode.out_corrected", "inputnode.dwi_filename")],
                ),
                (
                    eddy_fsl,
                    sdc,
                    [("outputnode.out_rotated_bvecs", "inputnode.b_vectors_filename")],
                ),
                # Bias correction
                (prepare_b0, bias, [("out_updated_bval", "in_bval")]),
                (
                    sdc,
                    bias,
                    [
                        ("outputnode.epi_corrected_dwi_image", "in_file"),
                        ("outputnode.rotated_b_vectors", "in_bvec"),
                    ],
                ),
                # Compute average b0 on corrected dataset (for brain mask extraction)
                (
                    prepare_b0,
                    compute_avg_b0,
                    [("out_updated_bval", "b_value_filename")],
                ),
                (bias, compute_avg_b0, [("out_file", "dwi_filename")]),
                # Compute b0 mask on corrected avg b0
                (compute_avg_b0, mask_avg_b0, [("out_b0_average", "in_file")]),
                # Print end message
                (init_node, print_end_message, [("image_id", "image_id")]),
                (mask_avg_b0, print_end_message, [("mask_file", "final_file")]),
                # Output node
                (bias, self.output_node, [("out_file", "preproc_dwi")]),
                (
                    sdc,
                    self.output_node,
                    [("outputnode.rotated_b_vectors", "preproc_bvec")],
                ),
                (prepare_b0, self.output_node, [("out_updated_bval", "preproc_bval")]),
                (mask_avg_b0, self.output_node, [("mask_file", "b0_mask")]),
            ]
        )
