import re
from itertools import groupby

multi_echo_ses = {'data': ['sub-01_ses-01_task-view_run-01_echo-1_bold.nii.gz',
                           'sub-01_ses-01_task-view_run-01_echo-2_bold.nii.gz',
                           'sub-01_ses-01_task-view_run-01_echo-3_bold.nii.gz',
                           'sub-01_ses-01_task-view_run-02_echo-1_bold.nii.gz',
                           'sub-01_ses-01_task-view_run-02_echo-2_bold.nii.gz',
                           'sub-01_ses-01_task-view_run-02_echo-3_bold.nii.gz',
                           'sub-01_ses-01_task-rest_run-01_bold.nii.gz',
                           'sub-01_ses-02_task-rest_run-01_bold.nii.gz'],
                  'grouping': [['sub-01_ses-01_task-view_run-01_echo-1_bold.nii.gz',
                                'sub-01_ses-01_task-view_run-01_echo-2_bold.nii.gz',
                                'sub-01_ses-01_task-view_run-01_echo-3_bold.nii.gz'],
                               ['sub-01_ses-01_task-view_run-02_echo-1_bold.nii.gz',
                                'sub-01_ses-01_task-view_run-02_echo-2_bold.nii.gz',
                                'sub-01_ses-01_task-view_run-02_echo-3_bold.nii.gz'],
                               'sub-01_ses-01_task-rest_run-01_bold.nii.gz',
                               'sub-01_ses-02_task-rest_run-01_bold.nii.gz']}

ds000205 = {'data': ['sub-07_task-functionallocalizer_run-01_bold.nii.gz',
                     'sub-07_task-view_run-01_bold.nii.gz',
                     'sub-07_task-view_run-02_bold.nii.gz'],
            'grouping': ['sub-07_task-functionallocalizer_run-01_bold.nii.gz',
                         'sub-07_task-view_run-01_bold.nii.gz',
                         'sub-07_task-view_run-02_bold.nii.gz']}

ds000210 = {'data': ['sub-02_task-cuedSGT_run-01_echo-1_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-01_echo-2_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-01_echo-3_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-02_echo-1_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-02_echo-2_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-02_echo-3_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-03_echo-1_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-03_echo-2_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-03_echo-3_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-04_echo-1_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-04_echo-2_bold.nii.gz',
                     'sub-02_task-cuedSGT_run-04_echo-3_bold.nii.gz',
                     'sub-02_task-rest_run-01_echo-1_bold.nii.gz',
                     'sub-02_task-rest_run-01_echo-2_bold.nii.gz',
                     'sub-02_task-rest_run-01_echo-3_bold.nii.gz'],
            'grouping': [['sub-02_task-cuedSGT_run-01_echo-1_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-01_echo-2_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-01_echo-3_bold.nii.gz'],
                         ['sub-02_task-cuedSGT_run-02_echo-1_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-02_echo-2_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-02_echo-3_bold.nii.gz'],
                         ['sub-02_task-cuedSGT_run-03_echo-1_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-03_echo-2_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-03_echo-3_bold.nii.gz'],
                         ['sub-02_task-cuedSGT_run-04_echo-1_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-04_echo-2_bold.nii.gz',
                          'sub-02_task-cuedSGT_run-04_echo-3_bold.nii.gz'],
                         ['sub-02_task-rest_run-01_echo-1_bold.nii.gz',
                          'sub-02_task-rest_run-01_echo-2_bold.nii.gz',
                          'sub-02_task-rest_run-01_echo-3_bold.nii.gz']]}


def _grp_echos(x):
    if '_echo-' not in x:
        return x
    echo = re.search("_echo-\\d*", x).group(0)
    return x.replace(echo, "_echo-?")


# mockup a subject_data dictionary
for mock_data in [multi_echo_ses, ds000205, ds000210]:
    subj_data = {'bold': mock_data['data']}

    if subj_data["bold"] is not []:
        bold_sess = subj_data["bold"]

        if any(['_echo-' in bold for bold in bold_sess]):
            ses_uids = [list(bold) for _, bold in groupby(bold_sess, key=_grp_echos)]
            ses_uids = list(map(lambda x: x[0] if len(x) == 1 else x, ses_uids))
        else:
            ses_uids = bold_sess

        subj_data.update({"bold": ses_uids})
        assert subj_data['bold'] == mock_data['grouping']