#!/usr/bin/env python3
#
# (c) 2017 Fetal-Neonatal Neuroimaging & Developmental Science Center
#                   Boston Children's Hospital
#
#              http://childrenshospital.org/FNNDSC/
#                        dev@babyMRI.org
#

import sys, os
sys.path.insert(1, os.path.join(os.path.dirname(__file__), '../mgz2imagetree'))

import mgz2imagetree
from    argparse            import RawTextHelpFormatter
from    argparse            import ArgumentParser
import  pudb

import  pfmisc
from    pfmisc._colors      import Colors
from    pfmisc              import other

str_version = "1.1.2"
str_desc = Colors.CYAN + """

                      _____ _                            _                 
                     / __  (_)                          | |                
 _ __ ___   __ _ ____`' / /'_ _ __ ___   __ _  __ _  ___| |_ _ __ ___  ___ 
| '_ ` _ \ / _` |_  /  / / | | '_ ` _ \ / _` |/ _` |/ _ \ __| '__/ _ \/ _ \\
| | | | | | (_| |/ / ./ /__| | | | | | | (_| | (_| |  __/ |_| | |  __/  __/
|_| |_| |_|\__, /___|\_____/_|_| |_| |_|\__,_|\__, |\___|\__|_|  \___|\___|
            __/ |                              __/ |                       
           |___/                              |___/                        
  
                        Path-File MGZ labels filter

        Recursively walk down a directory tree and extract MGZ labels,
        creating image slices for each individual label preserving directory structure in output tree.

                             -- version """ + \
             Colors.YELLOW + str_version + Colors.CYAN + """ --

        'mgz2imagetree' is a customizable and friendly MGZ labels filter. 
         Input trees are reconstructed in an output
        directory, preserving directory structure. Each node tree contains
        label directories containing image slices (in png/jpeg format) on the corresponding output directory for each input subject.


""" + Colors.NO_COLOUR

