From 8b0c053e9aa101ea48157da36c7557b9951f3511 Mon Sep 17 00:00:00 2001
From: vschaffn <valentin.schaffner@cs-soprasteria.com>
Date: Thu, 31 Oct 2024 09:59:40 +0100
Subject: [PATCH] feat: update cli for coregistration

---
 xdem/__init__.py | 46 +++++++++++++++++++++++++++++++++++++++++++---
 xdem/xdem_cli.py | 42 ++++++++++++++++++++++++------------------
 2 files changed, 67 insertions(+), 21 deletions(-)

diff --git a/xdem/__init__.py b/xdem/__init__.py
index f7de0ef7..6b65b226 100644
--- a/xdem/__init__.py
+++ b/xdem/__init__.py
@@ -1,3 +1,8 @@
+import logging
+import os
+
+import geoutils
+
 from xdem import (  # noqa
     coreg,
     dem,
@@ -8,6 +13,7 @@
     terrain,
     volume,
 )
+from xdem.coreg.workflows import dem_coregistration
 from xdem.ddem import dDEM  # noqa
 from xdem.dem import DEM  # noqa
 from xdem.demcollection import DEMCollection  # noqa
@@ -24,8 +30,42 @@
     )
 
 
-def run(reference_dem: str, dem_to_be_aligned: str, verbose: str) -> None:
+def coregister(ref_dem_path: str, tba_dem_path: str) -> None:
     """
-    Function to compare DEMs
+    Function to compare and coregister Digital Elevation Models (DEMs).
+
+    This function verifies the existence of the provided DEM paths,
+    loads the reference DEM and the DEM to be aligned, and performs
+    coregistration. The aligned DEM and an inlier mask are then saved
+    to disk.
+
+    :param ref_dem_path: Path to the reference DEM file.
+    :param tba_dem_path: Path to the DEM that needs to be aligned to the reference.
+    :return:
+    :raises FileNotFoundError: if the reference DEM or the DEM to be aligned does not exist.
     """
-    print("hello world")
+    # Verify that both DEM paths exist
+    if not os.path.exists(ref_dem_path):
+        raise FileNotFoundError(f"Reference DEM path does not exist: {ref_dem_path}")
+    if not os.path.exists(tba_dem_path):
+        raise FileNotFoundError(f"DEM to be aligned path does not exist: {tba_dem_path}")
+
+    logging.info("Loading DEMs: %s, %s", ref_dem_path, tba_dem_path)
+
+    # Load the reference and secondary DEMs
+    reference_dem, to_be_aligned_dem = geoutils.raster.load_multiple_rasters([ref_dem_path, tba_dem_path])
+
+    # Execute coregistration
+    logging.info("Starting coregistration...")
+    coreg_dem, coreg_method, out_stats, inlier_mask = dem_coregistration(
+        to_be_aligned_dem, reference_dem, "aligned_dem.tiff"
+    )
+
+    # Save outputs
+    logging.info("Saving aligned DEM and inlier mask...")
+    inlier_rst = coreg_dem.copy(new_array=inlier_mask)
+    inlier_rst.save("inlier_mask.tiff")
+
+    # Print the coregistration details
+    print("Coregistration statistics:\n", out_stats)
+    logging.info("Coregistration completed")
diff --git a/xdem/xdem_cli.py b/xdem/xdem_cli.py
index 89e96d63..796cb8d1 100644
--- a/xdem/xdem_cli.py
+++ b/xdem/xdem_cli.py
@@ -1,4 +1,5 @@
 import argparse
+import logging
 from argparse import ArgumentParser
 
 import argcomplete
@@ -12,20 +13,7 @@ def get_parser() -> ArgumentParser:
 
     :return: parser
     """
-    parser = argparse.ArgumentParser(
-        description="Compare Digital Elevation Models",
-        fromfile_prefix_chars="@",
-    )
-
-    parser.add_argument(
-        "reference_dem",
-        help="path to a reference dem",
-    )
-
-    parser.add_argument(
-        "dem_to_be_aligned",
-        help="path to a second dem",
-    )
+    parser = argparse.ArgumentParser(prog="xdem", description="xDEM command-line interface")
 
     parser.add_argument(
         "--loglevel",
@@ -41,6 +29,13 @@ def get_parser() -> ArgumentParser:
         version=f"%(prog)s {xdem.__version__}",
     )
 
+    subparsers = parser.add_subparsers(title="Subcommands", dest="command")
+
+    # Subcommand for coregistration
+    coregister_parser = subparsers.add_parser("coregister", help="Coregister two DEMs")
+    coregister_parser.add_argument("reference_dem", help="path to a reference dem")
+    coregister_parser.add_argument("dem_to_be_aligned", help="path to a second dem")
+
     return parser
 
 
@@ -51,10 +46,21 @@ def main() -> None:
     parser = get_parser()
     argcomplete.autocomplete(parser)
     args = parser.parse_args()
-    try:
-        xdem.run(args.reference_dem, args.dem_to_be_aligned, args.loglevel)
-    except Exception as e:
-        print(f"Error: {e}")
+
+    # Show help if no subcommand is provided
+    if not args.command:
+        parser.print_help()
+        return
+
+    # Set the logging configuration
+    logging.basicConfig(level=args.loglevel)
+
+    # Handle coregister subcommand
+    if args.command == "coregister":
+        try:
+            xdem.coregister(args.reference_dem, args.dem_to_be_aligned)
+        except Exception as e:
+            print(f"Error: {e}")
 
 
 if __name__ == "__main__":