# coding: utf8

import clinica.pipelines.engine as cpe


class T1VolumeCreateDartel(cpe.Pipeline):
    """T1VolumeCreateDartel - Create new Dartel template.

    Returns:
        A clinica pipeline object containing the T1VolumeCreateDartel pipeline.
    """
    def check_pipeline_parameters(self):
        """Check pipeline parameters."""
        from clinica.utils.group import check_group_label

        if 'group_id' not in self.parameters.keys():
            raise KeyError('Missing compulsory group_id key in pipeline parameter.')
        if 'dartel_tissues' not in self.parameters.keys():
            self.parameters['dartel_tissues'] = [1, 2, 3]

        check_group_label(self.parameters['group_id'])

    def check_custom_dependencies(self):
        """Check dependencies that can not be listed in the `info.json` file."""
        pass

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

        Returns:
            A list of (string) input fields name.
        """

        return ['dartel_inputs']

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

        Returns:
            A list of (string) output fields name.
        """

        return ['final_template_file', 'template_files', 'dartel_flow_fields']

    def build_input_node(self):
        """Build and connect an input node to the pipeline."""
        import os
        import sys
        from colorama import Fore
        import nipype.pipeline.engine as npe
        import nipype.interfaces.utility as nutil
        from clinica.utils.inputs import clinica_file_reader
        from clinica.utils.input_files import t1_volume_dartel_input_tissue
        from clinica.utils.exceptions import ClinicaException
        from clinica.utils.stream import cprint
        from clinica.utils.ux import print_groups_in_caps_directory, print_images_to_process, print_begin_image

        representative_output = os.path.join(self.caps_directory,
                                             'groups',
                                             'group-' + self.parameters['group_id'],
                                             't1',
                                             'group-' + self.parameters['group_id'] + '_template.nii.gz'
                                             )
        if os.path.exists(representative_output):
            cprint("%sDARTEL template for %s already exists. Currently, Clinica does not propose to overwrite outputs "
                   "for this pipeline.%s" % (Fore.YELLOW, self.parameters['group_id'], Fore.RESET))
            print_groups_in_caps_directory(self.caps_directory)
            sys.exit(0)

        # Check that there is at least 2 subjects
        if len(self.subjects) <= 1:
            raise ClinicaException('%sThis pipeline needs at least 2 images to create DARTEL template but '
                                   'Clinica only found %s.%s' % (Fore.RED, len(self.subjects), Fore.RESET))

        read_parameters_node = npe.Node(name="LoadingCLIArguments",
                                        interface=nutil.IdentityInterface(fields=self.get_input_fields(),
                                                                          mandatory_inputs=True))
        all_errors = []
        d_input = []
        for tissue_number in self.parameters['dartel_tissues']:
            try:
                current_file = clinica_file_reader(self.subjects,
                                                   self.sessions,
                                                   self.caps_directory,
                                                   t1_volume_dartel_input_tissue(tissue_number))
                d_input.append(current_file)
            except ClinicaException as e:
                all_errors.append(e)

        # Raise all errors if some happened
        if len(all_errors) > 0:
            error_message = 'Clinica faced errors while trying to read files in your BIDS or CAPS directories.\n'
            for msg in all_errors:
                error_message += str(msg)
            raise RuntimeError(error_message)

        # d_input is a list of size len(self.parameters['dartel_tissues'])
        #     Each element of this list is a list of size len(self.subjects)
        read_parameters_node.inputs.dartel_inputs = d_input

        if len(self.subjects):
            print_images_to_process(self.subjects, self.sessions)
            cprint('Computational time for DARTEL creation will depend on the number of images.')
            print_begin_image('group-' + self.parameters['group_id'])

        self.connect([
            (read_parameters_node, self.input_node, [('dartel_inputs', 'dartel_input_images')])
        ])

    def build_output_node(self):
        """Build and connect an output node to the pipeline."""
        import os.path as op
        import nipype.pipeline.engine as npe
        import nipype.interfaces.io as nio
        import re
        from clinica.utils.filemanip import zip_nii

        # Writing flowfields into CAPS
        # ============================
        write_flowfields_node = npe.MapNode(name='write_flowfields_node',
                                            iterfield=['container', 'flow_fields'],
                                            interface=nio.DataSink(infields=['flow_fields']))
        write_flowfields_node.inputs.base_directory = self.caps_directory
        write_flowfields_node.inputs.parameterization = False
        write_flowfields_node.inputs.container = ['subjects/' + self.subjects[i] + '/' + self.sessions[i] +
                                                  '/t1/spm/dartel/group-' + self.parameters['group_id']
                                                  for i in range(len(self.subjects))]
        write_flowfields_node.inputs.regexp_substitutions = [
            (r'(.*)_Template(\.nii(\.gz)?)$', r'\1\2'),
            (r'(.*)c1(sub-.*)(\.nii(\.gz)?)$', r'\1\2_segm-graymatter\3'),
            (r'(.*)c2(sub-.*)(\.nii(\.gz)?)$', r'\1\2_segm-whitematter\3'),
            (r'(.*)c3(sub-.*)(\.nii(\.gz)?)$', r'\1\2_segm-csf\3'),
            (r'(.*)c4(sub-.*)(\.nii(\.gz)?)$', r'\1\2_segm-bone\3'),
            (r'(.*)c5(sub-.*)(\.nii(\.gz)?)$', r'\1\2_segm-softtissue\3'),
            (r'(.*)c6(sub-.*)(\.nii(\.gz)?)$', r'\1\2_segm-background\3'),
            (r'(.*)r(sub-.*)(\.nii(\.gz)?)$', r'\1\2\3'),
            (r'(.*)_dartelinput(\.nii(\.gz)?)$', r'\1\2'),
            (r'(.*)flow_fields/u_(sub-.*)_segm-.*(\.nii(\.gz)?)$',
             r'\1\2_target-' + re.escape(self.parameters['group_id']) + r'_transformation-forward_deformation\3'),
            (r'trait_added', r'')
        ]

        # Writing templates into CAPS
        # ===========================
        write_template_node = npe.Node(nio.DataSink(), name='write_template_node')
        write_template_node.inputs.parameterization = False
        write_template_node.inputs.base_directory = self.caps_directory
        write_template_node.inputs.container = op.join('groups/group-' + self.parameters['group_id'], 't1')
        write_template_node.inputs.regexp_substitutions = [
            (r'(.*)final_template_file/.*(\.nii(\.gz)?)$',
             r'\1group-' + re.escape(self.parameters['group_id']) + r'_template\2'),
            (r'(.*)template_files/.*([0-9])(\.nii(\.gz)?)$',
             r'\1group-' + re.escape(self.parameters['group_id']) + r'_iteration-\2_template\3')
        ]

        self.connect([
            (self.output_node, write_flowfields_node, [(('dartel_flow_fields', zip_nii, True), 'flow_fields')]),
            (self.output_node, write_template_node, [(('final_template_file', zip_nii, True), 'final_template_file'),
                                                     (('template_files', zip_nii, True), 'template_files')])
        ])

    def build_core_nodes(self):
        """Build and connect the core nodes of the pipeline."""
        import nipype.interfaces.spm as spm
        import nipype.pipeline.engine as npe
        import nipype.interfaces.utility as nutil
        from clinica.utils.filemanip import unzip_nii
        from clinica.utils.spm import spm_standalone_is_available, use_spm_standalone

        if spm_standalone_is_available():
            use_spm_standalone()

        # Unzipping
        # =========
        unzip_node = npe.MapNode(nutil.Function(input_names=['in_file'],
                                                output_names=['out_file'],
                                                function=unzip_nii),
                                 name='unzip_node', iterfield=['in_file'])

        # DARTEL template
        # ===============
        dartel_template = npe.Node(spm.DARTEL(),
                                   name='dartel_template')

        # Connection
        # ==========
        self.connect([
            (self.input_node, unzip_node,    [('dartel_input_images', 'in_file')]),
            (unzip_node, dartel_template,    [('out_file', 'image_files')]),
            (dartel_template, self.output_node, [('dartel_flow_fields', 'dartel_flow_fields'),
                                                 ('final_template_file', 'final_template_file'),
                                                 ('template_files', 'template_files')])
        ])