def synopsis(ab_shortOnly = False):
    scriptName = os.path.basename(sys.argv[0])
    shortSynopsis =  '''
    NAME

	    mgz2imagetree 
        - process mgz files and create labelled directories in a similar tree structure in the output directory.

    SYNOPSIS

            mgz2imagetree                                                  \\
                    [-I|--inputDir <inputDir>]                                  \\
                    [-i|--inputFile <inputFile>]                                \\
                    [-O|--outputDir <outputDir>]                                \\
                    [-o|--outputFileStem <outputFileStem>]                      \\
                    [--outputLeafDir <outputLeafDirFormat>]                     \\
                    [-t|--outputFileType <outputFileType>]                      \\
                    [--feature <MGZFileToConvertToLabelledSegments>]            \\        
                    [--image <MGZFileImageFile>]                                \\
                    [--label <prefixForLabelDirectories>]                       \\
                    [-n|--normalize]                                            \\
                    [-l|--lookuptable <LUTfile>]                                \\
                    [-s|--skipLabelValueList <ListOfLabelNumbersToSkip>]        \\
                    [-f|--filterLabelValueList <ListOfVoxelValuesToInclude>]    \\
                    [-w|--wholeVolume <wholeVolDirName>]                        \\
                    [-v <verbosity>]                                            \\
                    [--version]                                                 \\
                    [-x|--man]                                                  \\
                    [-y|--synopsis]

    '''

    description =  '''
    DESCRIPTION

        `mgz2imagetree` filters the labels in mgz files for each subject within 
        the inputdir tree structure and copies the resultant label directories containing
        the image slices in png/jpg format into the outputdir with a similar structure as inputdir.
         
        The script accepts an <inputDir>, and then from this point an os.walk
        is performed to extract all the subdirs. Each subdir is examined for
        MGZ files (in the simplest sense by a file extension mapping).

        An image conversion is performed on the mgz files using the library "mgz2imgslices"
        which creates individual label direcctories containing 256 png/jpg slices and 
        the whole volume too with all label values. The results are stored in a tree structure 
        similar to the inputdir.

    ARGS

        [-I|--inputDir <inputDir>]
        Input directory to examine. By default, the first file in this
        directory is examined for its tag information.

        [-i|--inputFile <inputFile>]
        An optional <inputFile> specified relative to the <inputDir>. If 
        specified, then do not perform a directory walk, but convert only 
        this file.

        [-O|--outputDir <outputDir>]
        The output root directory that will contain a tree structure identical
        to the input directory, and each "leaf" node will contain the analysis
        results.

        [--outputLeafDir <outputLeafDirFormat>]
        If specified, will apply the <outputLeafDirFormat> to the output
        directories containing data. This is useful to blanket describe
        final output directories with some descriptive text, such as 
        'anon' or 'preview'. 

        This is a formatting spec, so 

            --outputLeafDir 'preview-%%s'

        where %%s is the original leaf directory node, will prefix each
        final directory containing output with the text 'preview-' which
        can be useful in describing some features of the output set.

        -o|--outputFileStem <outputFileStem>
        The output file stem to store data. This should *not* have a file
        extension, or rather, any "." in the name are considered part of 
        the stem and are *not* considered extensions.

        [-t|--outputFileType <outputFileType>]
        The output file type. If different to <outputFileStem> extension,
        will override extension in favour of <outputFileType>.
        Should be a ``png`` or ``jpg``.

        [--feature <MGZFileToConvertToLabelledSegments>]
        The feature file containing the cortical strip which needs to be filtered 
        into its constituent cortical segments using ``mgz2imgslices``
        
        [--image <MGZFileImageFile>]
        The raw 3D mgz image file that needs to be split and stored
        as slices in the corresponding subject's output directory.

        [--label <prefixForLabelDirectories>]
        Prefixes the string <prefixForLabelDirectories> to each filtered 
        directory name. This is mostly for possible downstream processing, 
        allowing a subsequent operation to easily determine which of the output
        directories correspond to labels.

        [-n|--normalize]
        If specified, will normalize the output image pixel values to
        0 and 1, otherwise pixel image values will retain the value in
        the original input volume.

        [-l|--lookuptable <LUTfile>]
        If passed, perform a lookup on the filtered voxel label values
        according to the contents of the <LUTfile>. This <LUTfile> should
        conform to the FreeSurfer lookup table format (documented elsewhere).
        Note that the special <LUTfile> string ``__val__`` can be passed which
        effectively means "no <LUTfile>". In this case, the numerical voxel
        values are used for output directory names. This special string is
        really only useful for scripted cases of running this application when
        modifying the CLI is more complex than simply setting the <LUTfile> to
        ``__val__``.

        [-s|--skipLabelValueList <ListOfLabelNumbersToSkip>]
        If specified as a comma separated string of label numbers,
        will not create directories of those label numbers.
        
        [-f|--filterLabelValueList <ListOfVoxelValuesToInclude>]
        The logical inverse of the [skipLabelValueList] flag. If specified,
        only filter the comma separated list of passed voxel values from the
        input volume.
        The default value of "-1" implies all voxel values should be filtered.

        [-w|--wholeVolume <wholeVolDirName>]
        If specified, creates a diretory called <wholeVolDirName> (within the
        outputdir) containing PNG/JPG images files of the entire input.
        This effectively really creates a PNG/JPG conversion of the input
        mgz file.
        Values in the image files will be the same as the original voxel
        values in the ``mgz``, unless the [--normalize] flag is specified
        in which case this creates a single-value mask of the input image.

        [-r|--rawDirName]
        The name that the user wants to use to store the slices of the 
        raw mgz image file
        
        [-x|--man]
        Show full help.

        [-y|--synopsis]
        Show brief help.

        [--version]
        If specified, print the version number and exit.

        [-v|--verbosity <level>]
        Set the app verbosity level. 

            0: No internal output;
            1: Run start / stop output notification;
            2: As with level '1' but with simpleProgress bar in 'pftree';
            3: As with level '2' but with list of input dirs/files in 'pftree';
            5: As with level '3' but with explicit file logging for
                    - read
                    - analyze
                    - write

    EXAMPLES:

    #WIP
    '''
    if ab_shortOnly:
        return shortSynopsis
    else:
        return shortSynopsis + description


parser  = ArgumentParser(description = str_desc, formatter_class = RawTextHelpFormatter)

parser.add_argument("-I", "--inputDir",
                    help    = "input dir",
                    dest    = 'inputDir')
parser.add_argument("-i", "--inputFile",
                    help    = "input file",
                    dest    = 'inputFile',
                    default = '')
