#!/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"]

# 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 handleArguments(args):
    """Split up arguments into input cameras and arguments."""
    # TODO: This code is buggy. A boolean option followed by a value will be taken
    # as the value of the boolean option but that one does not take a value.
    # As of now, all boolean options must either be followed by other options
    # or be at the end.
    cameras = []
    optionsList  = []

    # Loop through all entries.
    iterable = iter(range(0, len(args)))
    for i in iterable:
        a = args[i]
        if (i < len(args)-1): # Don't load the next value when we are at the end!
            n = args[i+1]
        else:
            n = '-' # This will just cause us to get out of the loop

        if asp_cmd_utils.isCmdOption(a):      # This is the start of an option.
            optionsList.append(a)  # Record this entry.

            # The next entry is the start of another option so this one has no values.
            if asp_cmd_utils.isCmdOption(n): 
                continue

            optionsList.append(n)  # Otherwise record the next entry has a value.
            next(iterable)        # Skip the next entry in the loop.

            if (a == '--crop-win'):  # These arguments have four values, not just one.
                # Add the additional three arguments and skip them in the loop.
                optionsList.append(args[i+2])           
                optionsList.append(args[i+3])
                optionsList.append(args[i+4])
                next(iterable)
                next(iterable)
                next(iterable)

        else: # This is one of the three positional arguments
            cameras.append(a)

    # Return the two lists
    return (cameras, optionsList)

def runSfs(options, cameras, outputFolder, outputName):
    """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 + cameras + extraArgs # Append other options
    willRun = True
    if options.resume:
        finalDem = tilePrefix + '-DEM-final.tif'
        verbose = True
        # See if it is a valid tif file
        ans = asp_system_utils.run_with_return_code(['gdalinfo',
                                                     finalDem], verbose) 
        if ans == 0:
            print("Will skip this tile, as file exists: " + finalDem)
            willRun = False
            
    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, inFilePrefix, outFilePrefix):

    # 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)
        outputDems.append(tilePrefix + '-' + inFilePrefix + '.tif')
         
    # Mosaic the outputs using dem_mosaic
    dem_mosaic_path = asp_system_utils.bin_path('dem_mosaic')
    dem_mosaic_args = ['--weights-exponent', '2', '--use-centerline-weights',
                       '-o', options.output_prefix]
    cmd = timeCmd + [dem_mosaic_path] + outputDems + dem_mosaic_args
    asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)

    # Rename from dem_mosaic convention to sfs convention
    finalDem = options.output_prefix + '-' + outFilePrefix + '.tif'
    print("Renaming to: " + finalDem)
    os.rename(options.output_prefix + '-tile-0.tif', finalDem)

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):

    cameras = []
    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 = 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("--num-processes",  dest="numProcesses", type=int, default=None,
                                          help="Number of processes to use per machine (the default program tries to choose best).")

    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('--threads',  dest='threads', default=1, type=int,
                      help='How many threads each process should use. The sfs executable is single-threaded in most of its execution, so a large number will not help here.')

    parser.add_argument("--resume", action="store_true", default=False,
                      dest="resume", help="Only run tiles for which the final DEM is missing or invalid.")

    parser.add_argument("--suppress-output", action="store_true", default=False,
                                            dest="suppressOutput",  help="Suppress output of sub-calls.")

    # 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)

    # Split up arguments into input cameras and the rest
    cameras, optionsList = handleArguments(args)

    # 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. " +
                      "Use the sfs tool directly if this is desired.\n" );

    if options.resume:
        if '--compute-exposures-only' in options.extraArgs or \
               '--estimate-slope-errors' in argsIn or \
               '--estimate-height-errors' not in argsIn:
            parser.print_help()
            parser.error("The --resume option cannot be used unless this " + \
                         "tool is invoked to produce a DEM.\n" )
        
    # Check the cameras.
    if len(cameras) < 1:
        parser.print_help()
        parser.error("Missing inputs.\n" );

    # Any additional arguments need to be forwarded to the sfs function
    options.extraArgs = optionsList


    # 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

    # 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, cameras, outputFolder, outputName) 

    # 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] + cameras + options.extraArgs + ['--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
    settings = asp_system_utils.run_and_parse_output( sfsPath, 
                                                      cameras + \
                                                      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
    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] + cameras + 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']

    # 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 + cameras + 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)

    if '--estimate-slope-errors' not in argsIn and '--estimate-height-errors' not in argsIn:
        # Mosaic the computed DEM
        mosaic_results(tileList, outputFolder, outputName, options,
                       'DEM-final', 'DEM-final')
        
        # Mosaic the albedo computed albedo
        if '--float-albedo' in argsIn:
            mosaic_results(tileList, outputFolder, outputName, options,
                           'comp-albedo-final', 'albedo-final')

    elif '--estimate-slope-errors' in argsIn:
        # Here we just mosaic the slope errors
        mosaic_results(tileList, outputFolder, outputName, options,
                       'slope-error', 'slope-error')
        
    elif '--estimate-height-errors' in argsIn:
        # Here we just mosaic the height errors
        mosaic_results(tileList, outputFolder, outputName, options,
                       'height-error', 'height-error')

    endTime = time.time()
    print("Finished in " + str(endTime - startTime) + " seconds.")

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