#!/usr/bin/env python
# __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 version of mapproject to
greatly increase the speed of map projecting ISIS images.
'''

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. First try to find things in libexec,
# and if that fails, then in bin. Otherwise mapproject_single
# cannot be found.
os.environ["PATH"] = libexecpath + os.pathsep + \
                     basepath + os.pathsep    + \
                     os.environ["PATH"]

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

def generateTileList(fullWidth, fullHeight, tileSize):
    """Generate a full list of tiles for this image"""

    numTilesX = int(math.ceil(fullWidth  / float(tileSize)))
    numTilesY = int(math.ceil(fullHeight / float(tileSize)))

    tileList = []
    for r in range(0, numTilesY):
        for c in range(0, numTilesX):

            # Starting pixel positions for the tile
            tileStartY = r * tileSize
            tileStartX = c * tileSize

            # Determine the size of this tile
            thisWidth  = tileSize
            thisHeight = tileSize
            if (r == numTilesY-1): # If the last row
                thisHeight = fullHeight - tileStartY # Height is last remaining pixels
            if (c == numTilesX-1): # If the last col
                thisWidth  = fullWidth  - tileStartX # Width is last remaining pixels

            # Get the end pixels for this tile
            tileStopY  = tileStartY + thisHeight # Stop values are exclusive
            tileStopX  = tileStartX + thisWidth

            # Create a name for this tile
            # - Tile format is tile_col_row_width_height_.tif
            tileString = generateTileName(tileStartX, tileStartY, tileStopX, tileStopY)

            tileList.append((tileStartX, tileStartY, tileStopX, tileStopY, tileString))

    return (numTilesX, numTilesY, tileList)

def writeSingleTile(options):
    """Writes a single tile according to the options"""

    # Determine the name of the tile we need to write
    tileName = generateTileName(options.pixelStartX, options.pixelStartY,
                                options.pixelStopX, options.pixelStopY)
    tilePath = os.path.join(options.workDir, tileName)

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

    # If the user passed in the t_pixelwin argument, reconcile it with the internal ROI we need.
    extraArgs = []
    i = 0
    while i < len(options.extraArgs):
        arg = options.extraArgs[i]

        if arg == '--t_pixelwin':
            # Skip this element and grab the user provided pixel bounds
            userStartX = int(options.extraArgs[i+1])
            userStartY = int(options.extraArgs[i+2])
            userStopX  = 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
        else:
            extraArgs.append(arg)
            i += 1

    # Just call the command for a single tile.
    cmd = ['mapproject_single', '--t_pixelwin', 
           str(startX), str(startY), str(stopX), str(stopY),
           options.demPath, options.imagePath, options.cameraPath, tilePath]
    if options.noGeoHeaderInfo:
        cmd += ['--no-geoheader-info']
    cmd = cmd + extraArgs # Append other options
    asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput,
                                    realTimeOutput = True)

    if options.convertTiles: # Make uint8 version of the tile for debugging
        tilePathU8 = os.path.splitext(tilePath)[0] + 'U8.tif'
        cmd = ['gdal_translate', '-ot', 'byte', '-scale', tilePath, tilePathU8]
        asp_system_utils.executeCommand(cmd, suppressOutput=options.suppressOutput)

    return 0

def parse_gdal_metadata(filename, meta_label):
    """
    Seek out the metadata for the given label in a tif file and return
    that metadata as a map having key and value pairs. Use gdalinfo
    to query the tif file.
    """
    args = [filename]
    sep = ""
    verbose = False
    gdal_settings = asp_system_utils.run_and_parse_output("gdalinfo", args, sep,
                                                          verbose, return_full_lines = True)
    meta_dict = {}
    parseBlock = False
    
    for v in sorted(gdal_settings.keys()):

        if gdal_settings[v] == meta_label + ':':
            # Just reached the line with the label, so need to parse
            # the block of next several lines which start with a
            # space. Notice how above there is a column at the end of
            # the label line.
            parseBlock = True 
            continue
        
        if len(gdal_settings[v]) == 0 or gdal_settings[v][0] != " ":
            # Another block started
            if parseBlock:
                # Past the current block, time to stop
                break
            else:
                # Did not reach the desired block yet
                continue

        if len(gdal_settings[v]) > 0 and parseBlock:
            # In the desired block
            m = re.match(r"^\s*(.*?)\s*=\s*(.*?)\s*$", gdal_settings[v])
            if not m:
                continue
            meta_dict[m.group(1)] = m.group(2)

    return meta_dict

def format_metadata(meta_label, meta_dict):
    """
    Format the metadata with given label and entries as the GDAL VRT
    format expects it.
    """

    # From meta_label = "RPC Metadata" extract "RPC".
    vrt_domain = ""
    m = re.match(r"^([\w]+)\s+Metadata", meta_label)
    if m:
        vrt_domain = m.group(1)

    # The VRT format is quite particular
    # https://gdal.org/drivers/raster/vrt.html
    if vrt_domain == "":
        meta_lines = ["  <Metadata>\n"]
    else:
        meta_lines = ["  <Metadata domain=\"" + vrt_domain + "\">\n"]
        
    for v in sorted(meta_dict.keys()):
        line = "    <MDI key=\"" + v + "\">" + meta_dict[v] + "</MDI>\n"
        meta_lines.append(line)
    meta_lines.append("  </Metadata>\n")

    return meta_lines
    
#------------------------------------------------------------------------------

# If there is an input _RPC.TXT or .RPB, create such an output file as well,
# to store there the RPC coefficients. Need this for stereo later on. 
# For some reason, this function must be invoked after the writing
# of the output tif file has been competed. Otherwise something wipes it.
def maybe_copy_rpc(inputPath, outputPath):
    for ext in ['_RPC.TXT', '_rpc.txt', '.RPB', '.rpb']:
        input_rpc  = os.path.splitext(inputPath )[0] + ext
        output_rpc = os.path.splitext(outputPath)[0] + ext
        if os.path.exists(input_rpc):
            print("Copied " + input_rpc + " to " + output_rpc)
            shutil.copy(input_rpc, output_rpc)

def main(argsIn):

    relOutputPath = ""
    try:
        # Get the help text from the base C++ tool so we can append it to the python help
        cmd = ['mapproject_single',  '--help']
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
        baseHelp, err = p.communicate()
    except OSError:
        print("Error: Unable to find required mapproject_single 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 'mapproject_single' program:\n" + \
                   baseHelp[vEnd:]

    # Use a parser that ignores unknown options. Those will be passed to mapproject_single.
    usage  = "mapproject [options] <dem> <camera-image> <camera-model> <output-image>\n" + \
             "Instead of the DEM file, a datum can be provided, such as\n" + \
             "WGS84, NAD83, NAD27, D_MOON, D_MARS, and MOLA."
    parser = argparse.ArgumentParser(usage=usage, epilog=baseHelpText,
                                     formatter_class=argparse.RawTextHelpFormatter)

    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='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('--tile-size',  dest='tileSize', default=None, type=int,
      help = 'Size of square tiles to break up processing up '   + \
      'into. Each tile is run by an individual process. The '    + \
      'default is 1024 pixels for ISIS cameras, as then each '   + \
      'process is single-threaded, and 5120 for other cameras, ' + \
      'as such a process is multi-threaded, and disk I/O '       + \
      'becomes a bigger consideration.')
    
    # Directory where the job is running
    parser.add_argument('--work-dir',  dest='workDir', default=None,
      help='Working directory to assemble the tiles in.')

    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("--no-geoheader-info", action="store_true", default=False,
      dest="noGeoHeaderInfo",
      help="Suppress writing some auxiliary information in geoheaders.")

    # DEBUG options
    parser.add_argument("--keep", action="store_true", dest="keep", default=False,
      help="Do not delete the temporary files.")
    parser.add_argument("--convert-tiles",  action="store_true", dest="convertTiles",
      help="Generate a uint8 version of each tile")
    parser.add_argument("--cog", action="store_true", dest="cog", default=False,
      help="Write a cloud-optimized GeoTIFF (COG).")
    parser.add_argument('-v', '--version', dest='version',  default=False, action='store_true',
      help='Display the version of software.')
    
    # PRIVATE options
    # These specify the tile location to request, bypassing the need to query mapproject.
    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_mapproject 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)

    # Run an initial query to parse and validate the user input
    sep = ","
    verbose = False
    progName = 'mapproject_single'
    settings = asp_system_utils.run_and_parse_output(progName, args + ['--parse-options'],
                                                     sep, verbose)
    options.demPath    = settings["dem"][0]
    options.imagePath  = settings["image"][0]
    options.cameraPath = settings["camera"][0]
    options.outputPath = settings["output_file"][0]

    # Remove these entries from the arguments.
    # TODO(oalexan1): This is fragile. What if we have the image file both on its own
    # and as an argument to some option? That is not the case for now.
    for v in [options.demPath, options.imagePath, options.cameraPath, options.outputPath]:
        if v != "" and v in args:
            args.remove(v)
    
    # This may contain --t_pixelwin, to be parsed later.
    options.extraArgs = args

    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 \
                    options.workDir is not None)
    
    if spawnedCopy: 
        # This copy was spawned to process a single tile
        return writeSingleTile(options) 

    # Now we are in the main process, not in a spawned copy

    # See if this is the ISIS session. The image file or camera file
    # must end in .cub (both are possible), but the camera file must
    # not end in .json, as then the CSM session is used.
    camExt =  os.path.splitext(options.cameraPath)[1].lower()
    isIsis = (asp_image_utils.isIsisFile(options.imagePath) or
              asp_image_utils.isIsisFile(options.cameraPath)) \
              and (camExt != '.json')
    
    # If the user did not set the tile size, then for ISIS use small
    # tiles, to have them run in parallel as individual processes,
    # since each process is necessarily single-threaded. For other
    # cameras use bigger tiles, as each process is multi-threaded, and
    # then file I/O is a bigger consideration.
    if options.tileSize is None:
        if isIsis:
            options.tileSize = 1024
        else:
            options.tileSize = 5120
            
    # See if the input tif file has RPC coefficients embedded in it.
    # In that case need to save them later in the output mapprojected
    # image.
    rpc_label = 'RPC Metadata'
    rpc_dict  = parse_gdal_metadata(options.imagePath, rpc_label)
    rpc_lines = format_metadata(rpc_label, rpc_dict)
     
    # Handle --query-pixel 
    if '--query-pixel' in options.extraArgs:
      cmd = ['mapproject_single', options.demPath, options.imagePath, options.cameraPath,
              options.outputPath] + options.extraArgs
      p = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
      ans, err = p.communicate()
      print(ans)
      return 0

    # Handle --query-projection        
    query_only = False
    if '--query-projection' in options.extraArgs:
        query_only = True
        # Wipe this, it will be added later right below
        asp_cmd_utils.wipe_option(options.extraArgs, '--query-projection', 0)

    # Call mapproject on the input data using subprocess and record output
    cmd = ['mapproject_single',  '--query-projection', options.demPath,
                options.imagePath, options.cameraPath, options.outputPath]
    cmd = cmd + options.extraArgs # Append other options
    if options.noGeoHeaderInfo:
        cmd += ['--no-geoheader-info']
    print(" ".join(cmd))
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
    projectionInfo, err = p.communicate()
    if not options.suppressOutput:
        print(projectionInfo)

    if query_only:
        return 0
    
    # Now find the image size in the output
    startPos    = projectionInfo.find('Output image size:')
    widthStart  = projectionInfo.find('width:', startPos)
    heightStart = projectionInfo.find('height:', widthStart)
    heightEnd   = projectionInfo.find(')', heightStart)
    if startPos < 0:
        print('Error computing projected area, quitting!')
        return -1

    # Extract values
    fullWidth   = int(projectionInfo[widthStart+7  : heightStart-1])
    fullHeight  = int(projectionInfo[heightStart+8 : heightEnd])
    print('Output image size is ' + str(fullWidth) + ' by ' + str(fullHeight) + ' pixels.')

    # For now we just break up the image into a user-specified tile size (default 1000x1000)
    numTilesX, numTilesY, tileList = generateTileList(fullWidth, fullHeight, options.tileSize)
    numTiles = numTilesX * numTilesY

    print('Splitting into ' + str(numTilesX) + ' by ' + str(numTilesY) + ' tiles.')

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

    # Make a temporary directory to store the tiles
    if options.workDir:
        tempFolder = options.workDir
    else: # No folder provided, create a default one
        tempFolder = os.path.join(outputFolder, outputName.replace('.', '_') + '_tiles/')
    asp_file_utils.createFolder(tempFolder)

    # Generate a text file that contains the boundaries for each tile
    argumentFilePath = os.path.join(tempFolder, 'argumentList.txt')
    with open(argumentFilePath, 'w') as argumentFile:
        for tile in tileList:
            argumentFile.write(str(tile[0]) + '\t' + str(tile[1]) + '\t'
                               + str(tile[2]) + '\t' + str(tile[3]) + '\n')

    # Indicate to GNU Parallel that there are multiple tab-separated
    # variables in the text file we just wrote
    parallelArgs = ['--colsep', "\\t"]

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

    processesPerCpu = 1

    # 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 the optimal number of processes if the user did not specify
    if not options.numProcesses:
        options.numProcesses = cpusPerNode * processesPerCpu

    # Note: mapproject can run with multiple threads on non-ISIS data but we don't use that
    # functionality here since we call mapproject 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
    mapproject_path = asp_system_utils.libexec_path('mapproject')
    commandList   = [python_path, mapproject_path,
                     '--pixelStartX', '{1}',
                     '--pixelStartY', '{2}',
                     '--pixelStopX',  '{3}',
                     '--pixelStopY',  '{4}',
                     '--work-dir', tempFolder,
                     options.demPath,
                     options.imagePath, options.cameraPath,
                     options.outputPath]
    if '--threads' not in options.extraArgs:
        commandList = commandList + ['--threads', '8'] # If not specified use 8 threads
    if options.convertTiles:
        commandList = commandList + ['--convert-tiles']
    if options.suppressOutput:
        commandList = commandList + ['--suppress-output']
    commandList = commandList + options.extraArgs # Append other options

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

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

    # Find the tiles that were generated
    tiles = []
    for tile in tileList:
        outTile = os.path.join(tempFolder, tile[4])
        if os.path.exists(outTile):
            tiles.append(outTile)
        else:
            print("Warning: Skipping non-existing file: ", outTile)
    if len(tiles) == 0:
        print("Removing: " + tempFolder)
        asp_file_utils.removeFolderIfExists(tempFolder)
        raise Exception("No mapprojected tif tiles were generated. Examine the " + \
                        "output above for more information.")

    # Build a gdal VRT file which is composed of all the processed tiles
    vrtPath = os.path.join(tempFolder, 'mosaic.vrt')
    cmd = "gdalbuildvrt -resolution highest " + vrtPath + " " + " ".join(tiles)
    ret = subprocess.call(["gdalbuildvrt", "-resolution", "highest", vrtPath]+tiles)
    if ret != 0:
        print("gdalbuildvrt error ", ret)
        exit(1)

    # Modify the vrt to append some metadata from the original tiles or the input file
    if not options.noGeoHeaderInfo:
        meta_label = 'Metadata'
        meta_dict  = parse_gdal_metadata(tiles[0], meta_label)
        meta_lines = format_metadata(meta_label, meta_dict)

        # Insert the metadata after the first line in the vrt.
        f = open(vrtPath, "r")
        lines = f.readlines()
        lines_before = [ lines[0] ]
        lines_after  = lines[1:len(lines)]
        f.close()

        lines = lines_before + meta_lines

        if len(rpc_dict) > 0:
            lines += rpc_lines

        lines += lines_after
        
        f = open(vrtPath, "w")
        f.writelines(lines)
        f.close()

    # Convert VRT file to final output file
    cmd = "gdal_translate -co BIGTIFF=IF_SAFER "
    if options.cog:
        # Use GDAL COG driver for cloud-optimized output. This ignores the
        # INTERLEAVE, TILED options.
        cmd += "-of COG -co COMPRESS=DEFLATE -co BLOCKSIZE=512 "
    else:
        # Standard tiled GeoTIFF
        cmd += "-co INTERLEAVE=BAND -co TILED=yes -co compress=lzw -co BLOCKXSIZE=256 -co BLOCKYSIZE=256 "
    cmd += vrtPath + " " + options.outputPath
    print(cmd)
    ans = os.system(cmd)

    # Clean up temporary files
    if not options.keep:
        print("Removing: " + tempFolder)
        asp_file_utils.removeFolderIfExists(tempFolder)

    if ans == 0: 
        print("Wrote: " + options.outputPath)
        maybe_copy_rpc(options.imagePath, options.outputPath)

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

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