#!python
# Create a table of key BOLD fMRI parameters for a BIDS dataset
#
# AUTHOR : Mike Tyszka
# PLACE  : Caltech
# DATES  : 2021-05-06 JMT From scratch
#          2025-05-03 JMT Fix handling of missing fields

import os.path as op
import sys
import argparse
import json
import bids

import pandas as pd
import nibabel as nib

from tabulate import tabulate

from glob import glob


def main():

    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Create table of key BOLD fMRI parameters for a BIDS dataset')
    parser.add_argument('-d', '--dataset', default='.', help='BIDS dataset directory')
    args = parser.parse_args()

    bids_dir = op.realpath(args.dataset)

    print('')
    print('BIDS BOLD Parameter Table')
    print('-------------------------')
    print('BIDS dataset directory : {}'.format(bids_dir))

    print('Searching for BOLD fMRI images')
    bold_json_list = sorted(glob(op.join(bids_dir, 'sub-*', 'ses-*', 'func', '*_bold.json')))

    if len(bold_json_list) < 1:
        print('* No BOLD fMRI images found - exiting')
        sys.exit(1)

    bold_info_list = []
    
    for bold_json_path in bold_json_list:

        bold_dir = op.dirname(bold_json_path)
        bold_json_fname = op.basename(bold_json_path)
        bold_nii_fname = bold_json_fname.replace('.json', '.nii.gz')
        bold_nii_path = op.join(bold_dir, bold_nii_fname)

        print('  Processing {}'.format(bold_nii_fname))

        # Load Nifti header
        bold_nii = nib.load(bold_nii_path)
        bold_hdr = bold_nii.header

        # Image dimensions
        dim = bold_hdr['dim']
        vox = bold_hdr['pixdim']

        # Parse BIDS filename
        keys = bids.layout.parse_file_entities(bold_json_path)

        # Load JSON metadata
        with open(bold_json_path, 'r') as fd:
            meta = json.load(fd)

        # In-plane GRAPPA factor if present
        if 'ParallelReductionFactorInPlane' in meta:
            r = meta['ParallelReductionFactorInPlane']
        else:
            r = 1

        #  Multiband acceleration factor if present
        if 'MultibandAccelerationFactor' in meta:
            m = meta['MultibandAccelerationFactor']
        else:
            m = 1

        # Handle commonly absent fields
        if not 'session' in keys:
            keys['session'] = '-'
        if not 'task' in keys:
            keys['task'] = '-'
        if not 'run' in keys:
            keys['run'] = 1

        # Add to running list
        bold_info_list.append([
            bold_nii_fname,
            keys['subject'],
            keys['session'],
            keys['task'],
            keys['run'],
            meta['RepetitionTime'],
            meta['EchoTime'],
            meta['FlipAngle'],
            r, m,
            dim[1], dim[2], dim[3], dim[4],
            vox[1], vox[2], vox[3]
        ])
               
    # Convert to dataframe and save as CSV
    df = pd.DataFrame(
        bold_info_list,
        columns=[
           'Filename',
           'Subject',
           'Session',
           'Task',
           'Run',
           'TR_secs',
           'TE_secs',
           'FlipAngle_degs',
           'R', 'M',
           'nx', 'ny', 'nz', 'nt',
           'vx_mm', 'vy_mm', 'vz_mm'
        ]
    )

    csv_fname = op.join(bids_dir, 'bidsdump.csv')
    print('')
    print('Saving BIDS BOLD info to {}'.format(csv_fname))
    df.to_csv(csv_fname, index=False)

    # Pretty print dataframe as a table to stdout
    print('')
    print(tabulate(df, headers='keys', tablefmt='pretty'))


if "__main__" in __name__:

    main()
