#!/usr/bin/env python
# __BEGIN_LICENSE__
#  Copyright (c) 2009-2026, 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__

"""
stereo_tile - Process a single tile for distributed stereo.
"""

import sys, argparse, subprocess, re, os, math, time, glob, shutil, platform, shlex
import os.path as P

# Set up the path to Python modules about to load
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)
# Import after setting up the paths
import asp_system_utils, asp_string_utils, asp_cmd_utils, asp_stereo_utils, asp_alg_utils
asp_system_utils.verify_python_version_is_supported()

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

# Ensure we can look up the needed ASP libraries
if 'ASP_LIBRARY_PATH' in os.environ:
    os.environ['LD_LIBRARY_PATH'] = os.environ['ASP_LIBRARY_PATH']

def symlinkStatsFiles(outPrefix, tileSubdir, tilePrefix):
    '''Symlink stats files from parent directory to tile subdir.'''
    for f in glob.glob(outPrefix + '*stats.tif'):
        if os.path.isdir(f):
            continue
        relSrc = os.path.relpath(f, tileSubdir)
        dstF = f.replace(outPrefix, tilePrefix)
        if not os.path.lexists(dstF):
            os.symlink(relSrc, dstF)

def getTileByIndex(outPrefix, tileIndex):
    '''Read the tile list and return (BBox, padding) for the given tile index.'''
    tiles = asp_stereo_utils.readDistTileList(outPrefix)

    if tileIndex < 0 or tileIndex >= len(tiles):
        asp_system_utils.raise_error('Tile index ' + str(tileIndex) +
                                     ' out of range [0, ' + str(len(tiles) - 1) + ']')

    return tiles[tileIndex]

def runDemMosaicBlock(masterFile, blockIndex, opt):
    '''Run dem_mosaic on a block of DEMs specified in a list. The master file
    has the input list and output DEM name at the given index.'''
    with open(masterFile, 'r') as f:
        lines = [l.strip() for l in f if l.strip()]

    if blockIndex < 0 or blockIndex >= len(lines):
        asp_system_utils.raise_error('Block index ' + str(blockIndex) +
                                     ' out of range [0, ' +
                                     str(len(lines) - 1) + ']')

    parts = lines[blockIndex].split()
    if len(parts) != 2:
        asp_system_utils.raise_error('Invalid master file line: ' +
                                     lines[blockIndex])

    demListFile = parts[0]
    outputDem = parts[1]

    demMosaicPath = asp_system_utils.bin_path('dem_mosaic')
    commandList = [demMosaicPath, '--dem-list', demListFile, '-o', outputDem]

    if opt.dryrun:
        print('Command: ' + ' '.join(commandList))
    else:
        if opt.verbose:
            print(' '.join(commandList))
        subprocess.call(commandList)

def runPoint2dem(tilePrefix, opt):
    '''Run point2dem on a tile's point cloud.'''
    pcFile = tilePrefix + '-PC.tif'
    commandList = [asp_system_utils.bin_path('point2dem')]
    if opt.point2dem_options:
        commandList += shlex.split(opt.point2dem_options)

    # If --orthoimage is in the options with no argument, insert L.tif after it
    if opt.point2dem_options and '--orthoimage' in opt.point2dem_options:
        lFile = tilePrefix + '-L.tif'
        try:
            orthoIdx = commandList.index('--orthoimage')
            commandList.insert(orthoIdx + 1, lFile)
        except ValueError:
            pass

    commandList.append(pcFile)

    if opt.dryrun:
        print('Command: ' + ' '.join(commandList))
    else:
        if opt.verbose:
            print(' '.join(commandList))
        subprocess.call(commandList)

