#!/usr/bin/env python
# -*- coding: utf-8 -*-
# __BEGIN_LICENSE__
#  Copyright (c) 2009-2013, United States Government as represented by the
#  Administrator of the National Aeronautics and Space Administration. All
#  rights reserved.
#
#  The NGT platform is licensed under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance with the
#  License. You may obtain a copy of the License at
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
# __END_LICENSE__

'''
This tool implements a multi-process and multi-machine version of sfs. The input
DEM gets split into tiles with padding, sfs runs on each tile, and the outputs
are mosaicked.
'''

import sys
import os, glob, re, shutil, subprocess, string, time, errno, argparse, math

# The path to the ASP python files
basepath    = os.path.abspath(sys.path[0])
pythonpath  = os.path.abspath(basepath + '/../Python')  # for dev ASP
libexecpath = os.path.abspath(basepath + '/../libexec') # for packaged ASP
sys.path.insert(0, basepath) # prepend to Python path
sys.path.insert(0, pythonpath)
sys.path.insert(0, libexecpath)

from asp_alg_utils import *

import asp_file_utils, asp_system_utils, asp_cmd_utils, asp_image_utils, asp_string_utils
asp_system_utils.verify_python_version_is_supported()

# Prepend to system PATH
os.environ["PATH"] = libexecpath + os.pathsep + os.environ["PATH"]

# This is explained in asp_system_utils.py.
if 'ASP_LIBRARY_PATH' in os.environ:
    os.environ['LD_LIBRARY_PATH'] = os.environ['ASP_LIBRARY_PATH']

# Measure the memory usage on Linux and elapsed time
timeCmd = []
if 'linux' in sys.platform:
    timeCmd = ['/usr/bin/time', '-f', 'elapsed=%E memory=%M (kb)']

def generateTileList(sizeX, sizeY, tileSize, padding):
    """Generate a full list of tiles for this image"""

    Lx = genSegmentList(sizeX, tileSize, padding)
    Ly = genSegmentList(sizeY, tileSize, padding)

    tileList = []
    for x in range(0, len(Lx)-1):
        for y in range(0, len(Ly)-1):

            begX = Lx[x]; endX = Lx[x+1]
            begY = Ly[y]; endY = Ly[y+1]

            # apply the padding
            begX = begX - padding
            if (begX < 0): begX = 0
            begY = begY - padding
            if (begY < 0): begY = 0
            endX = endX + padding
            if (endX > sizeX): endX = sizeX
            endY = endY + padding
            if (endY > sizeY): endY = sizeY
            
            # Create a name for this tile
            # - Tile format is tile_col_row_width_height_.tif
            tileString = generateTileDir(begX, begY, endX, endY)
            tileList.append((begX, begY, endX, endY, tileString))
            
    return (len(Lx)-1, len(Ly)-1, tileList)

def generateTilePrefix(outputFolder, tileName, outputName):
    return os.path.join(outputFolder, tileName, outputName)

def runSfs(options, outputFolder, outputName, perTileFiles):
    """Run sfs in a single tile."""

    # Determine the name of the tile we need to write
    tileName = generateTileDir(options.pixelStartX, options.pixelStartY,
                               options.pixelStopX, options.pixelStopY)
    tilePrefix = generateTilePrefix(outputFolder, tileName, outputName)

    # Bounds for this tile
    startX = int(options.pixelStartX)
    stopX  = int(options.pixelStopX)
    startY = int(options.pixelStartY)
    stopY  = int(options.pixelStopY)

    extraArgs = []
    i = 0
    while i < len(options.extraArgs):
        arg = options.extraArgs[i]

        if arg == '--crop-win':
            # If the user passed in the crop-win argument, reconcile
            # it with the internal ROI we need.  Skip this element and
            # grab the user provided pixel bounds.
            userStartX = int(options.extraArgs[i+1])
            userStopX  = int(options.extraArgs[i+2])
            userStartY = int(options.extraArgs[i+3])
            userStopY  = int(options.extraArgs[i+4])

            # Reconcile the bounding box
            if (startX < userStartX): startX = userStartX
            if (stopX  > userStopX ): stopX  = userStopX
            if (startY < userStartY): startY = userStartY
            if (stopY  > userStopY ): stopY  = userStopY

            # If there is no tile left to generate then we can just return here
            if ((startX > stopX) or (startY > stopY)):
                return 0
            i += 5
        if arg == '-o' and i + 1 < len(options.extraArgs):
            # Replace the final output directory with the one for the given tile
            extraArgs.append(arg)
            extraArgs.append(tilePrefix)
            i += 2
        else:
            extraArgs.append(arg)
            i += 1

    # Call the command for a single tile
    cmd = timeCmd + ['sfs',  '--crop-win', str(startX), str(startY), str(stopX), str(stopY)]

    cmd = cmd + extraArgs # Append other options
    willRun = True
    
    if options.resume:
        # If all files that need to be computed exist, skip this tile
        willRun = False
        for perTileFile in perTileFiles:

            fullFilePath = tilePrefix + '-' + perTileFile
            if (not asp_system_utils.is_valid_image(fullFilePath)):
                willRun = True

        if not willRun:
            print("Will skip tile: " + tileName + ", as found valid: " + " ".join(perTileFiles))
        else:
            print("Will run tile: " + tileName)
            
    if willRun:
        (out, err, status) = asp_system_utils.executeCommand(cmd,
                                                             suppressOutput=options.suppressOutput)
        write_cmd_output(tilePrefix, cmd, out, err, status)
        
    return 0

