#!/bin/python3

"""
About this Script

The `build_task_info()` function creates the relevant `task_information.json` file
that is critical to running first- and second-level models using GLM-Express

Ian Richard Ferguson | Stanford University
"""

import os, json, sys
from bids import BIDSLayout

def derive_tasks(bids_root):
      """
      Feeds user-supplied BIDS root into BIDSLayout object to obtain list of functional tasks

      Parameters
            bids_root: str | Relative path to the top of your BIDS project

      Returns
            List of functional tasks
      """
      
      # Feed bids_root into BIDSLayout object
      # NOTE: This will throw an error if your project is not BIDS-compliant
      bids = BIDSLayout(bids_root)
      
      # Return functional tasks in list
      return bids.get_tasks()


def build_task_info(bids_root):
      """
      This function constructs the `task_information.json` file that is required
      for the Subject and GroupLevel objects to run successfully

      Parameters
            bids_root: str | Relative path to the top of your BIDS project
      """

      # Derive list of tasks in your BIDS project
      functional_tasks = derive_tasks(bids_root)

      # Empty dictionary to append into
      output = {}

      # Loop through tasks and create dictionaries for each
      for task in functional_tasks:

            output[task] = {}                                           # Empty dictionary to append into

            output[task]['block_identifier'] = 'block_type'             # Column in events file corresponding to Block identifier (if applicable) 
            output[task]['condition_identifier'] = 'trial_type'         # Column in events file corresponding to Trial Type identifier
            output[task]['confound_regressors'] = []                    # Regressors to include from fmriprep confounds file
            output[task]['design_contrasts'] = 'default'                # Contrasts of interest to run in first-level design
            output[task]['excludes'] = []                               # List of subject IDs to exclude (for batching / SLURM scheduling)
            output[task]['group_level_regressors'] = []                 # Regressors to include at second-level (NOTE: In development)
            output[task]['modulators'] = []                             # Parametric modulators to include in the first-level analysis (should be columns in events file)
            output[task]['tr'] = 1.                                     # Repetition time for the given first-level task (you can update this in the Subject object too)

      # Save task file locally
      with open(os.path.join('./task_information.json'), 'w') as outgoing:
            json.dump(output, outgoing, indent=6)


if __name__ == "__main__":
      try:
            root = sys.argv[1]
      except:
            raise OSError(f'No command line argument supplied')

      build_task_info(root)
