"""
ingest_dem.py — turn a REAL DEM (and optional soil raster) into the full
analysis + cinematic pipeline. Drop a GeoTIFF in and run.

    python3 ingest_dem.py path/to/dem.tif [--erode] [--flyover] [--soil soil.tif]

Accepts: GeoTIFF/.tif (rasterio), ESRI ASCII .asc, .xyz point cloud, .npy.
Produces: elevation + hillshade + slope maps, flow-accumulation rivers, real-
world stats (area in km²/mi², relief, mean slope), a baked .npz, and — with
flags — an erosion run and a voxel-space flyover over the true terrain.

Get a DEM (see DATA_SOURCES.md): Copernicus GLO-30 or SRTM GL1 (30 m) for the
Ein Auja / eastern Judean flank clip through OpenTopography, as GeoTIFF.
"""
import sys, os, numpy as np
from scipy.ndimage import zoom, gaussian_filter, distance_transform_edt
import geomorph as G


# ---------------------------------------------------------------- loading ---
def load_dem(path):
    """Return (elev[m] float32, meta dict). Fills nodata by nearest-neighbour."""
    ext = os.path.splitext(path)[1].lower()
    meta = {"path": path, "px_m": None, "bbox": None, "crs": None}
    if ext in (".tif", ".tiff"):
        import rasterio
        with rasterio.open(path) as ds:
            z = ds.read(1).astype(np.float32)
            nod = ds.nodata
            b, crs = ds.bounds, ds.crs
            meta["crs"] = str(crs)
            meta["bbox"] = (b.left, b.bottom, b.right, b.top)
            # pixel size in metres (project degrees at scene-centre latitude)
            px, py = abs(ds.transform.a), abs(ds.transform.e)
            if crs and crs.is_geographic:
                latc = 0.5 * (b.bottom + b.top)
                meta["px_m"] = (px * 111320 * np.cos(np.radians(latc)),
                                py * 110540)
            else:
                meta["px_m"] = (px, py)
            if nod is not None:
                z = np.where(z == nod, np.nan, z)
    elif ext == ".npy":
        z = np.load(path).astype(np.float32)
    elif ext == ".asc":
        z = np.loadtxt(path, skiprows=6, dtype=np.float32)
    elif ext == ".xyz":
        pts = np.loadtxt(path)
        # grid an x,y,z point cloud onto a regular raster
        xs, ys = np.unique(pts[:, 0]), np.unique(pts[:, 1])
        z = np.full((len(ys), len(xs)), np.nan, np.float32)
        xi = np.searchsorted(xs, pts[:, 0]); yi = np.searchsorted(ys, pts[:, 1])
        z[yi, xi] = pts[:, 2]; z = z[::-1]
    else:
        raise SystemExit(f"unsupported DEM format: {ext}")

    if np.isnan(z).any():                      # fill voids by nearest valid
        idx = distance_transform_edt(np.isnan(z), return_distances=False,
                                     return_indices=True)
        z = z[tuple(idx)]
    return z.astype(np.float32), meta


# ------------------------------------------------------------- reporting ----
def describe(z, meta):
    H, W = z.shape
    print(f"DEM {W}x{H}  elev[{z.min():.1f}, {z.max():.1f}] m  relief {np.ptp(z):.0f} m")
    if meta["px_m"]:
        pxm = float(np.mean(meta["px_m"]))
        area_km2 = (pxm * W) * (pxm * H) / 1e6
        print(f"pixel ~{pxm:.0f} m  extent {pxm*W/1000:.1f} x {pxm*H/1000:.1f} km"
              f"  area {area_km2:.0f} km^2 ({area_km2*0.386:.0f} mi^2)")
    if meta["bbox"]:
        print(f"bbox {tuple(round(v,4) for v in meta['bbox'])}  crs {meta['crs']}")
    gy, gx = np.gradient(z)
    slope_deg = np.degrees(np.arctan(np.hypot(gx, gy) /
                                     (np.mean(meta['px_m']) if meta['px_m'] else 30)))
    print(f"mean slope {slope_deg.mean():.1f} deg  max {slope_deg.max():.1f} deg")
    return slope_deg


# ------------------------------------------------------------------- main ---
def main():
    if len(sys.argv) < 2:
        raise SystemExit(__doc__)
    path = sys.argv[1]
    do_erode = "--erode" in sys.argv
    do_fly = "--flyover" in sys.argv
    soil = sys.argv[sys.argv.index("--soil") + 1] if "--soil" in sys.argv else None

    z, meta = load_dem(path)
    slope_deg = describe(z, meta)

    # cap working resolution for tractable compute
    MAXW = 1024
    if max(z.shape) > MAXW:
        s = MAXW / max(z.shape)
        z = zoom(z, s, order=1)
        print(f"resampled to {z.shape} for processing")

    hn = (z - z.min()) / (np.ptp(z) + 1e-9)          # normalised for rendering
    acc = G.flow_accumulation(gaussian_filter(hn, 0.6))

    G.save_png(G.colorize(hn, accum=acc, stretch=True, hs_z=3.0), "dem_01_elevation_rivers.png")
    G.save_png(np.repeat(G.hillshade(hn, z=3.5)[..., None], 3, 2), "dem_02_hillshade.png")
    sd = (slope_deg - slope_deg.min()) / (np.ptp(slope_deg) + 1e-9)
    if sd.shape != hn.shape:
        sd = zoom(sd, (hn.shape[0]/sd.shape[0], hn.shape[1]/sd.shape[1]), order=1)
    G.save_png(np.stack([sd, 0.5*sd, 1 - sd], -1), "dem_03_slope.png")

    if soil:                                          # optional soil overlay
        sz, _ = load_dem(soil)
        sz = zoom(sz, (hn.shape[0]/sz.shape[0], hn.shape[1]/sz.shape[1]), order=1)
        sn = (sz - sz.min()) / (np.ptp(sz) + 1e-9)
        G.save_png(np.stack([0.75*sn+0.15, 0.55*sn+0.1, 0.3*(1-sn)+0.1], -1),
                   "dem_04_soil.png")
        print("wrote dem_04_soil.png")

    np.savez_compressed("baked_dem.npz", h=hn.astype(np.float32),
                        acc=acc.astype(np.float32), elev_m=z.astype(np.float32))

    if do_erode:
        print("running erosion on the real terrain...")
        h2, water, speed, vx, vy, _ = G.hydraulic_erosion(
            hn, iters=180, rain=0.03, evap=0.012, Ks=0.9, Kd=0.5, Kt=0.05,
            max_step=0.03, snapshot_every=0)
        acc2 = G.flow_accumulation(h2)
        G.save_png(G.colorize(h2, accum=acc2, stretch=True, hs_z=3.5),
                   "dem_05_eroded_rivers.png")
        np.savez_compressed("baked_dem_eroded.npz", h=h2, water=water,
                            vx=vx, vy=vy, speed=speed, acc=acc2)
        print("wrote dem_05_eroded_rivers.png")

    print("done — see dem_*.png and baked_dem.npz")
    if do_fly:
        print("for the flyover, point flyover.py at baked_dem*.npz")


if __name__ == "__main__":
    main()