parser.add_argument("-o", "--outputFileStem",
                    help    = "output file stem",
                    default = "",
                    dest    = 'outputFileStem')
parser.add_argument("-O", "--outputDir",
                    help    = "output image directory",
                    dest    = 'outputDir',
                    default = '')
parser.add_argument("-t", "--outputFileType",
                    help    = "list of output report types",
                    dest    = 'outputFileType',
                    default = 'raw')
parser.add_argument("--feature",
                    help    = "name of mgz file with features (aparc+aseg)",
                    dest    = 'feature',
                    default = 'aparc+aseg.mgz')
parser.add_argument("--image",
                    help    = "name of entire image file (brain.mgz)",
                    dest    = 'image',
                    default = 'brain.mgz')
parser.add_argument('--label',
                    help='prefix a label to all the label directories',
                    dest='label',
                    default = 'label'
                    )
parser.add_argument('-n', '--normalize',
                    help='normalize the pixels of output image files',
                    dest='normalize',
                    default = False
                    )
parser.add_argument('-l', '--lookuptable',
                    help='file contain text string lookups for voxel values',
                    dest='lookuptable',
                    default = '__val__'
                    )
parser.add_argument('-s', '--skipLabelValueList',
                    help='Comma separated list of voxel values to skip',
                    dest='skipLabelValueList',
                    default = ''
                    )
parser.add_argument('-f', '--filterLabelValueList',
                    help='Comma separated list of voxel values to include',
                    dest='filterLabelValueList',
                    default = "-1"
                    )
parser.add_argument('-w', '--wholeVolume',
                    help='Converts entire mgz volume to png/jpg instead of individually masked labels',
                    dest='wholeVolume',
                    default = 'wholeVolume'
                    )
parser.add_argument('-r', '--rawDirName',
                    help='name of directory that will contain raw mgz image slices',
                    dest='rawDirName',
                    default = 'rawDirName'
                    )
parser.add_argument("-x", "--man",
                    help    = "man",
                    dest    = 'man',
                    action  = 'store_true',
                    default = False)
parser.add_argument("-y", "--synopsis",
                    help    = "short synopsis",
                    dest    = 'synopsis',
                    action  = 'store_true',
                    default = False)
parser.add_argument("--outputLeafDir",
                    help    = "formatting spec for output leaf directory",
                    dest    = 'outputLeafDir',
                    default = "")
parser.add_argument("--printElapsedTime",
                    help    = "print program run time",
                    dest    = 'printElapsedTime',
                    action  = 'store_true',
                    default = False)
parser.add_argument("-v", "--verbosity",
                    help    = "verbosity level for app",
                    dest    = 'verbosity',
                    default = "0")
parser.add_argument('--version',
                    help    = 'if specified, print version number',
                    dest    = 'b_version',
                    action  = 'store_true',
                    default = False)

args = parser.parse_args()

if args.man or args.synopsis:
    print(str_desc)
    if args.man:
        str_help     = synopsis(False)
    else:
        str_help     = synopsis(True)
    print(str_help)
    sys.exit(1)

if args.b_version:
    print("Version: %s" % str_version) 
    sys.exit(1)

# pudb.set_trace()
mgz2imagetree = mgz2imagetree.mgz2imagetree(
                        inputDir            = args.inputDir,
                        inputFile           = args.inputFile,
                        outputDir           = args.outputDir,
                        outputFileStem      = args.outputFileStem,
                        outputLeafDir       = args.outputLeafDir,
                        outputFileType      = args.outputFileType,
                        feature             = args.feature,
                        image               = args.image,
                        label               = args.label,
                        normalize           = args.normalize,
                        lookuptable         = args.lookuptable,
                        skipLabelValueList  = args.skipLabelValueList,
                        filterLabelValueList= args.filterLabelValueList,
                        wholeVolume         = args.wholeVolume,
                        rawDirName          = args.rawDirName,
                        man                 = args.man,
                        synopsis            = args.synopsis,
                        verbosity           = args.verbosity,
                        version             = str_version   
                        )

# And now run it!
d_mgz2imagetree = mgz2imagetree.run(timerStart = True)

if args.printElapsedTime: 
    mgz2imagetree.dp.qprint(
                                "Elapsed time = %f seconds" % 
                                d_mgz2imagetree['runTime']
                            )

sys.exit(0)