def mosaic_results(tileList, outputFolder, outputName, options, inFile, outFile):

    # Create the list of final DEMs that get created at the end 
    outputDems = []
    for tile in tileList:
        tileName = tile[4]
        tilePrefix = generateTilePrefix(outputFolder, tileName, outputName)
        # TODO(oalexan1): At some point this command will become too big.
        # Need to keep these in a list.
        outputDems.append(tilePrefix + '-' + inFile)
         
    # Mosaic the outputs using dem_mosaic
    finalDem = options.output_prefix + '-' + outFile
    dem_mosaic_path = asp_system_utils.bin_path('dem_mosaic')
    dem_mosaic_args = ['--weights-exponent', '2', '--use-centerline-weights',
                       '-o', finalDem]
    cmd = timeCmd + [dem_mosaic_path] + outputDems + dem_mosaic_args
    asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)

def write_cmd_output(output_prefix, cmd, out, err, status):
    logFile = output_prefix + '-cmd-log.txt'
    print("Saving log in: " + logFile)
    with open(logFile, 'w') as f:
        f.write(" ".join(cmd))
        f.write(out)
        f.write(err)
        f.write('Command return status: ' + str(status))

def main(argsIn):

    sfsPath = asp_system_utils.bin_path('sfs')

    try:
        # Get the help text from the base C++ tool so we can append it to the python help
        cmd = [sfsPath,  '--help']
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
        baseHelp, err = p.communicate()
    except OSError:
        print("Error: Unable to find the required sfs tool!")
        return -1

    # Extract the version and help text
    vStart  = baseHelp.find('[ASP')
    vEnd    = baseHelp.find(']', vStart)+1
    version = baseHelp[vStart:vEnd]
    baseHelpText = "Help options for the underlying 'sfs' program:\n" + baseHelp[vEnd:]

    # Use parser that ignores unknown options
    usage  = "parallel_sfs -i <input DEM> -n <max iterations> -o <output prefix> <images> [other options]"

    parser = argparse.ArgumentParser(usage=usage, epilog=baseHelpText,
                                     formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('-i', '--input-dem',  dest='input_dem', default='',
                  help='The input DEM to refine using SfS.')

    parser.add_argument('-o', '--output-prefix',  dest='output_prefix', default='',
                  help='Prefix for output filenames.')

    parser.add_argument('--tile-size',  dest='tileSize', default=300, type=int,
                        help='Size of approximately square tiles to break up processing into (not counting the padding).')

    parser.add_argument('--padding',  dest='padding', default=50, type=int,
                                        help='How much to expand a tile in each direction. This helps with reducing artifacts in the final mosaicked SfS output.')

    parser.add_argument("--processes",  dest="numProcesses", type=int, default=None,
                        help="Number of processes to use on each node (the default is for the " + \
                        "program to choose).")

    parser.add_argument("--num-processes",  dest="numProcesses2", type=int, default=None,
                        help="Same as --processes. Used for backwards compatibility.")

    parser.add_argument('--nodes-list', dest='nodesListPath', default=None,
                                        help='A file containing the list of computing nodes, one per line. If not provided, run on the local machine.')

    parser.add_argument('--parallel-options', dest='parallel_options', 
                        default='--sshdelay 0.2',
                        help='Options to pass directly to GNU Parallel.')

    parser.add_argument('--threads',  dest='threads', default=8, type=int,
                        help='How many threads each process should use. This will be changed ' + \
                        'to 1 for ISIS cameras when --use-approx-camera-models is not set, ' + \
                        'as ISIS is single-threaded. Not all parts of the computation ' + \
                        'benefit from parallelization.')

    parser.add_argument("--resume", action="store_true", default=False, dest="resume",
                        help= "Resume a partially done run. Only process the tiles for which " + \
                        "the desired per-tile output files are missing or invalid (as "  + \
                        "checked by gdalinfo).")
    
    parser.add_argument("--suppress-output", action="store_true", default=False,
                        dest="suppressOutput",  help="Suppress output of sub-calls.")

    parser.add_argument('-v', '--version',        dest='version', default=False,
                        action='store_true', help='Display the version of software.')
    
    # Private options
    parser.add_argument('--pixelStartX', dest='pixelStartX', default=None, type=int,
                                        help=argparse.SUPPRESS)
    parser.add_argument('--pixelStartY', dest='pixelStartY', default=None, type=int,
                                        help=argparse.SUPPRESS)
    parser.add_argument('--pixelStopX',  dest='pixelStopX', default=None, type=int,
                                        help=argparse.SUPPRESS)
    parser.add_argument('--pixelStopY',  dest='pixelStopY', default=None, type=int,
                                        help=argparse.SUPPRESS)

    # This call handles all the parallel_sfs specific options.
    (options, args) = parser.parse_known_args(argsIn)

    if options.version:
        asp_system_utils.print_version_and_exit()

    if not args and not options.version:
        parser.print_help()
        sys.exit(1)
        
    # The parallel_sfs tool must not take the option --crop-win as it will
    # work on the entire input DEM.
    if '--crop-win' in argsIn:
        parser.print_help()
        parser.error("parallel_sfs cannot take the --crop-win option. " +
                      "Crop the input DEM using gdal_translate.\n" );

    if     '--float-exposure'          in argsIn or \
           '--float-cameras'           in argsIn or \
           '--float-all-cameras'       in argsIn or \
           '--float-haze'              in argsIn or \
           '--float-reflectance-model' in argsIn or \
           '--float-sun-position'      in argsIn:
        parser.print_help()
        parser.error("Cannot float exposures, cameras, or other per-tile quantities, except " +
                     "for the DEM heights or albedo in parallel_sfs, as different results " +
                     "will be obtained in different tiles. If desired to float these, do that " +
                     "on a clip where all images have good coverage, and using the sfs tool.\n");
        
    # Any additional arguments need to be forwarded to the sfs function
    options.extraArgs = args

    # Pass to the sfs executable the -i and -o options we filtered out
    options.extraArgs = ['-i', options.input_dem, '-o', options.output_prefix,
                         '--threads', str(options.threads) ] + options.extraArgs

    # See what kind of outputs to generate
    perTileFiles = []
    mosaickedFiles = []
    if ('--estimate-slope-errors' not in argsIn) and ('--estimate-height-errors' not in argsIn):
        # Computed DEM
        perTileFiles  += ['DEM-final.tif']
        mosaickedFiles += ['DEM-final.tif']
        
        # Albedo
        if '--float-albedo' in argsIn:
            perTileFiles  += ['comp-albedo-final.tif']
            mosaickedFiles += ['albedo-final.tif']
             
        # Dem with nodata
        if '--save-dem-with-nodata' in argsIn:
            perTileFiles  += ['DEM-nodata-final.tif']
            mosaickedFiles += ['DEM-nodata-final.tif']
            
    elif '--estimate-slope-errors' in argsIn:
        # Slope errors
        perTileFiles  += ['slope-error.tif']
        mosaickedFiles += ['slope-error.tif']
        
    elif '--estimate-height-errors' in argsIn:
        # Here we just mosaic the height errors
        perTileFiles  += ['height-error.tif']
        mosaickedFiles += ['height-error.tif']

    if options.resume:
        if '--compute-exposures-only' in options.extraArgs:
            parser.print_help()
            parser.error("The --resume option cannot be used to compute the exposures.\n")

    # Handle the situation that both --processes and --num-processes can happen
    if (options.numProcesses is not None) and (options.numProcesses2 is not None):
        raise Exception("Cannot set both --processes and --num-processes.")
    if options.numProcesses is None and (options.numProcesses2 is not None):
        # Copy over --num-processes to --processes
        options.numProcesses = options.numProcesses2

    # Set up output folder
    outputFolder = os.path.dirname(options.output_prefix)
    if outputFolder == '':
        outputFolder = './' # Handle calls in same directory
    outputName = os.path.basename(options.output_prefix)
    asp_file_utils.createFolder(outputFolder)

    startTime = time.time()

    # Determine if this is a main copy or a spawned copy
    spawnedCopy = ((options.pixelStartX is not None) and (options.pixelStartY is not None) and
                   (options.pixelStopX  is not None) and (options.pixelStopY  is not None) and
                   outputFolder)

    if spawnedCopy:
        # This copy was spawned to process a single tile. Just call a
        # function to handle this and then we are done.
        return runSfs(options, outputFolder, outputName, perTileFiles) 

    # Otherwise this is the original called process and there are multiple steps to go through

    # Compute the exposures based on the full DEM if not specified and
    # missing. Save them to disk. They will be read when SfS runs on
    # each small DEM tile.
    if '--image-exposures-prefix' not in options.extraArgs:
        print("Computing exposures.\n")
        cmd = timeCmd + [sfsPath] + options.extraArgs
        # Add the option to compute the exposures, if not there already
        if '--compute-exposures-only' not in options.extraArgs:
            cmd += ['--compute-exposures-only']
        asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)
        options.extraArgs += ['--image-exposures-prefix', options.output_prefix]

    if '--compute-exposures-only' in options.extraArgs:
        print("Finished computing exposures.")
        return
    
    # What is the size of the DEM on which to do SfS
    sep = ","
    verbose = False
    print("Running initial query")
    settings = asp_system_utils.run_and_parse_output(sfsPath, 
                                                     options.extraArgs + ['--query'],
                                                     sep, verbose)
    sizeX = int(settings['dem_cols'][0])
    sizeY = int(settings['dem_rows'][0])

    # Compute the number and size of tiles in x and y (width and height).
    # Make all tiles about the same size.
    # For now we just break up the image into a user-specified tile size
    print("Generating the list of tiles") # this can take a while so print a note
    numTilesX, numTilesY, tileList = generateTileList(sizeX, sizeY,
                                                      options.tileSize, options.padding)
    numTiles = numTilesX * numTilesY

    print('Splitting into ' + str(numTilesX) + ' by ' + str(numTilesY) + \
          ' = ' + str(numTilesX*numTilesY) + ' tiles.\n')
    
    # If there is only going to be one output tile, just use the non-parallel call
    if (numTilesX*numTilesY == 1):
        cmd = timeCmd + [sfsPath] + options.extraArgs
        (out, err, status) = asp_system_utils.executeCommand(cmd,
                                                             suppressOutput=options.suppressOutput)
        write_cmd_output(options.output_prefix, cmd, out, err, status)
        return 0
    
    # Generate a text file that contains the boundaries for each tile
    argumentFilePath = os.path.join(outputFolder, 'argumentList.txt')
    argumentFile     = open(argumentFilePath, 'w')
    for tile in tileList:
        argumentFile.write( str(tile[0]) + '\t' + str(tile[1]) + '\t' \
                            + str(tile[2]) + '\t' + str(tile[3]) + '\n')
    argumentFile.close()

    # Indicate to GNU Parallel that there are multiple tab-seperated
    # variables in the text file we just wrote
    parallelArgs = ['--colsep', "\\t", '--will-cite', '--env', 'ASP_DEPS_DIR',
                    '--env', 'PATH', '--env', 'LD_LIBRARY_PATH',
                    '--env', 'ASP_LIBRARY_PATH']

    if options.parallel_options is not None:
        parallelArgs += options.parallel_options.split(' ')

    # Get the number of available nodes and CPUs per node
    numNodes = asp_system_utils.getNumNodesInList(options.nodesListPath)

    # We assume all machines have the same number of CPUs (cores)
    cpusPerNode = asp_system_utils.get_num_cpus()

    # TODO: What is a good number here?
    processesPerCpu = 2

    # Set the optimal number of processes if the user did not specify
    if not options.numProcesses:
        options.numProcesses = cpusPerNode * processesPerCpu

    # Note: sfs can run with multiple threads on non-ISIS data but we don't use that
    #       functionality here since we call sfs with one tile at a time.

    # No need for more processes than their are tiles!
    if options.numProcesses > numTiles:
        options.numProcesses = numTiles

    # Build the command line that will be passed to GNU parallel
    # - The numbers in braces will receive the values from the text file we wrote earlier
    # - The output path used here does not matter since spawned copies compute the correct tile path.
    python_path = sys.executable # children must use same Python as parent
    # We use below the libexec_path to call python, not the shell script
    parallel_sfs_path = asp_system_utils.libexec_path('parallel_sfs')
    commandList   = [python_path, parallel_sfs_path,
                     '--pixelStartX', '{1}',
                     '--pixelStartY', '{2}',
                     '--pixelStopX',  '{3}',
                     '--pixelStopY',  '{4}',
                     '--threads', str(options.threads)
                     ]
    if options.suppressOutput:
        commandList = commandList + ['--suppress-output']

    if options.resume:
        commandList.append('--resume')
        
    commandList   = commandList + options.extraArgs # Append other options
    commandString = asp_string_utils.argListToString(commandList)

    # Use GNU parallel call to distribute the work across computers
    # - This call will wait until all processes are finished
    asp_system_utils.runInGnuParallel(options.numProcesses, commandString,
                                      argumentFilePath, parallelArgs,
                                      options.nodesListPath, True)#not options.suppressOutput)

    # Mosaic the results
    for it in range(len(perTileFiles)):
        mosaic_results(tileList, outputFolder, outputName, options,
                       perTileFiles[it], mosaickedFiles[it])
        
    endTime = time.time()
    print("Finished in " + str(endTime - startTime) + " seconds.")

if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
