#!/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, baseFileNames):
    """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 baseFileName in baseFileNames:

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

        if not willRun:
            print("Will skip tile: " + tileName + ", as found valid: " + " ".join(baseFileNames))
        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, baseFile):

    # Final mosaicked result
    mosaicFile = options.output_prefix + '-' + baseFile
    
    # Create the list of DEMs to mosaic
    demList = mosaicFile
    demList = re.sub(r'\.\w+$', '', demList)
    demList += '-list.txt'
    with open(demList, 'w') as f:
        for tile in tileList:
            tileName = tile[4]
            tilePrefix = generateTilePrefix(outputFolder, tileName, outputName)
            demFile = tilePrefix + '-' + baseFile
            f.write(demFile + '\n')
                 
    # 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',
                       '--dem-list', demList, '-o', mosaicFile]
    cmd = timeCmd + [dem_mosaic_path] + 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 program.")
        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('--run-dir', dest='runDir', default=None,
      help = 'Directory in which the script is running.')
    
    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)

    # This can handle spaces in directory names, unlike when GNU parallel is in charge    
    if options.runDir is not None:
      os.chdir(options.runDir)
        
    # 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-haze'              in argsIn or \
        '--float-reflectance-model' in argsIn:
        print("Warning: Floating exposures and other quantities per tile may result " +
              "in non-globally consistent results.")
        
    # 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

    # If --save-covariances is on, also turn on --save-variances
    if '--save-covariances' in argsIn and '--save-variances' not in argsIn:
        argsIn.append('--save-variances')
        options.extraArgs.append('--save-variances')

    # See what kind of outputs to generate
    baseFileNames = []
    if '--estimate-height-errors' not in argsIn:
        # Computed DEM
        baseFileNames  += ['DEM-final.tif']
        
        # Albedo
        if '--float-albedo' in argsIn:
            baseFileNames  += ['albedo-final.tif']
             
        # Dem with nodata
        if '--save-dem-with-nodata' in argsIn:
            baseFileNames  += ['DEM-nodata-final.tif']
        
        # Variances and covariances. Can be up to 8 of them.
        if '--save-variances' in argsIn:
            prefixes = ['DEM']
            if '--float-albedo' in argsIn:
                prefixes += ['albedo']
            suffixes = ['-variance.tif']
            if '--save-covariances' in argsIn:
                suffixes += ['-left-covariance.tif', '-right-covariance.tif',
                             '-bottom-covariance.tif', '-top-covariance.tif']
            for prefix in prefixes:
                for suffix in suffixes:
                    baseFileNames += [prefix + suffix]
            
    elif '--estimate-height-errors' in argsIn:
        # Here we just mosaic the height errors
        baseFileNames  += ['height-error.tif']

    if options.resume:
        if '--compute-exposures-only' in options.extraArgs \
        or '--estimate-exposure-haze-albedo' in options.extraArgs:
            parser.print_help()
            parser.error("The --resume option cannot be used to compute the " + \
                         "exposures, haze, albedo.\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, baseFileNames) 

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

    # Compute the exposures, haze, initial albedo 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.
    estimatedExposureHazeAlbedo = False
    if '--image-exposures-prefix' not in options.extraArgs:
        print("Computing exposures (and haze and initial albedo, if applicable).\n")
        cmd = timeCmd + [sfsPath] + options.extraArgs
        # Add the option to compute the exposures, if not there already
        # TODO(oalexan1): Must handle the case when there is also haze
        # and albedo to compute. Must also test if only haze or if only
        # floating albedo is enabled.
        if '--compute-exposures-only' not in options.extraArgs:
            cmd += ['--compute-exposures-only']
        if '--estimate-exposure-haze-albedo' not in options.extraArgs:
            # Not strictly needed as it is equivalent to the above.
            cmd += ['--estimate-exposure-haze-albedo']
        asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)
        options.extraArgs += ['--image-exposures-prefix', options.output_prefix]
        estimatedExposureHazeAlbedo = True

    # If the haze prefix is not set, and have a non-zero number of haze coeffs,
    # must estimate the haze, unless done already
    if '--haze-prefix' not in options.extraArgs:
      # TODO(oalexan1): Must test!
      numHazeCoeffs = asp_cmd_utils.option_val(options.extraArgs, '--num-haze-coeffs')
      if numHazeCoeffs is not None and numHazeCoeffs != '0':
        if not estimatedExposureHazeAlbedo:
            print("Estimating haze.\n")
            cmd = timeCmd + [sfsPath] + options.extraArgs
            if '--estimate-exposure-haze-albedo' not in options.extraArgs:
              cmd += ['--estimate-exposure-haze-albedo']
            asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)
            estimatedExposureHazeAlbedo = True 
        # By now the haze should have been computed in either case
        options.extraArgs += ['--haze-prefix', options.output_prefix]
     
    # Same about estimating albedo, if we want to float albedo and was not provided
    if '--input-albedo' not in options.extraArgs and '--float-albedo' in options.extraArgs:
        if not estimatedExposureHazeAlbedo:
            print("Estimating albedo.\n")
            cmd = timeCmd + [sfsPath] + options.extraArgs
            if '--estimate-exposure-haze-albedo' not in options.extraArgs:
                cmd += ['--estimate-exposure-haze-albedo']
            asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)
            estimatedExposureHazeAlbedo = True
        # By now the albedo should have been computed in either case
        options.extraArgs += ['--input-albedo', options.output_prefix + "-albedo-estim.tif"]
        
    if '--compute-exposures-only' in options.extraArgs or \
       '--estimate-exposure-haze-albedo' in options.extraArgs:
        print("Finished computing exposure, haze, albedo.")
        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')
    
    # 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-separated
    # 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', '--env', 'ISISROOT']

    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 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')
        
    # Prepend the parallel arguments and append the extra arguments
    commandList = parallelArgs + commandList + options.extraArgs

    # Append the run directory
    runDir = asp_system_utils.escape_token(os.getcwd())
    commandList = commandList + ['--run-dir', runDir]

    # Use GNU parallel call to distribute the work across computers. Wait until
    # all processes are finished.
    verbose = True
    asp_system_utils.runInGnuParallel(options.numProcesses, 
                                      argumentFilePath, 
                                      commandList,
                                      options.nodesListPath, verbose)

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

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