#!/usr/bin/env python

import os
import argparse
import surfa as sf
from freesurfer import subregions


description = f'''
Cross-sectional and longitudinal segmentation for the following
structures: {", ".join(subregions.structure_names)}. To segment
the thalamic nuclei, for example, in a cross-sectional analysis:

    segment_subregions thalamus --cross subj

Similarly, for a longitudinal analysis, run:

    segment_subregions thalamus --long-base base

Timepoints are extracted from the base-tps file in the `base`
subject. Output segmentations and computed structure volumes
will be saved to the subject's `mri` subdirectory.
'''

parser = argparse.ArgumentParser(description=description)
parser.add_argument('structure', help=f'Structure to segment. Options are: {", ".join(subregions.structure_names)}.')
parser.add_argument('--cross', help='Subject to segment in cross-sectional analysis.')
parser.add_argument('--long-base', help='Base subject for longitudinal analysis. Timepoints are extracted from the base-tps file.')
parser.add_argument('--sd', help='Specify subjects directory (will override SUBJECTS_DIR env variable).')
parser.add_argument('--suffix', default='', help='Optional output file suffix.')
parser.add_argument('--temp-dir', help='Use alternative temporary directory. This will get deleted unless --debug is enabled.')
parser.add_argument('--out-dir', help='Use alternative output directory (only for cross-sectional). Default is the subject\'s `mri` directory.')
parser.add_argument('--debug', action='store_true', help='Write intermediate debugging outputs.')
parser.add_argument('--threads', type=int, default=1, help='Number of threads to use. Defaults to 1.')
args = parser.parse_args()

# Make sure freesurfer has been sourced
if not os.environ.get('FREESURFER_HOME'):
    sf.system.fatal('FREESURFER_HOME must be set!')

# Specify the maximum number of threads the GEMS code will use
subregions.set_thread_count(args.threads)

# Sanity check on process type
if args.cross and args.long_base:
    sf.system.fatal('Cannot specify arguments for both cross-sectional and longitudinal processing')
if not args.cross and not args.long_base:
    sf.system.fatal('Must specify inputs for either cross-sectional or longitudinal processing')

# Get the subjects directory
subjects_dir = args.sd if args.sd is not None else os.environ.get('SUBJECTS_DIR')
if subjects_dir is None:
    sf.system.fatal('Subjects directory must be set by --sd or the SUBJECTS_DIR env var')
    
# Input volumes types. This should be extended later on to accept multiple and variable inputs images
input_image_types = ['norm.mgz']
input_seg_type = 'aseg.mgz'

# Might need to run on each hemi
if args.structure == 'hippo-amygdala':
    sides = ['left', 'right']
else:
    sides = [None]

# Loop over sides of the brain (if necessary)
for side in sides:

    # Configure initial parameters
    parameters = {
        'debug': args.debug,
        'fileSuffix': args.suffix,
    }

    # Make sure to indicate hemi if necessary
    if side is not None:
        print(f'\nProcessing {side} hemisphere\n')
        parameters['side'] = side

    # Let's make a utility function to configure subject-specific parameters
    # to avoid redundancy with the longitudinal options
    def config_subj_params(subject, out_dir=None, temp_subdir=''):
        subjdir = os.path.join(subjects_dir, subject)
        subjParameters = parameters.copy()
        temp_dir = os.path.join(args.temp_dir, temp_subdir) if args.temp_dir else None
        subjParameters.update({
            'inputImageFileNames': [os.path.join(subjdir, 'mri', i) for i in input_image_types],
            'inputSegFileName': os.path.join(subjdir, 'mri', input_seg_type),
            'outDir': out_dir if out_dir else os.path.join(subjdir, 'mri'),
            'tempDir': temp_dir,
        })
        if args.structure == 'hippo-amygdala':
            # Provide some extra data for hippocampal subregions
            subjParameters['wmParcFileName'] = os.path.join(subjdir, 'mri', 'wmparc.mgz')
        return subjParameters

    if args.cross:
        # Cross-sectional analysis
        parameters = config_subj_params(args.cross, out_dir=args.out_dir)
        subregions.run_cross_sectional(args.structure, parameters)
    else:
        # Longitudinal analysis
        baseParameters = config_subj_params(args.long_base, temp_subdir=f'base')
        with open(os.path.join(subjects_dir, args.long_base, 'base-tps'), 'r') as file:
            long_tps = [f'{tp}.long.{args.long_base}' for tp in file.read().splitlines() if tp]
        print('Using longitudinal timepoints:', ' '.join(long_tps), '\n')
        tpParameters = [config_subj_params(subj, temp_subdir=f'tp{t:02d}') for t, subj in enumerate(long_tps)]
        subregions.run_longitudinal(args.structure, baseParameters, tpParameters)
