#!/usr/bin/env python

import os
import sys
import json
import scipy.io
import argparse
import surfa as sf
import freesurfer.samseg as samseg


# ------ Parse Command Line Arguments ------

parser = argparse.ArgumentParser()

default_threads = int(os.environ.get('OMP_NUM_THREADS', 1))

# Custom action to extend arguments into a list so that both of these are allowed:
# command --input 1 2 3
# command --input 1 --input 2 --input 3
class ExtendAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        items = getattr(namespace, self.dest) or []
        items.extend(values)
        setattr(namespace, self.dest, items)
parser.register('action', 'extend', ExtendAction)

# required
parser.add_argument('-o', '--output', metavar='DIR', dest='outputDirectory', help='Output directory.', required=True)
parser.add_argument('-i', '--input', nargs='+', action='extend', metavar='FILE', dest='inputFileNames', help='Input image(s).', required=True)
# optional processing options
parser.add_argument('-m', '--mode', nargs='+', help='Output basenames for the input image mode.')
parser.add_argument('--threads', type=int, default=default_threads, help='Number of threads to use. Defaults to current OMP_NUM_THREADS or 1.')
parser.add_argument('--reg-only', action='store_true', default=False, help='Only perform initial affine registration.')
parser.add_argument('-r', '--reg', metavar='FILE', help='Skip initial affine registration and read transform from file.')
parser.add_argument('--init-reg', metavar='FILE', help='Initial affine registration.')
parser.add_argument('-a', '--atlas', metavar='DIR', help='Point to an alternative atlas directory.')
parser.add_argument('--gmm', metavar='FILE', help='Point to an alternative GMM file.')
parser.add_argument('--ignore-unknown', action='store_true', help='Ignore final priors corresponding to unknown class.')
parser.add_argument('--options', metavar='FILE', help='Override advanced options via a json file.')
parser.add_argument('--pallidum-separate', action='store_true', default=False, help='Move pallidum outside of global white matter class. Use this flag when T2/flair is used.')
parser.add_argument('--mesh-stiffness', type=float, help='Override mesh stiffness.')
parser.add_argument('--smooth-wm-cortex-priors', type=float, help='Sigma value to smooth the WM and cortex atlas priors.')
parser.add_argument('--bias-field-smoothing-kernel', type=float, help='Distance in mm of sinc function center to first zero crossing.')
# optional lesion options
parser.add_argument('--lesion', action='store_true', default=False, help='Enable lesion segmentation (requires tensorflow).')
parser.add_argument('--threshold', type=float, default=0.3, help='Lesion threshold for final segmentation. Lesion segmentation must be enabled.')
parser.add_argument('--samples', type=int, default=50, help='Number of samples for lesion segmentation. Lesion segmentation must be enabled.')
parser.add_argument('--burnin', type=int, default=50, help='Number of burn-in samples for lesion segmentation. Lesion segmentation must be enabled.')
parser.add_argument('--lesion-pseudo-samples', nargs=2, type=float, default=[500, 500], help='Lesion pseudo samples mean and variance.')
parser.add_argument('--lesion-rho', type=float, default=50, help='Lesion ratio.')
parser.add_argument('--lesion-mask-structure', default='Cortex', help='Intensity mask brain structure. Lesion segmentation must be enabled.')
parser.add_argument('--lesion-mask-pattern', type=int, nargs='+', help='Lesion mask list (set value for each input volume): -1 below lesion mask structure mean, +1 above, 0 no mask. Lesion segmentation must be enabled.')
parser.add_argument('--random-seed', type=int, default=12345, help='Random seed.')
# optional options for segmenting 3D reconstructions of photo volumes
parser.add_argument('--dissection-photo', default=None, help='Use this flag for 3D reconstructed photos, and specify hemispheres that are present in the volumes: left, right, or both')

# optional debugging options
parser.add_argument('--history', action='store_true', default=False, help='Save history.')
parser.add_argument('--save-posteriors', nargs='*', help='Save posterior volumes to the "posteriors" subdirectory.')
parser.add_argument('--save-probabilities', action='store_true', help='Save final modal class probabilities to "probabilities" subdirectory.')
parser.add_argument('--showfigs', action='store_true', default=False, help='Show figures during run.')
parser.add_argument('--save-mesh', action='store_true', help='Save the final mesh in template space.')
parser.add_argument('--save-warp', action='store_true', help='Save the image->template warp field.')
parser.add_argument('--movie', action='store_true', default=False, help='Show history as arrow key controlled time sequence.')

args = parser.parse_args()

# ------ Initial Setup ------

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

# Start the process timer
timer = samseg.Timer()

# Create the output folder
os.makedirs(args.outputDirectory, exist_ok=True)

# Specify the maximum number of threads the GEMS code will use
samseg.setGlobalDefaultNumberOfThreads(args.threads)

