import os
from glob import glob


def rename_func(basename):
    name, ext = os.path.splitext(basename)
    assert ext in ['.npy', '.tsv'], f'{basename} {ext}'
    info = {}
    chunks = os.path.basename(name).split('_')
    if chunks[-1] == 'timeseries':
        extra = '_timeseries'
        chunks = chunks[:-1]
    else:
        extra = ''

    for chunk in chunks:
        key, val = chunk.split('-')
        info[key] = val

    for key in info:
        if key not in ['sub', 'ses', 'task', 'run', 'desc']:
            raise ValueError
    sid, ses, task, run = info['sub'], info['ses'], info['task'], info['run']

    if ses == 'movie' and task == 'movie':
        new_run = int(run)
        new_task = 'forrest'
    elif task in ['objectcategories', 'retmapccw', 'retmapclw', 'retmapcon', 'retmapexp', 'movielocalizer']:
        new_run = int(run)
        new_task = task
    else:
        raise ValueError

    label = f'sub-{sid}_task-{new_task}_run-{new_run:02d}'
    if 'desc' in info:
        label += f'_desc-{info["desc"]}'
    new_name = label + extra + ext
    return new_name


if __name__ == '__main__':
    fns = sorted(glob(os.path.join('*', 'resampled', '*', '*', '*', 'sub-*.npy')))
    print(len(fns))
    for fn in fns:
        parts = []
        head = fn
        while head:
            head, tail = os.path.split(head)
            parts.insert(0, tail)
        fmriprep_version, _, space, structure, resample, basename = parts

        out_basename = rename_func(basename)
        out_fn = os.path.join(fmriprep_version, 'renamed', space, structure, resample, out_basename)
        relpath = os.path.relpath(fn, os.path.dirname(out_fn))
        print(relpath)
        os.makedirs(os.path.dirname(out_fn), exist_ok=True)
        if not os.path.lexists(out_fn):
            os.symlink(relpath, out_fn)

    fns = sorted(glob(os.path.join('*', 'confounds', 'sub-*.npy')) + glob(os.path.join('*', 'confounds', 'sub-*.tsv')))
    print(len(fns))
    for fn in fns:
        parts = []
        head = fn
        while head:
            head, tail = os.path.split(head)
            parts.insert(0, tail)
        print(parts)
        fmriprep_version, _, basename = parts
        out_basename = rename_func(basename)
        out_fn = os.path.join(fmriprep_version, 'renamed_confounds', out_basename)
        relpath = os.path.relpath(fn, os.path.dirname(out_fn))
        print(relpath)
        os.makedirs(os.path.dirname(out_fn), exist_ok=True)
        if not os.path.lexists(out_fn):
            os.symlink(relpath, out_fn)