reproject_crs.py

"""Reproject photogrammetry outputs (LAS / DEM / orthomosaic) between CRSs.

Defaults perform a WGS 84 -> NAD83(CSRS) / UTM 12N + CGVD2013 transform with
the CGG2013 geoid correction (~+19.7 m in central Alberta). All EPSG codes
and the geoid step are configurable via CLI flags so the tool works for any
target CRS.

Steps applied (each can be skipped):
  1. Horizontal datum shift   (--source-las-epsg -> --target-horiz-epsg)
  2. Vertical geoid correction (--no-geoid to skip; uses PROJ-managed grids)
  3. Target CRS stamped on outputs (--target-compound-epsg)

Inputs (filenames default to cloud.las / dem.tif / ortho.tif; override with
--las-name / --dem-name / --ortho-name):
  - LAS point cloud  (full 3D transform with optional geoid correction)
  - DEM GeoTIFF      (horizontal reprojection + geoid correction on Z values)
  - Orthomosaic TIFF (horizontal reprojection only)

Usage:
  reproject-crs --input-dir my_metashape_output
  reproject-crs --input-dir my_metashape_output --no-geoid --target-horiz-epsg 32612
  reproject-crs --input-dir my_metashape_output --suffix _csrs --dry-run
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

import laspy
import numpy as np
import pyproj
import rasterio
from pyproj import CRS, Transformer
from rasterio.crs import CRS as RioCRS
from rasterio.warp import Resampling, calculate_default_transform, reproject


# ---------------------------------------------------------------------------
# CRS defaults (overridable via CLI flags). These defaults assume a WGS 84
# UTM 12N source and an Alberta NAD83(CSRS) + CGVD2013 target.
# ---------------------------------------------------------------------------
SRC_LAS_EPSG = 32612       # WGS 84 / UTM zone 12N
SRC_RASTER_EPSG = 4326     # WGS 84 geographic
DST_HORIZ_EPSG = 2956      # NAD83(CSRS) / UTM zone 12N
DST_COMPOUND_EPSG = 6655   # NAD83(CSRS) / UTM zone 12N + CGVD2013 height
SRC_GEOG3D_EPSG = 4955     # NAD83(CSRS) geographic 3D (for geoid lookup)

APPLY_GEOID = True         # If False, the vertical step is skipped entirely.

GEOID_GRID = "ca_nrc_CGG2013n83.tif"
GEOID_GRID_A = "ca_nrc_CGG2013an83.tif"
GEOID_URL = f"https://cdn.proj.org/{GEOID_GRID}"
GEOID_URL_A = f"https://cdn.proj.org/{GEOID_GRID_A}"


# ---------------------------------------------------------------------------
# Grid management
# ---------------------------------------------------------------------------

def ensure_geoid_grid() -> None:
    """Verify the CGG2013 geoid grid is available to PROJ. Download if missing."""
    user_dir = Path(pyproj.datadir.get_user_data_dir())
    user_dir.mkdir(parents=True, exist_ok=True)

    for name, url in [(GEOID_GRID, GEOID_URL), (GEOID_GRID_A, GEOID_URL_A)]:
        grid_path = user_dir / name
        if grid_path.exists() and grid_path.stat().st_size > 1_000_000:
            print(f"  Grid {name}: found ({grid_path.stat().st_size:,} bytes)")
            continue
        print(f"  Downloading {name} from cdn.proj.org ...")
        import urllib.request
        urllib.request.urlretrieve(url, grid_path)
        print(f"  Saved to {grid_path} ({grid_path.stat().st_size:,} bytes)")


def _init_pyproj() -> None:
    """Enable PROJ network and (when geoid correction is on) verify grids."""
    pyproj.network.set_network_enabled(True)
    if APPLY_GEOID:
        print("Checking geoid grids...")
        ensure_geoid_grid()
    else:
        print("Geoid correction disabled (--no-geoid); skipping grid check.")


# ---------------------------------------------------------------------------
# Transformers (lazy-init, module-level singletons)
# ---------------------------------------------------------------------------

_transformers: dict[str, Transformer] = {}


def _get_transformers() -> dict[str, Transformer]:
    if not _transformers:
        _transformers["horiz"] = Transformer.from_crs(
            f"EPSG:{SRC_LAS_EPSG}", f"EPSG:{DST_HORIZ_EPSG}", always_xy=True
        )
        _transformers["to_geo"] = Transformer.from_crs(
            f"EPSG:{SRC_LAS_EPSG}", "EPSG:4326", always_xy=True
        )
        if APPLY_GEOID:
            _transformers["vert"] = Transformer.from_crs(
                CRS.from_epsg(SRC_GEOG3D_EPSG),
                CRS.from_epsg(DST_COMPOUND_EPSG),
                always_xy=True,
            )
    return _transformers


# ---------------------------------------------------------------------------
# LAS reprojection
# ---------------------------------------------------------------------------

def reproject_las(src_path: Path, dst_path: Path, *, dry_run: bool = False) -> dict:
    """Reproject LAS point cloud from EPSG:32612 to EPSG:6655."""
    print(f"\n{'[DRY RUN] ' if dry_run else ''}Reprojecting LAS: {src_path.name}")
    t0 = time.time()

    las = laspy.read(str(src_path))
    n_pts = len(las.points)
    print(f"  Read {n_pts:,} points ({src_path.stat().st_size / 1e9:.1f} GB)")

    # Extract coordinates as float64
    x = np.array(las.x, dtype=np.float64)
    y = np.array(las.y, dtype=np.float64)
    z = np.array(las.z, dtype=np.float64)

    # Sample before transform
    sample_idx = [0, n_pts // 2, n_pts - 1]
    before = [(x[i], y[i], z[i]) for i in sample_idx]

    transformers = _get_transformers()

    # Step 1: Horizontal datum shift
    print(f"  Step 1/2: Horizontal datum shift (EPSG:{SRC_LAS_EPSG} -> EPSG:{DST_HORIZ_EPSG})...")
    x_new, y_new = transformers["horiz"].transform(x, y)

    # Step 2: Vertical geoid correction (optional)
    if APPLY_GEOID:
        print(f"  Step 2/2: Vertical geoid correction (ellipsoidal -> EPSG:{DST_COMPOUND_EPSG})...")
        lon, lat = transformers["to_geo"].transform(x, y)
        _, _, z_new = transformers["vert"].transform(lon, lat, z)
        z_shift = np.median(z_new - z)
        if abs(z_shift) < 1.0:
            print(f"  WARNING: Median Z shift is only {z_shift:.3f}m — geoid may not have applied")
    else:
        print("  Step 2/2: skipped (--no-geoid).")
        z_new = z
        z_shift = 0.0

    # Sample after transform
    after = [(x_new[i], y_new[i], z_new[i]) for i in sample_idx]

    if dry_run:
        dt = time.time() - t0
        return _build_report("LAS", src_path, None, n_pts, before, after, z_shift, dt)

    # Build output LAS
    print(f"  Writing {dst_path.name}...")
    header = laspy.LasHeader(
        point_format=las.header.point_format,
        version=las.header.version,
    )
    # Compute new offsets from data range
    header.offsets = np.array([
        np.floor(x_new.min()),
        np.floor(y_new.min()),
        np.floor(z_new.min()),
    ])
    header.scales = las.header.scales

    # Add CRS as WKT VLR (LAS 1.4 style, also works in 1.2+).
    # Stamp the compound CRS when geoid was applied; otherwise the 2D target.
    stamp_epsg = DST_COMPOUND_EPSG if APPLY_GEOID else DST_HORIZ_EPSG
    wkt = CRS.from_epsg(stamp_epsg).to_wkt()
    wkt_vlr = laspy.VLR(
        user_id="LASF_Projection",
        record_id=2112,
        description="OGC Coordinate System WKT",
        record_data=wkt.encode("utf-8") + b"\x00",
    )
    # GeoKeyDirectory VLR for EPSG code (broad compatibility)
    geo_keys = _build_geokey_vlr(stamp_epsg)
    header.vlrs = [wkt_vlr, geo_keys]

    out_las = laspy.LasData(header)

    # Copy all point attributes from source
    for dim in las.point_format.dimension_names:
        if dim not in ("X", "Y", "Z"):
            setattr(out_las, dim, getattr(las, dim))

    out_las.x = x_new
    out_las.y = y_new
    out_las.z = z_new

    out_las.write(str(dst_path))
    dt = time.time() - t0
    print(f"  Written {dst_path} ({dst_path.stat().st_size / 1e9:.1f} GB, {dt:.0f}s)")

    return _build_report("LAS", src_path, dst_path, n_pts, before, after, z_shift, dt)


def _build_geokey_vlr(epsg: int) -> laspy.VLR:
    """Build a minimal GeoKeyDirectoryVLR stamping the EPSG code."""
    import struct
    # GeoKeyDirectory: version=1, revision=1, minor_revision=0, numberOfKeys=1
    # Key: GTModelTypeGeoKey(1024)=1 (projected), ProjectedCSTypeGeoKey(3072)=epsg
    data = struct.pack(
        "<" + "H" * 12,
        1, 1, 0, 2,       # header: version, rev, minor, nkeys
        1024, 0, 1, 1,    # GTModelTypeGeoKey = ModelTypeProjected
        3072, 0, 1, epsg,  # ProjectedCSTypeGeoKey
    )
    return laspy.VLR(
        user_id="LASF_Projection",
        record_id=34735,
        description="GeoKeyDirectoryTag",
        record_data=data,
    )


# ---------------------------------------------------------------------------
# DEM reprojection
# ---------------------------------------------------------------------------

def reproject_dem(src_path: Path, dst_path: Path, *, dry_run: bool = False) -> dict:
    """Reproject DEM from EPSG:4326 to EPSG:2956 and apply geoid correction to Z values."""
    print(f"\n{'[DRY RUN] ' if dry_run else ''}Reprojecting DEM: {src_path.name}")
    t0 = time.time()

    with rasterio.open(str(src_path)) as src:
        src_crs = src.crs
        print(f"  Source: {src.width}x{src.height}, CRS={src_crs}")

        # Compute geoid undulation at DEM centroid (only if geoid is enabled)
        cx = (src.bounds.left + src.bounds.right) / 2
        cy = (src.bounds.bottom + src.bounds.top) / 2
        if APPLY_GEOID:
            geoid_n = _compute_geoid_undulation(cx, cy)
            print(f"  Centroid: ({cy:.6f}N, {cx:.6f}E)")
            print(f"  Geoid undulation N = {geoid_n:.3f}m (H = h - N, so +{-geoid_n:.3f}m)")
        else:
            geoid_n = 0.0
            print(f"  Centroid: ({cy:.6f}N, {cx:.6f}E)")
            print("  Geoid correction disabled (--no-geoid); Z values pass through unchanged.")

        # Read source data
        src_data = src.read(1)
        src_nodata = src.nodata
        z_before_sample = float(np.nanmedian(src_data[src_data != src_nodata])) if src_nodata else float(np.nanmedian(src_data))

        # Compute target transform and dimensions
        dst_crs = RioCRS.from_epsg(DST_HORIZ_EPSG)
        transform, width, height = calculate_default_transform(
            src_crs, dst_crs, src.width, src.height, *src.bounds
        )

        if dry_run:
            z_after = z_before_sample - geoid_n
            dt = time.time() - t0
            return {
                "type": "DEM", "src": str(src_path), "dst": None,
                "shape": f"{width}x{height}", "geoid_N": geoid_n,
                "z_before_median": z_before_sample, "z_after_median": z_after,
                "duration_s": dt,
            }

        # Reproject horizontally
        print(f"  Reprojecting raster grid (EPSG:{src_crs.to_epsg() if src_crs else SRC_RASTER_EPSG} -> EPSG:{DST_HORIZ_EPSG})...")
        dst_meta = src.meta.copy()
        dst_meta.update({
            "crs": dst_crs,
            "transform": transform,
            "width": width,
            "height": height,
        })

        dst_data = np.empty((height, width), dtype=src_data.dtype)
        reproject(
            source=src_data,
            destination=dst_data,
            src_transform=src.transform,
            src_crs=src_crs,
            dst_transform=transform,
            dst_crs=dst_crs,
            src_nodata=src_nodata,
            dst_nodata=src_nodata,
            resampling=Resampling.bilinear,
        )

    # Apply geoid correction to pixel values: H = h - N (skip if disabled)
    if APPLY_GEOID:
        print("  Applying geoid correction to elevation values...")
        if src_nodata is not None:
            valid = dst_data != src_nodata
            dst_data[valid] = dst_data[valid] - geoid_n
        else:
            dst_data = dst_data - geoid_n

    z_after_sample = float(np.nanmedian(dst_data[dst_data != src_nodata])) if src_nodata else float(np.nanmedian(dst_data))

    # Write output with target CRS (compound when geoid was applied, 2D otherwise)
    print(f"  Writing {dst_path.name}...")
    dst_meta["count"] = 1
    stamp_epsg = DST_COMPOUND_EPSG if APPLY_GEOID else DST_HORIZ_EPSG
    dst_meta["crs"] = RioCRS.from_wkt(CRS.from_epsg(stamp_epsg).to_wkt())

    with rasterio.open(str(dst_path), "w", **dst_meta) as dst:
        dst.write(dst_data, 1)

    dt = time.time() - t0
    print(f"  Written {dst_path} ({dst_path.stat().st_size / 1e6:.0f} MB, {dt:.0f}s)")

    return {
        "type": "DEM", "src": str(src_path), "dst": str(dst_path),
        "shape": f"{width}x{height}", "geoid_N": geoid_n,
        "z_before_median": z_before_sample, "z_after_median": z_after_sample,
        "duration_s": dt,
    }


def _compute_geoid_undulation(lon: float, lat: float) -> float:
    """Compute geoid undulation N at a geographic point.

    Uses a reference ellipsoidal height (0m) and returns N = h - H,
    where H is the orthometric height from the geoid model.
    """
    transformers = _get_transformers()
    # Transform geographic 3D (with h=0) to compound CRS
    t_geo = Transformer.from_crs(
        CRS.from_epsg(SRC_GEOG3D_EPSG), CRS.from_epsg(DST_COMPOUND_EPSG), always_xy=True
    )
    _, _, h_ortho = t_geo.transform(lon, lat, 0.0)
    # N = h - H, so with h=0: N = -H
    return -h_ortho


# ---------------------------------------------------------------------------
# Orthomosaic reprojection
# ---------------------------------------------------------------------------

def reproject_ortho(src_path: Path, dst_path: Path, *, dry_run: bool = False) -> dict:
    """Reproject orthomosaic from EPSG:4326 to EPSG:2956 (horizontal only)."""
    print(f"\n{'[DRY RUN] ' if dry_run else ''}Reprojecting ortho: {src_path.name}")
    t0 = time.time()

    with rasterio.open(str(src_path)) as src:
        print(f"  Source: {src.width}x{src.height}, {src.count} bands, CRS={src.crs}")

        dst_crs = RioCRS.from_epsg(DST_HORIZ_EPSG)
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds
        )

        if dry_run:
            dt = time.time() - t0
            return {
                "type": "Ortho", "src": str(src_path), "dst": None,
                "shape": f"{width}x{height}", "bands": src.count,
                "duration_s": dt,
            }

        dst_meta = src.meta.copy()
        dst_meta.update({
            "crs": dst_crs,
            "transform": transform,
            "width": width,
            "height": height,
        })

        print(f"  Reprojecting {src.count} bands (EPSG:{src.crs.to_epsg() if src.crs else SRC_RASTER_EPSG} -> EPSG:{DST_HORIZ_EPSG})...")
        with rasterio.open(str(dst_path), "w", **dst_meta) as dst:
            for band_idx in range(1, src.count + 1):
                src_data = src.read(band_idx)
                dst_data = np.empty((height, width), dtype=src_data.dtype)
                reproject(
                    source=src_data,
                    destination=dst_data,
                    src_transform=src.transform,
                    src_crs=src.crs,
                    dst_transform=transform,
                    dst_crs=dst_crs,
                    src_nodata=src.nodata,
                    dst_nodata=src.nodata,
                    resampling=Resampling.bilinear,
                )
                dst.write(dst_data, band_idx)
                if src.count > 1:
                    print(f"    Band {band_idx}/{src.count} done")

    dt = time.time() - t0
    print(f"  Written {dst_path} ({dst_path.stat().st_size / 1e6:.0f} MB, {dt:.0f}s)")

    return {
        "type": "Ortho", "src": str(src_path), "dst": str(dst_path),
        "shape": f"{width}x{height}", "bands": src.count,
        "duration_s": dt,
    }


# ---------------------------------------------------------------------------
# Verification
# ---------------------------------------------------------------------------

def print_report(reports: list[dict]) -> None:
    """Print a summary verification report."""
    print("\n" + "=" * 70)
    print("REPROJECTION SUMMARY")
    print("=" * 70)
    target_label = f"EPSG:{DST_COMPOUND_EPSG}" if APPLY_GEOID else f"EPSG:{DST_HORIZ_EPSG}"
    print(f"Target CRS: {target_label}")
    print()

    for r in reports:
        print(f"--- {r['type']} ---")
        print(f"  Source: {r['src']}")
        print(f"  Output: {r.get('dst', '(dry run)')}")

        if r["type"] == "LAS":
            print(f"  Points: {r['n_points']:,}")
            print(f"  Sample coordinates (first, middle, last):")
            for label, before, after in zip(
                ["First", "Middle", "Last"], r["before"], r["after"]
            ):
                dx = after[0] - before[0]
                dy = after[1] - before[1]
                dz = after[2] - before[2]
                print(f"    {label}: ({before[0]:.1f}, {before[1]:.1f}, {before[2]:.2f})")
                print(f"         -> ({after[0]:.1f}, {after[1]:.1f}, {after[2]:.2f})")
                print(f"            dX={dx:+.3f}  dY={dy:+.3f}  dZ={dz:+.3f}")
            print(f"  Median Z shift: {r['z_shift_median']:+.3f}m (geoid correction)")

        elif r["type"] == "DEM":
            print(f"  Output size: {r['shape']}")
            print(f"  Geoid undulation N: {r['geoid_N']:.3f}m")
            print(f"  Median elevation: {r['z_before_median']:.2f}m (before) -> {r['z_after_median']:.2f}m (after)")
            print(f"  Shift: {r['z_after_median'] - r['z_before_median']:+.3f}m")

        elif r["type"] == "Ortho":
            print(f"  Output size: {r['shape']}, {r['bands']} bands")
            print(f"  Horizontal reprojection only (no Z data)")

        print(f"  Duration: {r['duration_s']:.0f}s")
        print()


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _build_report(
    rtype: str, src: Path, dst: Path | None,
    n_pts: int, before: list, after: list,
    z_shift: float, dt: float,
) -> dict:
    return {
        "type": rtype, "src": str(src), "dst": str(dst) if dst else "(dry run)",
        "n_points": n_pts, "before": before, "after": after,
        "z_shift_median": z_shift, "duration_s": dt,
    }


def main() -> None:
    global SRC_LAS_EPSG, SRC_RASTER_EPSG, DST_HORIZ_EPSG, DST_COMPOUND_EPSG, SRC_GEOG3D_EPSG, APPLY_GEOID

    parser = argparse.ArgumentParser(
        description="Reproject photogrammetry outputs (LAS / DEM / orthomosaic) "
                    "between coordinate reference systems."
    )
    parser.add_argument(
        "--input-dir", required=True, type=Path,
        help="Directory containing the LAS / DEM / ortho input files",
    )
    parser.add_argument(
        "--las-name", default="cloud.las",
        help="LAS filename inside --input-dir (default: cloud.las)",
    )
    parser.add_argument(
        "--dem-name", default="dem.tif",
        help="DEM filename inside --input-dir (default: dem.tif)",
    )
    parser.add_argument(
        "--ortho-name", default="ortho.tif",
        help="Orthomosaic filename inside --input-dir (default: ortho.tif)",
    )
    parser.add_argument(
        "--suffix", default="_reprojected",
        help="Suffix for output filenames (default: _reprojected)",
    )
    parser.add_argument(
        "--source-las-epsg", type=int, default=SRC_LAS_EPSG,
        help=f"Source EPSG for the LAS point cloud (default: {SRC_LAS_EPSG} = WGS 84 / UTM 12N)",
    )
    parser.add_argument(
        "--source-raster-epsg", type=int, default=SRC_RASTER_EPSG,
        help=f"Fallback source EPSG for rasters with no embedded CRS (default: {SRC_RASTER_EPSG})",
    )
    parser.add_argument(
        "--target-horiz-epsg", type=int, default=DST_HORIZ_EPSG,
        help=f"Target horizontal EPSG (default: {DST_HORIZ_EPSG} = NAD83(CSRS) / UTM 12N)",
    )
    parser.add_argument(
        "--target-compound-epsg", type=int, default=DST_COMPOUND_EPSG,
        help=f"Target compound (3D) EPSG used when --no-geoid is NOT set (default: {DST_COMPOUND_EPSG} = NAD83(CSRS) / UTM 12N + CGVD2013)",
    )
    parser.add_argument(
        "--source-geog3d-epsg", type=int, default=SRC_GEOG3D_EPSG,
        help=f"Source geographic-3D EPSG used for geoid lookup (default: {SRC_GEOG3D_EPSG} = NAD83(CSRS) 3D)",
    )
    parser.add_argument(
        "--no-geoid", action="store_true",
        help="Skip the vertical geoid correction step (horizontal reprojection only)",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Compute transforms and report without writing output files",
    )
    parser.add_argument("--skip-las", action="store_true", help="Skip LAS reprojection")
    parser.add_argument("--skip-dem", action="store_true", help="Skip DEM reprojection")
    parser.add_argument("--skip-ortho", action="store_true", help="Skip orthomosaic reprojection")
    args = parser.parse_args()

    # Apply CLI overrides to module-level config (read by helper functions).
    SRC_LAS_EPSG = args.source_las_epsg
    SRC_RASTER_EPSG = args.source_raster_epsg
    DST_HORIZ_EPSG = args.target_horiz_epsg
    DST_COMPOUND_EPSG = args.target_compound_epsg
    SRC_GEOG3D_EPSG = args.source_geog3d_epsg
    APPLY_GEOID = not args.no_geoid

    input_dir = args.input_dir.resolve()
    if not input_dir.is_dir():
        sys.exit(f"Error: {input_dir} is not a directory")

    sfx = args.suffix
    las_src = input_dir / args.las_name
    dem_src = input_dir / args.dem_name
    ortho_src = input_dir / args.ortho_name
    las_dst = input_dir / f"{Path(args.las_name).stem}{sfx}.las"
    dem_dst = input_dir / f"{Path(args.dem_name).stem}{sfx}.tif"
    ortho_dst = input_dir / f"{Path(args.ortho_name).stem}{sfx}.tif"

    target_label = f"EPSG:{DST_COMPOUND_EPSG}" if APPLY_GEOID else f"EPSG:{DST_HORIZ_EPSG}"
    print(f"Reproject photogrammetry outputs -> {target_label}")
    print(f"Input directory: {input_dir}")
    if args.dry_run:
        print("Mode: DRY RUN (no files will be written)")
    print()

    _init_pyproj()

    reports: list[dict] = []

    if not args.skip_las and las_src.exists():
        reports.append(reproject_las(las_src, las_dst, dry_run=args.dry_run))
    elif not las_src.exists():
        print(f"\n  Skipping LAS: {las_src} not found")

    if not args.skip_dem and dem_src.exists():
        reports.append(reproject_dem(dem_src, dem_dst, dry_run=args.dry_run))
    elif not dem_src.exists():
        print(f"\n  Skipping DEM: {dem_src} not found")

    if not args.skip_ortho and ortho_src.exists():
        reports.append(reproject_ortho(ortho_src, ortho_dst, dry_run=args.dry_run))
    elif not ortho_src.exists():
        print(f"\n  Skipping ortho: {ortho_src} not found")

    if reports:
        print_report(reports)


if __name__ == "__main__":
    main()