#!/usr/bin/python3
'''
created on Wed 19th 2020
author: Christian Olaf Haeusler
'''


from glob import glob
import argparse
import os
import re
import subprocess


SUBMIT_CONTENT = [
    '# auto-generate file (generate_1st_level_design.py) -- do not modify!',
    'universe = vanilla',
    'output = condor_logs/$(CLUSTER).$(PROCESS).out',
    'error = condor_logs/$(CLUSTER).$(PROCESS).err',
    'log = condor_logs/$(CLUSTER).$(PROCESS).log',
    'getenv = True',
    'request_cpus = 1',
    'request_memory = 5000',
    'should_transfer_files = NO',
    'transfer_executable = False',
    'initialdir = /data/project/studyforrest_ppa',
    'executable = $ENV(FSLDIR)/bin/feat',
    '',
]


def parse_arguments():
    '''
    '''
    parser = argparse.ArgumentParser(
        description='rename and copy the auditory \
        confound files to onset directories')

    parser.add_argument('-fmri',
                        default='inputs/studyforrest-data-aligned/sub-01/in_bold3Tp2/sub-01_task-aomovie_run-1_bold.nii.gz',
                        help='example 4d fmri file/path from sub-01')

    parser.add_argument('-design',
                        default='code/1st-lvl_audio-ppa-grp.fsf',
                        help='the design template (for sub-01)')

    args = parser.parse_args()

    fmriEx = args.fmri
    designFile = args.design

    return fmriEx, designFile


# main program #
if __name__ == "__main__":
    # prepare the header for the condor .submit file
    submit_content = [line + '\n' for line in SUBMIT_CONTENT]

    # with launch_ipdb_on_exception():
    # call FSL first
    os.system('. /etc/fsl/5.0/fsl.sh')

    fmri_example, design_file = parse_arguments()

    # for 3T data input
    if 'sub-01' in fmri_example:
        subjs_pattern = fmri_example.replace('sub-01', 'sub-*')
        # search for subjects
        pathes = sorted(glob(subjs_pattern))
        subjs = [re.search(r'sub-\d{2}', path) for path in pathes]
        subjs = [int(subj.group()[-2:]) for subj in subjs]

        # filter out sub-10 that is not in intersection of subjects for
        # both movie and audio-description (missing in audio-description)
        if 10 in subjs:
            subjs.remove(10)

        runs_pattern = fmri_example.replace('run-1', 'run-*')
        pathes = sorted(glob(runs_pattern))
        runs = [re.search(r'run-\d', path) for path in pathes]
        runs = [(run.group()) for run in runs]
        runs = [int(run[-1]) for run in runs]
    else:
        print('subject string not found')

    # open design file
    with open(design_file, 'r') as f:
        lines = f.readlines()
    # use the design file as template
    template = [line.strip() for line in lines]

    for sub in subjs:
        sub = str(sub)
        print('\nsub', sub)
        for run in runs:
            run = str(run)
            print('run', run)

            # number of fmri volumes for current run
            fmri_file = fmri_example.replace('sub-01', 'sub-' + sub.zfill(2))
            fmri_file = fmri_file.replace('run-1', 'run-' + run)
            info =  subprocess.check_output('fslinfo %s' % fmri_file, shell=True)
            dim4 = info.split()[9]
            vol = str(dim4).split("'")[1]

            design_file_content = []
            for lineNo, line in enumerate(template):
                if 'sub-01' in line or 'run-1' in line:
                    line = line.replace('sub-01', 'sub-' + sub.zfill(2))
                    line = line.replace('run-1', 'run-' + run)

                # correct bold_dico_moco.txt (motion correction for audio data)
                elif 'sub001' in line or 'run001' in line:
                   line = line.replace('sub001', 'sub' + sub.zfill(3))
                   line = line.replace('run001', 'run' + run.zfill(3))

                elif 'set fmri(npts) ' in line:
                    line_splitted = line.split()
                    line_splitted[-1] = ' ' + vol
                    line = ' '.join(line_splitted)

                elif 'set fmri(evtitle' in line:
                    # check for the corresponding EV3 file
                    # it's name comes always 31 lines later in the file
                    ev = line.split('"')[1]
                    ev_file_line = template[lineNo + 31]
                    ev_file = ev_file_line.split('"')[1]
                    ev_file = ev_file.replace('sub-01', 'sub-' + sub.zfill(2))
                    ev_file = ev_file.replace('run-1', 'run-' + run)

                    if os.path.isfile(ev_file):
                        with open(ev_file, 'r') as f:
                            content = f.readlines()
                            length = len(content)
                    # check the length of the event onset file
                        # use canonical HRF if it contains lines/events
                        if length > 0:
                            fmri_shape = '3'
                            # use 0 regressor otherwise
                        else:
                            fmri_shape = '10'
                            print('no events in', ev_file)
                            print('choosing null events for ev %s' % ev)
                    else:
                        print('could not find', ev_file)
                        fmri_shape = '10'
                        print('choosing null events for ev %s' % ev)


                elif 'set fmri(shape' in line:
                    line_splitted = line.split()
                    line_splitted[-1] = fmri_shape
                    line = ' '.join(line_splitted)

                design_file_content.append(line)

            design_file_content = [line + '\n' for line in design_file_content]

            # writing
            fsl_design = os.path.basename(design_file).split('.')[0]
            analysis = fsl_design.split('lvl_')[1]
            out_name = 'run-%s_1st_%s.fsf' % (run, analysis)
            out_path = os.path.join('sub-' + sub.zfill(2), out_name)

            try:
                with open(out_path, 'w') as the_file:
                    the_file.writelines(design_file_content)

                    submit_arg = 'arguments = sub-%s/%s\nqueue\n' % (sub.zfill(2), out_name)
                    submit_content.append(submit_arg)
            except IOError:
                print(sub, 'is not in intersection of movie and audio-description')


        # write the condor_submit file
        submit_name = 'compute_1st-lvl_%s.submit' %  analysis
        submit_path = os.path.join(os.path.dirname(design_file), submit_name)

        with open(submit_path, 'w') as the_file:
            the_file.writelines(submit_content)

        os.makedirs('condor_logs', exist_ok=True)