# coding: utf8


import clinica.pipelines.engine as cpe


class T1FreeSurferTemplate(cpe.Pipeline):
    """FreeSurfer Longitudinal template class

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

    @staticmethod
    def get_processed_images(caps_directory, subjects, sessions):
        import os
        from clinica.utils.inputs import clinica_file_reader
        from clinica.utils.input_files import T1_FS_T_DESTRIEUX
        from clinica.utils.longitudinal import get_long_id
        from clinica.utils.participant import get_unique_subjects
        import re

        [list_participant_id, list_list_session_ids] = get_unique_subjects(subjects, sessions)
        list_long_id = [get_long_id(list_session_ids) for list_session_ids in list_list_session_ids]

        image_ids = []
        if os.path.isdir(caps_directory):
            t1_freesurfer_files = clinica_file_reader(
                list_participant_id, list_long_id,
                caps_directory, T1_FS_T_DESTRIEUX, False
            )
            image_ids = [re.search(r'(sub-[a-zA-Z0-9]+)_(long-[a-zA-Z0-9]+)', file).group()
                         for file in t1_freesurfer_files]
        return image_ids

    def check_pipeline_parameters(self):
        """Check pipeline parameters."""
        pass

    def check_custom_dependencies(self):
        """Check dependencies that cannot be listed in `info.json`
        """
        pass

    def get_input_fields(self):
        """Specify the list of possible inputs of this pipeline.

        Note:
            The list of inputs of the T1FreeSurferTemplate pipeline is:
                * participant_id (str): Participant ID (e.g. "sub-CLNC01")
                * list_session_ids (List[str]: List of session IDs associated to `participant_id`
                    (e.g. ["ses-M00"] or ["ses-M00", "ses-M18", "ses-M36"])

        Returns:
            A list of (string) input fields name.
        """
        return ['participant_id', 'list_session_ids']

    def get_output_fields(self):
        """Specify the list of possible outputs of this pipeline.

        Note:
            The list of inputs of the T1FreeSurferTemplate pipeline is:
                * image_id (str): Image ID (e.g. "sub-CLNC01_long-M00M18")

        Returns:
            A list of (string) output fields name.
        """
        return ['image_id']

    def build_input_node(self):
        """Build and connect an input node to the pipeline."""
        import os
        from colorama import Fore

        import nipype.interfaces.utility as nutil
        import nipype.pipeline.engine as npe

        from clinica.utils.exceptions import ClinicaException, ClinicaCAPSError
        from clinica.utils.filemanip import extract_subjects_sessions_from_filename
        from clinica.utils.inputs import clinica_file_reader
        from clinica.utils.input_files import T1_FS_DESTRIEUX
        from clinica.utils.longitudinal import get_long_id, read_sessions, get_participants_long_id
        from clinica.utils.participant import get_unique_subjects, unique_subjects_sessions_to_subjects_sessions
        from clinica.utils.stream import cprint
        from .longitudinal_utils import extract_participant_long_ids_from_filename, save_part_sess_long_ids_to_tsv

        # Display image(s) already present in CAPS folder
        # ===============================================
        output_ids = self.get_processed_images(self.caps_directory, self.subjects, self.sessions)
        processed_participants, processed_long_sessions = extract_participant_long_ids_from_filename(output_ids)
        if len(processed_participants) > 0:
            cprint("%sClinica found %s participant(s) already processed in CAPS directory:%s" %
                   (Fore.YELLOW, len(processed_participants), Fore.RESET))
            for p_id, l_id in zip(processed_participants, processed_long_sessions):
                cprint("%s\t%s | %s%s" % (Fore.YELLOW, p_id, l_id, Fore.RESET))
            if self.overwrite_caps:
                output_folder = "<CAPS>/subjects/<participant_id>/<long_id>/freesurfer_unbiased_template/"
                cprint("%s\nOutput folders in %s will be recreated.\n%s" % (Fore.YELLOW, output_folder, Fore.RESET))
            else:
                cprint("%s\nParticipant(s) will be ignored by Clinica.\n%s" % (Fore.YELLOW, Fore.RESET))
                input_ids = [p_id + '_' + s_id for p_id, s_id in
                             zip(self.subjects, self.sessions)]
                processed_sessions_per_participant = [
                    read_sessions(self.caps_directory, p_id, l_id)
                    for (p_id, l_id) in zip(processed_participants, processed_long_sessions)
                ]
                participants, sessions = unique_subjects_sessions_to_subjects_sessions(
                    processed_participants,
                    processed_sessions_per_participant)
                processed_ids = [p_id + '_' + s_id for p_id, s_id in zip(participants, sessions)]
                to_process_ids = list(set(input_ids) - set(processed_ids))
                self.subjects, self.sessions = extract_subjects_sessions_from_filename(to_process_ids)

        # Check that t1-freesurfer has run on the CAPS directory
        try:
            clinica_file_reader(self.subjects,
                                self.sessions,
                                self.caps_directory,
                                T1_FS_DESTRIEUX)
        except ClinicaException as e:
            err_msg = 'Clinica faced error(s) while trying to read files in your CAPS directory.\n' + str(e)
            raise ClinicaCAPSError(err_msg)

        # Save subjects to process in <WD>/<Pipeline.name>/participants.tsv
        folder_participants_tsv = os.path.join(self.base_dir, self.name)
        long_ids = get_participants_long_id(self.subjects, self.sessions)
        save_part_sess_long_ids_to_tsv(self.subjects, self.sessions, long_ids, folder_participants_tsv)

        [list_participant_id, list_list_session_ids] = get_unique_subjects(self.subjects, self.sessions)
        list_long_id = [get_long_id(list_session_ids) for list_session_ids in list_list_session_ids]

        def print_images_to_process(unique_part_list, per_part_session_list, list_part_long_id):
            cprint('The pipeline will be run on the following %s participant(s):' % len(unique_part_list))
            for (part_id, list_sess_id, list_id) in zip(unique_part_list, per_part_session_list, list_part_long_id):
                sessions_participant = ', '.join(s_id for s_id in list_sess_id)
                cprint("\t%s | %s | %s" % (part_id, sessions_participant, list_id))

        if len(self.subjects):
            # TODO: Generalize long IDs to the message display
            print_images_to_process(list_participant_id, list_list_session_ids, list_long_id)
            cprint('List available in %s' % os.path.join(folder_participants_tsv, 'participants.tsv'))
            cprint('The pipeline will last approximately 10 hours per participant.')

        read_node = npe.Node(name="ReadingFiles",
                             iterables=[
                                 ('participant_id', list_participant_id),
                                 ('list_session_ids', list_list_session_ids),
                             ],
                             synchronize=True,
                             interface=nutil.IdentityInterface(
                                 fields=self.get_input_fields())
                             )
        self.connect([
            (read_node, self.input_node, [('participant_id', 'participant_id')]),
            (read_node, self.input_node, [('list_session_ids', 'list_session_ids')]),
        ])

    def build_output_node(self):
        """Build and connect an output node to the pipeline."""
        import nipype.pipeline.engine as npe
        import nipype.interfaces.utility as nutil
        from .t1_freesurfer_template_utils import save_to_caps
        import os

        save_to_caps = npe.Node(interface=nutil.Function(
            input_names=['source_dir', 'image_id', 'list_session_ids', 'caps_dir', 'overwrite_caps'],
            output_names=['image_id'],
            function=save_to_caps),
            name='SaveToCaps')
        save_to_caps.inputs.source_dir = os.path.join(self.base_dir, self.name, 'ReconAll')
        save_to_caps.inputs.caps_dir = self.caps_directory
        save_to_caps.inputs.overwrite_caps = False

        self.connect([
            (self.input_node, save_to_caps, [('list_session_ids', 'list_session_ids')]),
            (self.output_node, save_to_caps, [('image_id', 'image_id')]),
        ])

    def build_core_nodes(self):
        """Build and connect the core nodes of the pipeline."""
        import os
        import nipype.interfaces.utility as nutil
        import nipype.pipeline.engine as npe
        from .t1_freesurfer_template_utils import init_input_node, run_recon_all_base, move_subjects_dir_to_source_dir

        # Nodes declaration
        # =================
        # Initialize the pipeline
        init_input = npe.Node(
            interface=nutil.Function(
                input_names=['caps_dir', 'participant_id', 'list_session_ids', 'output_dir'],
                output_names=['image_id', 'subjects_dir', 'flags'],
                function=init_input_node),
            name='0-InitPipeline')
        init_input.inputs.caps_dir = self.caps_directory
        init_input.inputs.output_dir = os.path.join(self.base_dir, self.name, 'ReconAll')

        # Run recon-all command
        recon_all = npe.Node(
            interface=nutil.Function(
                input_names=['subjects_dir', 'subject_id', 'flags', 'directive'],
                output_names=['subject_id'],
                function=run_recon_all_base),
            name='1-SegmentationReconAll')
        recon_all.inputs.directive = '-all'

        # Move $SUBJECT_DIR to source_dir (1 time point case)
        move_subjects_dir = npe.Node(
            interface=nutil.Function(
                input_names=['subjects_dir', 'source_dir', 'subject_id'],
                output_names=['subject_id'],
                function=move_subjects_dir_to_source_dir),
            name='2-MoveSubjectsDir')
        move_subjects_dir.inputs.source_dir = os.path.join(self.base_dir, self.name, 'ReconAll')

        # Connections
        # ===========
        self.connect([
            # Initialize the pipeline
            (self.input_node, init_input, [('participant_id', 'participant_id')]),
            (self.input_node, init_input, [('list_session_ids', 'list_session_ids')]),
            # Run recon-all command
            (init_input, recon_all, [('subjects_dir', 'subjects_dir'),
                                     ('image_id', 'subject_id'),
                                     ('flags',  'flags')]),
            # Move $SUBJECT_DIR to source_dir (1 time point case)
            (init_input, move_subjects_dir, [('subjects_dir', 'subjects_dir')]),
            (recon_all, move_subjects_dir, [('subject_id', 'subject_id')]),
            # Output node
            (move_subjects_dir, self.output_node, [('subject_id', 'image_id')]),
        ])
