#!/usr/bin/env python3
"""Read a plain float32 GeoTIFF, report its georeferencing and elevation stats.

Usage:
    Geo/tools/dem_stats.py Geo/raw/dem-3dep-1m-26911-2026-11-06.tif
    Geo/tools/dem_stats.py <tif> --sample 376515.389,3774224.083

Why this exists: GDAL is not runnable on the host (its flatpak libraries do not
resolve outside the sandbox), so this reads the strip-based GeoTIFF that the
3DEP export service produces with the standard library only. It handles
uncompressed, single-sample, strip-organised float32 images with the usual
GeoTIFF tags; it is a small reader, not a general GeoTIFF implementation.

Values are orthometric metres (NAVD88 for 3DEP CONUS). Nodata is excluded from
the statistics.
"""

import argparse
import struct
import sys

TYPE_SIZES = {1: 1, 2: 1, 3: 2, 4: 4, 5: 8, 7: 1, 11: 4, 12: 8, 16: 8}


def read_tiff(path):
    with open(path, "rb") as handle:
        blob = handle.read()

    if blob[:2] == b"II":
        endian = "<"
    elif blob[:2] == b"MM":
        endian = ">"
    else:
        raise SystemExit("not a TIFF file")
    if struct.unpack(endian + "H", blob[2:4])[0] != 42:
        raise SystemExit("not a classic TIFF (BigTIFF is unsupported)")

    ifd_offset = struct.unpack(endian + "I", blob[4:8])[0]
    count = struct.unpack(endian + "H", blob[ifd_offset : ifd_offset + 2])[0]
    entries = {}
    for i in range(count):
        base = ifd_offset + 2 + i * 12
        tag, typ, n = struct.unpack(endian + "HHI", blob[base : base + 8])
        size = TYPE_SIZES.get(typ, 1) * n
        raw = blob[base + 8 : base + 8 + 4] if size <= 4 else blob[
            struct.unpack(endian + "I", blob[base + 8 : base + 12])[0] :
        ][:size]
        fmt = {1: "B", 3: "H", 4: "I", 7: "B", 11: "f", 12: "d", 2: "c"}.get(typ, "B")
        if typ == 2:
            values = raw[:size].split(b"\0")[0].decode("latin-1")
        elif fmt == "c":
            values = raw
        else:
            values = list(struct.unpack(endian + fmt * n, raw[: struct.calcsize(endian + fmt * n)]))
        entries[tag] = values

    width = entries[256][0]
    height = entries[257][0]
    bits = entries.get(258, [8])[0]
    compression = entries.get(259, [1])[0]
    sample_format = entries.get(339, [1])[0]
    if compression != 1 or bits != 32 or sample_format != 3:
        raise SystemExit(
            f"unsupported TIFF layout: compression={compression} bits={bits} sample_format={sample_format}"
        )

    if 322 in entries:
        # tiled layout: tags 322/323 give tile size, 324/325 the offsets and byte counts
        tile_w = entries[322][0]
        tile_h = entries[323][0]
        offsets = entries[324]
        counts = entries[325]
        across = (width + tile_w - 1) // tile_w
        down = (height + tile_h - 1) // tile_h
        if len(offsets) < across * down:
            raise SystemExit(f"tiled TIFF with {len(offsets)} tiles, expected {across * down}")
        grid = [[0.0] * width for _ in range(height)]
        for index, (offset, nbytes) in enumerate(zip(offsets, counts)):
            tile_row, tile_col = divmod(index, across)
            tile = struct.unpack(endian + "f" * (tile_w * tile_h), blob[offset : offset + nbytes])
            for r in range(tile_h):
                target_row = tile_row * tile_h + r
                if target_row >= height:
                    break
                base = r * tile_w
                for c in range(tile_w):
                    target_col = tile_col * tile_w + c
                    if target_col >= width:
                        break
                    grid[target_row][target_col] = tile[base + c]
        flat = [value for row in grid for value in row]
        return entries, width, height, flat

    # strip layout
    offsets = entries[273]
    counts = entries[279]
    rows_per_strip = entries.get(278, [height])[0]
    chunks = [blob[offset : offset + nbytes] for offset, nbytes in zip(offsets, counts)]
    pixels = b"".join(chunks)
    values = list(struct.unpack(endian + "f" * (width * height), pixels[: width * height * 4]))
    return entries, width, height, values


def geokeys(entries):
    keys = {}
    raw = entries.get(34735)
    if not raw:
        return keys
    # GeoKeyDirectory: version, revision, minor, count, then 4 shorts per key
    for i in range(4, len(raw), 4):
        key, tag_location, count, value = raw[i : i + 4]
        keys[key] = (tag_location, count, value)
    return keys


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("tif")
    ap.add_argument("--sample", action="append", default=[], help="E,N to sample, repeatable")
    args = ap.parse_args()

    entries, width, height, values = read_tiff(args.tif)
    scale = entries.get(33550)
    tiepoint = entries.get(33922)
    keys = geokeys(entries)
    nodata = None
    if 42113 in entries:
        try:
            nodata = float(entries[42113])
        except ValueError:
            nodata = None

    print(f"size: {width} x {height} px, float32, uncompressed")
    if nodata is not None:
        print(f"nodata: {nodata}")

    valid = [v for v in values if nodata is None or v != nodata]
    if valid:
        print(
            f"elevation: min {min(valid):.2f} m, max {max(valid):.2f} m, "
            f"relief {max(valid) - min(valid):.2f} m, mean {sum(valid)/len(valid):.2f} m"
        )
        print(f"valid pixels: {len(valid)} of {width*height}")

    if scale and tiepoint:
        sx, sy = scale[0], scale[1]
        px, py, pz, ex, ey = tiepoint[0], tiepoint[1], tiepoint[2], tiepoint[3], tiepoint[4]
        x0 = ex - px * sx
        y0 = ey + py * sy  # north-up: tiepoint row 0 is the top edge
        print(f"pixel size: {sx} x {sy} m")
        print(f"extent: E {x0:.1f} .. {x0 + width*sx:.1f}, N {y0 - height*sy:.1f} .. {y0:.1f}")
        crs = keys.get(3072, (None, None, None))[2]
        model = keys.get(1024, (None, None, None))[2]
        print(f"model type: {model} (1 = projected), projected CRS: EPSG:{crs}")

        for spec in args.sample:
            try:
                e, n = (float(x) for x in spec.split(","))
            except ValueError:
                raise SystemExit(f"bad --sample value: {spec!r}")
            col = int((e - x0) / sx)
            row = int((y0 - n) / sy)
            if not (0 <= col < width and 0 <= row < height):
                print(f"  sample E{e:.1f} N{n:.1f}: outside the raster")
                continue
            value = values[row * width + col]
            print(f"  sample E{e:.3f} N{n:.3f} -> row {row}, col {col}: {value:.3f} m")
    else:
        print("no pixel scale or tiepoint found; georeferencing missing", file=sys.stderr)


if __name__ == "__main__":
    main()