#!/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_dist - Distributed stereo processing.
"""

import sys, argparse, subprocess, re, os, time, glob, shutil, platform, math
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 asp_system_utils, asp_string_utils, asp_cmd_utils
asp_system_utils.verify_python_version_is_supported()
from asp_stereo_utils import * # must be after the path is altered above
from asp_dist_utils import *

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

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

if __name__ == '__main__':
    usage = '''stereo_dist [options] <images> [<cameras>]
                  <output_prefix>
        Distributed stereo processing.\n''' + asp_system_utils.get_asp_version()

    p = argparse.ArgumentParser(usage=usage)
    p.add_argument('-v', '--version', dest='version', default=False, action='store_true',
                   help='Display the version of software.')
    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('--tile-size', dest='tile_size', default=2048, type=int,
                   help='Size of each tile for distributed processing. This is before the padding is added.')
    p.add_argument('--tile-padding', dest='tile_padding', default=256, type=int,
                   help='Padding around each tile to avoid boundary artifacts.')
    p.add_argument('--threads-multiprocess', dest='threads_multi', default=None, type=int,
                   help='The number of threads to use per process when running multiple '
                   'processes.')
    p.add_argument('--threads-singleprocess', dest='threads_single', default=None, type=int,
                   help='The number of threads to use when running a single process '
                   '(PPRC and FLTR).')
    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, '
                   '8=mosaic.')
    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('--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('--processes', dest='processes', default=None, type=int,
                   help='The number of processes to use per node.')
    p.add_argument('--nodes-list', dest='nodes_list', default=None,
                   help='A file containing the list of computing nodes, one per line. '
                   'If not provided, run on the local machine.')
    p.add_argument('--dem', dest='dem', required=True,
                   help='Input DEM. Required.')
    p.add_argument('--mapproject', dest='mapproject', default=False,
                   action='store_true',
                   help='Mapproject the input images onto the DEM before '
                   'running stereo. The projection, grid size, and bounding '
                   'box are auto-determined to be consistent for both images. '
                   'These can be overridden with --t_srs, --tr, and --t_projwin.')
    p.add_argument('--t_srs', dest='t_srs', default=None,
                   help='Output projection for mapprojection. A PROJ string, '
                   'EPSG code, or path to a WKT file. Auto-determined if not set.')
    p.add_argument('--tr', dest='tr', default=None, type=float,
                   help='Output grid size (ground sample distance) for '
                   'mapprojection, in units of the projection. '
                   'Auto-determined as the finer of the two images if not set.')
    p.add_argument('--t_projwin', dest='t_projwin', default=None, type=float,
                   nargs=4, metavar=('XMIN', 'YMIN', 'XMAX', 'YMAX'),
                   help='Bounding box for mapprojection in projected coordinates: '
                   'xmin ymin xmax ymax. Auto-determined as the intersection of '
                   'the two images if not set.')
    p.add_argument('--parallel-options', dest='parallel_options', default='--sshdelay 0.2',
                   help='Options to pass directly to GNU Parallel.')

    # Warn about --threads being ignored
    if '--threads' in sys.argv:
        print("Ignoring the option --threads. Use --threads-multiprocess "
              "and --threads-singleprocess.")
        asp_cmd_utils.wipe_option(sys.argv, '--threads', 1)

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

    if opt.version:
        asp_system_utils.print_version_and_exit()

    if not args and not opt.version:
        p.print_help()
        asp_system_utils.raise_error('Missing input files.', code=2)

    # Validate --t_projwin if provided. Swap min/max if in wrong order,
    # as some tools print bounding boxes with inverted y axis.
    if opt.t_projwin is not None:
        if opt.t_projwin[0] > opt.t_projwin[2]:
            opt.t_projwin[0], opt.t_projwin[2] = opt.t_projwin[2], opt.t_projwin[0]
        if opt.t_projwin[1] > opt.t_projwin[3]:
            opt.t_projwin[1], opt.t_projwin[3] = opt.t_projwin[3], opt.t_projwin[1]

    # Validate that --t_srs, --tr, --t_projwin are only used with --mapproject
    if not opt.mapproject:
        for name in ['--t_srs', '--tr', '--t_projwin']:
            val = getattr(opt, name.lstrip('-').replace('-', '_'))
            if val is not None:
                asp_system_utils.raise_error(
                    name + ' can only be used with --mapproject.')

    # These are not supported as stereo_dist 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_dist.')

    # Pass the DEM to the stereo C++ tools. When mapprojecting, defer this
    # until after mapprojection, as stereo_parse will complain about
    # non-mapprojected images with --dem.
    if not opt.mapproject:
        args.extend(['--dem', opt.dem])

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

    # Set default threads if not specified
    if opt.threads_single is None:
        opt.threads_single = asp_system_utils.get_num_cpus()

    # Get the output prefix from stereo_parse
    sep = ","
    settings = asp_system_utils.run_and_parse_output("stereo_parse", args, sep, opt.verbose)
    outPrefix = settings['out_prefix'][0]

    # Mapproject both images onto the DEM with consistent parameters
    if opt.mapproject:
        (args, settings, outPrefix) = \
            mapprojAndUpdateArgs(opt, settings, outPrefix, sep)

    # Run stereo_tile on all tiles for pprc only (entry_point=0, stop_point=1)
    if opt.entry_point <= Step.pprc < opt.stop_point:
        # Run stereo_parse to generate the tile list
        tile_params = ['--stereo-dist-tile-params', str(opt.tile_size), str(opt.tile_padding)]
        stereo_run('stereo_parse', args, opt, extra_args=tile_params)
        # The tile layout already accounts for --proj-win, so remove it to avoid
        # conflict with per-tile --left-image-crop-win / --right-image-crop-win
        asp_cmd_utils.wipe_option(args, '--proj-win', 4)
        # Run stereo_pprc with full images and --stop-after-stats to compute statistics only
        asp_cmd_utils.wipe_option(args, '--threads', 1)
        pprc_args = ['--stop-after-stats', '--threads', str(opt.threads_single)]
        stereo_run('stereo_pprc', args, opt, extra_args=pprc_args)
        # Do the rest of preprocessing logic per tiles
        runStereoTiles(opt, args, outPrefix, Step.pprc, Step.corr)

    # Run stereo_tile on all tiles for corr only (entry_point=1, stop_point=2)
    if opt.entry_point <= Step.corr < opt.stop_point:
        runStereoTiles(opt, args, outPrefix, Step.corr, Step.blend)

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

    # Run stereo_tile on all tiles for rfne only (entry_point=3, stop_point=4)
    if opt.entry_point <= Step.rfne < opt.stop_point:
        runStereoTiles(opt, args, outPrefix, Step.rfne, Step.fltr)

    # Run stereo_tile on all tiles for fltr only (entry_point=4, stop_point=5)
    if opt.entry_point <= Step.fltr < opt.stop_point:
        runStereoTiles(opt, args, outPrefix, Step.fltr, Step.tri)

    # Run stereo_tile on all tiles for tri only (entry_point=5, stop_point=6)
    if opt.entry_point <= Step.tri < opt.stop_point:
        runStereoTiles(opt, args, outPrefix, Step.tri, Step.clean)

    # Cleanup step (entry_point=6, stop_point=7)
    if opt.entry_point <= Step.clean < opt.stop_point:
        runStereoTiles(opt, args, outPrefix, Step.clean, Step.dem)

    # Run point2dem on all tiles (entry_point=7, stop_point=8)
    if opt.entry_point <= Step.dem < opt.stop_point:
        runStereoTiles(opt, args, outPrefix, Step.dem, Step.mosaic)

    # Mosaic per-tile DEMs into final DEM (entry_point=8, stop_point=9)
    if opt.entry_point <= Step.mosaic < opt.stop_point:
        mosaicDems(opt, outPrefix)

    sys.exit(0)