# Remove previous cost log
costfile = os.path.join(args.outputDirectory, 'cost.txt')
if os.path.exists( costfile ):
    os.remove( costfile ) 

# Get the atlas directory
atlasDir = os.environ.get('SAMSEG_DATA_DIR')
if args.atlas:
    atlasDir = args.atlas
if not atlasDir:
    # Altas defaults
    if args.lesion:
        defaultAtlas = '20Subjects_smoothing2_down2_smoothingForAffine2_lesion'
    else:
        defaultAtlas = '20Subjects_smoothing2_down2_smoothingForAffine2'
    atlasDir = os.path.join(os.environ.get('FREESURFER_HOME'), 'average', 'samseg', defaultAtlas)

# Setup the visualization tool
visualizer = samseg.initVisualizer(args.showfigs, args.movie)

# ------ Prepare Samseg Parameters ------

# Load user options from a JSON file
userModelSpecifications = {}
userOptimizationOptions = {}
if args.options:
    with open(args.options) as f:
        userOptions = json.load(f)
    if userOptions.get('modelSpecifications') is not None: 
        userModelSpecifications = userOptions.get('modelSpecifications')
    if userOptions.get('optimizationOptions') is not None: 
        userOptimizationOptions = userOptions.get('optimizationOptions')

# Check if --save-posteriors was specified without any structure search string
if args.save_posteriors is not None and len(args.save_posteriors) == 0:
    savePosteriors = True
else:
    savePosteriors = args.save_posteriors

if args.smooth_wm_cortex_priors is not None:
    userModelSpecifications['whiteMatterAndCortexSmoothingSigma'] = args.smooth_wm_cortex_priors

if args.bias_field_smoothing_kernel is not None:
    userModelSpecifications['biasFieldSmoothingKernelSize'] = args.bias_field_smoothing_kernel

if args.mesh_stiffness is not None:
    userModelSpecifications['K'] = args.mesh_stiffness
    print("Mesh stiffness set to %g" % userModelSpecifications['K'])

# ------ Run Samseg ------

# If we are dealing with photos, we  skip rescaling of intensities and also force  ignoreUnknownPriors=True
if args.dissection_photo is None:
    intensityWM = 110
    ignoreUnknownPriors = args.ignore_unknown
else:
    intensityWM = None
    ignoreUnknownPriors = True

samseg_kwargs = dict(
    imageFileNames=args.inputFileNames,
    atlasDir=atlasDir,
    savePath=args.outputDirectory,
    userModelSpecifications=userModelSpecifications,
    userOptimizationOptions=userOptimizationOptions,
    targetIntensity=intensityWM,
    targetSearchStrings=[ 'Cerebral-White-Matter' ],
    visualizer=visualizer,
    saveHistory=args.history,
    saveMesh=args.save_mesh,
    savePosteriors=savePosteriors,
    saveWarp=args.save_warp,
    modeNames=args.mode,
    pallidumAsWM=(not args.pallidum_separate),
    saveModelProbabilities=args.save_probabilities,
    gmmFileName=args.gmm,
    ignoreUnknownPriors=ignoreUnknownPriors
)

if args.lesion:

    # If lesion mask pattern is not specified, assume inputs are T1-contrast
    lesion_mask_pattern = args.lesion_mask_pattern
    if lesion_mask_pattern is None:
        lesion_mask_pattern = [0] * len(args.inputFileNames)
        print('Defaulting lesion mask pattern to %s' % str(lesion_mask_pattern))

    # Delay import until here so that tensorflow doesn't get loaded unless needed
    from freesurfer.samseg.SamsegLesion import SamsegLesion
    samseg = SamsegLesion(**samseg_kwargs,
        numberOfPseudoSamplesMean=args.lesion_pseudo_samples[0],
        numberOfPseudoSamplesVariance=args.lesion_pseudo_samples[1],
        rho=args.lesion_rho,
        intensityMaskingSearchString=args.lesion_mask_structure,
        intensityMaskingPattern=lesion_mask_pattern,
        numberOfBurnInSteps=args.burnin,
        numberOfSamplingSteps=args.samples,
        threshold=args.threshold,
        randomSeed=args.random_seed
    )

else:
    samseg = samseg.Samseg(**samseg_kwargs, 
                           dissectionPhoto=args.dissection_photo,
                           nthreads=args.threads)

_, _, _, optimizationSummary = samseg.segment(
    costfile=costfile,
    timer=timer,
    transformFile=args.reg,
    initTransformFile=args.init_reg,
    reg_only=args.reg_only
)

# Save a summary of the optimization process
with open(costfile, 'a') as file:
    for multiResolutionLevel, item in enumerate(optimizationSummary):
        file.write('atlasRegistrationLevel%d %d %f\n' % (multiResolutionLevel, item['numberOfIterations'], item['perVoxelCost']))

timer.mark('run_samseg complete')