if __name__ == '__main__':
    usage = '''stereo_tile [options] <images> [<cameras>]
                  <output_prefix>
        Process a single tile for distributed stereo.\n''' + asp_system_utils.get_asp_version()

    p = argparse.ArgumentParser(usage=usage)
    p.add_argument('--tile-index', dest='tile_index', default=None, type=int,
                   help='Index of the tile to process (0-based).')
    p.add_argument('-e', '--entry-point', dest='entry_point', default=0, type=int,
                   help='Stereo pipeline entry point (default: 0). Values: 0=pprc, '
                   '1=corr, 2=blend, 3=rfne, 4=fltr, 5=tri, 6=cleanup, 7=dem.')
    p.add_argument('--stop-point', dest='stop_point', default=9, type=int,
                   help='Stereo pipeline stop point, stop before this step '
                   '(default: 9).')
    p.add_argument('--dem', dest='dem', default=None,
                   help='Input DEM for mapprojection. Required for tile processing.')
    p.add_argument('--point2dem-options', dest='point2dem_options', default='',
                   help='Options to pass to point2dem. Can pass --orthoimage '
                   'with no argument, and the L.tif file for each tile '
                   'will be autocompleted. Can also pass --errorimage.')
    p.add_argument('-s', '--stereo-file', dest='stereo_file',
                   default='./stereo.default',
                   help='Explicitly specify the stereo.default file to use. '
                   'This is optional.')
    p.add_argument('--corr-seed-mode', dest='seed_mode', default=None, type=int,
                   help='Correlation seed strategy. See stereo_corr for options.')
    p.add_argument('--sparse-disp-options', dest='sparse_disp_options',
                   default=None,
                   help='Options to pass directly to sparse_disp. Use quotes '
                   'around this string.')
    p.add_argument('--dem-mosaic-index', dest='dem_mosaic_index', default=None,
                   type=int,
                   help='Index of the mosaic block to process (0-based). '
                   'Requires --dem-mosaic-master. Invoked by stereo_dist for '
                   'parallel DEM mosaicking.')
    p.add_argument('--dem-mosaic-master', dest='dem_mosaic_master', default=None,
                   help='Master file enumerating for each block to mosaic the '
                   'input DEM list file and output DEM path, one block per line.')
    p.add_argument('--run-dir', dest='runDir', default=None,
                   help='Directory in which the script is running.')
    p.add_argument('--verbose', dest='verbose', default=False, action='store_true',
                   help='Display the commands being executed.')
    p.add_argument('--dry-run', dest='dryrun', default=False, action='store_true',
                   help='Do not launch the jobs, only print the commands that '
                   'should be run.')
    p.add_argument('-v', '--version', dest='version', default=False, action='store_true',
                   help='Display the version of the software.')

    (opt, args) = p.parse_known_args()
    args = asp_cmd_utils.clean_args(args)

    if opt.version:
        asp_system_utils.print_version_and_exit()

    # This can handle spaces in directory names, unlike when GNU parallel is in charge
    if opt.runDir is not None:
        os.chdir(opt.runDir)

    # Handle dem mosaic mode. Runs dem_mosaic on one block and exits. Hence this
    # aux tool does both mosaicking of DEMs or running stereo tiles depending on
    # the options.
    if opt.dem_mosaic_index is not None:
        if opt.dem_mosaic_master is None:
            asp_system_utils.raise_error('--dem-mosaic-master is required with ' +
                                         '--dem-mosaic-index.')
        runDemMosaicBlock(opt.dem_mosaic_master, opt.dem_mosaic_index, opt)
        sys.exit(0)

    # In tile processing mode there should be some positional arguments
    if not args and not opt.version:
        p.print_help()
        asp_system_utils.raise_error('Missing input files.', code=2)

    if opt.tile_index is None:
        asp_system_utils.raise_error('The option --tile-index is required.', code=2)

    if opt.dem is None:
        asp_system_utils.raise_error('The option --dem is required for tile processing.',
                                     code=2)

    # These are not supported as stereo_tile manages its own crop windows per tile
    for bad_opt in ['--left-image-crop-win', '--right-image-crop-win']:
        if bad_opt in args:
            asp_system_utils.raise_error('The option ' + bad_opt +
                                         ' is not supported by stereo_tile.')

    # Pass the DEM to the stereo C++ tools
    args.extend(['--dem', opt.dem])

    # Resolve corr-seed-mode from CLI, stereo file, or default
    sep = ","
    asp_stereo_utils.resolve_seed_mode(opt, args)

    if os.path.exists(opt.stereo_file):
        args.extend(['--stereo-file', opt.stereo_file])

    # Run stereo_parse to get settings
    settings = asp_system_utils.run_and_parse_output("stereo_parse", args, sep,
                                                     opt.verbose)
    outPrefix = settings['out_prefix'][0]

    print("Output prefix: " + outPrefix)
    print("Tile index: " + str(opt.tile_index))

    (tile, padding) = getTileByIndex(outPrefix, opt.tile_index)
    print("Tile: " + tile.name_str() + ", padding: " + str(padding))

    # Create tile subdir
    tileSubdir = outPrefix + '-' + tile.name_str()
    asp_system_utils.mkdir_p(tileSubdir)
    tilePrefix = tileSubdir + '/' + os.path.basename(outPrefix)

    # Symlink stats files to the subdir
    symlinkStatsFiles(outPrefix, tileSubdir, tilePrefix)

    # Common extra args for all steps
    extraArgs = ['--output-prefix-override', tilePrefix,
                  '--stereo-dist-mode',
                  '--left-image-crop-win',
                  str(tile.x), str(tile.y), str(tile.width), str(tile.height),
                  '--sgm-collar-size', str(padding)]

    Step = asp_stereo_utils.Step

    # Validate --tr is in point2dem options when DEM step will run
    if opt.entry_point <= Step.dem < opt.stop_point and '--tr' not in opt.point2dem_options:
        asp_system_utils.raise_error('The option --point2dem-options must include --tr to '
                                     'ensure consistent grid size across all tiles.')

    if opt.entry_point <= Step.pprc < opt.stop_point:
        asp_stereo_utils.stereo_run('stereo_pprc', args, opt, extra_args=extraArgs)

    if opt.entry_point <= Step.corr < opt.stop_point:
        # Compute low-res disparity (uses sparse_disp for seed mode 3)
        tile_args = args + ['--output-prefix-override', tilePrefix]
        asp_stereo_utils.calc_lowres_disp(tile_args, opt, sep)
        # Run full-res correlation
        asp_stereo_utils.stereo_run('stereo_corr', args, opt, extra_args=extraArgs)

    # Skip blend - in distributed stereo each tile goes to point cloud/DEM independently,
    # then only final DEMs are merged. No cross-talk needed until DEM merging.

    if opt.entry_point <= Step.rfne < opt.stop_point:
        asp_stereo_utils.stereo_run('stereo_rfne', args, opt, extra_args=extraArgs)

    if opt.entry_point <= Step.fltr < opt.stop_point:
        asp_stereo_utils.stereo_run('stereo_fltr', args, opt, extra_args=extraArgs)

    if opt.entry_point <= Step.tri < opt.stop_point:
        asp_stereo_utils.stereo_run('stereo_tri', args, opt, extra_args=extraArgs)

    # TODO(oalexan1): Implement cleanup step (delete intermediate files per tile)
    if opt.entry_point <= Step.clean < opt.stop_point:
        pass

    if opt.entry_point <= Step.dem < opt.stop_point:
        runPoint2dem(tilePrefix, opt)

    sys.exit(0)
