"""LiDAR point clouds in Python with laspy: stream USGS 3DEP lidar, read a LAZ file,
remove noise, build a DTM, DSM and canopy height model (CHM), and measure the
Washington Monument.

Data: USGS 3DEP lidar (MD_VA_Sandy_NCR_2014), streamed as Cloud Optimised Point
Cloud (COPC) files from the Microsoft Planetary Computer. No account needed.
Run with:  python lidar_dtm_chm.py
"""
import platform
import time
from importlib.metadata import version
from pathlib import Path

import laspy
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import planetary_computer
import pystac_client
import rasterio
import scipy
from matplotlib.colors import BoundaryNorm, LightSource, ListedColormap
from pyproj import Transformer
from rasterio.crs import CRS
from rasterio.transform import from_origin
from scipy import ndimage
from scipy.interpolate import griddata
from scipy.spatial import cKDTree

HERE = Path(__file__).parent
OUT = HERE / "output"
DATA = HERE / "data"
OUT.mkdir(exist_ok=True)
DATA.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})

print("=== versions ===")
print(f"python {platform.python_version()} | laspy {laspy.__version__} | lazrs {version('lazrs')} | "
      f"numpy {np.__version__} | scipy {scipy.__version__} | rasterio {rasterio.__version__}")

# ---------------------------------------------------------------- 1. find tiles
LON, LAT = -77.0353, 38.8895  # Washington Monument
EPSG = 26918                  # NAD83 / UTM zone 18N, the CRS of this lidar project
to_utm = Transformer.from_crs(4326, EPSG, always_xy=True)
to_ll = Transformer.from_crs(EPSG, 4326, always_xy=True)
mx, my = to_utm.transform(LON, LAT)
xmin, xmax, ymin, ymax = mx - 500, mx + 500, my - 400, my + 400
lon0, lat0 = to_ll.transform(xmin, ymin)
lon1, lat1 = to_ll.transform(xmax, ymax)

catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                    modifier=planetary_computer.sign_inplace)
items = list(catalog.search(collections=["3dep-lidar-copc"], bbox=[lon0, lat0, lon1, lat1]).items())
print("\n=== STAC search ===")
print(f"area of interest: {xmax - xmin:.0f} m x {ymax - ymin:.0f} m centred on the Washington Monument")
for it in items:
    print(f"{it.id}  points={it.properties['pc:count']:,}")

# ------------------------------------------------- 2. stream only the points we need
print("\n=== streaming COPC ===")
bounds = laspy.Bounds(mins=np.array([xmin, ymin]), maxs=np.array([xmax, ymax]))
parts, header = [], None
t = time.perf_counter()
for it in items:
    with laspy.CopcReader.open(it.assets["data"].href) as reader:
        if header is None:
            header = reader.header
            print(f"LAS {header.version}, point format {header.point_format.id}, "
                  f"CRS: {header.parse_crs().name}")
        pts = reader.query(bounds)
        print(f"{it.id[-20:]}: {len(pts):,} points in the area of interest")
        if len(pts):
            parts.append(pts)
print(f"streamed in {time.perf_counter() - t:.1f} s")

# Merge into one LasData and save it as an ordinary LAZ file.
# Each tile has its own scale/offset, so copy real-world x/y/z rather than the raw integers.
out_header = laspy.LasHeader(point_format=header.point_format.id, version=header.version)
out_header.offsets = np.array([xmin, ymin, 0.0])
out_header.scales = np.array([0.01, 0.01, 0.01])
out_header.add_crs(header.parse_crs())
merged = laspy.LasData(out_header)
merged.points = laspy.ScaleAwarePointRecord.zeros(sum(len(p) for p in parts), header=out_header)
for dim in header.point_format.dimension_names:
    if dim not in ("X", "Y", "Z"):
        merged[dim] = np.concatenate([np.asarray(p[dim]) for p in parts])
merged.x = np.concatenate([np.asarray(p.x) for p in parts])
merged.y = np.concatenate([np.asarray(p.y) for p in parts])
merged.z = np.concatenate([np.asarray(p.z) for p in parts])
laz_path = DATA / "national_mall.laz"
merged.write(laz_path)
print(f"wrote data/{laz_path.name} ({laz_path.stat().st_size / 1e6:.1f} MB)")

# ---------------------------------------------------------------- 3. read and inspect
print("\n=== laspy.read ===")
las = laspy.read(laz_path)
print(f"points: {las.header.point_count:,}")
print(f"dimensions: {list(las.point_format.dimension_names)}")
print(f"x {las.x.min():.2f}..{las.x.max():.2f}  y {las.y.min():.2f}..{las.y.max():.2f}  "
      f"z {las.z.min():.2f}..{las.z.max():.2f}")
area = (xmax - xmin) * (ymax - ymin)
print(f"density: {las.header.point_count / area:.1f} points/m²")

# Codes 17, 18 and 25 are this project's overlap classes: the standard class + 16.
NAMES = {1: "Unclassified", 2: "Ground", 7: "Noise", 9: "Water", 10: "Ignored ground",
         11: "Withheld", 17: "Overlap unclassified", 18: "Overlap ground", 25: "Overlap water"}
x, y, z = np.asarray(las.x), np.asarray(las.y), np.asarray(las.z)
cls = np.asarray(las.classification)
rn, nr = np.asarray(las.return_number), np.asarray(las.number_of_returns)
print("\nclass     points   share   z p1..p99 (m)   name")
for c, n in zip(*np.unique(cls, return_counts=True)):
    zp = np.percentile(z[cls == c], [1, 99])
    print(f"{c:>5}  {n:>9,}  {n / len(cls):6.1%}   {zp[0]:6.1f}..{zp[1]:6.1f}   {NAMES.get(c, '?')}")
print(f"\npoints from pulses with more than one return: {(nr > 1).mean():.1%}")

# Chunked reading: how you would process a file too big for memory.
counts = np.zeros(256, dtype=np.int64)
with laspy.open(laz_path) as f:
    for chunk in f.chunk_iterator(1_000_000):
        counts += np.bincount(np.asarray(chunk.classification), minlength=256)
print(f"chunked pass agrees with laspy.read: "
      f"{bool((counts[counts > 0] == np.unique(cls, return_counts=True)[1]).all())}")

GROUND = np.isin(cls, [2, 18])
WATER = np.isin(cls, [9, 25])

# ---------------------------------------------------------------- 4. remove noise
print("\n=== statistical outlier removal ===")
t = time.perf_counter()
tree = cKDTree(np.c_[x, y, z])
dist, _ = tree.query(np.c_[x, y, z], k=9, workers=-1)  # nearest 8 neighbours (+ the point itself)
mean_d = dist[:, 1:].mean(axis=1)
limit = mean_d.mean() + 3 * mean_d.std()
NOISE = (cls == 7) | (mean_d > limit)
print(f"mean neighbour distance {mean_d.mean():.2f} m, limit {limit:.2f} m, took {time.perf_counter() - t:.1f} s")
print(f"outliers: {NOISE.sum():,} ({NOISE.mean():.2%}); of the {(z > 200).sum()} points above 200 m, "
      f"{(NOISE & (z > 200)).sum()} are outliers")
keep = ~NOISE

# ---------------------------------------------------------------- 5. figures of the cloud
GROUPS = [("Ground (2, 18)", [2, 18], "#c9a66b"), ("Water (9, 25)", [9, 25], "#2c7fb8"),
          ("Unclassified: trees, buildings (1, 17)", [1, 17], "#2f7d3c"), ("Other (10, 11)", [10, 11], "#d7301f")]
rng = np.random.default_rng(0)
sub = rng.choice(len(x), size=min(1_000_000, len(x)), replace=False)
sub = sub[keep[sub]]
fig, ax = plt.subplots(figsize=(9, 7.4))
for label, codes, colour in GROUPS:
    m = np.isin(cls[sub], codes)
    ax.scatter(x[sub][m] - xmin, y[sub][m] - ymin, s=0.12, c=colour, label=label, rasterized=True, linewidths=0)
ax.set_aspect("equal")
ax.set_xlim(0, xmax - xmin)
ax.set_ylim(0, ymax - ymin)
ax.set_xlabel("metres east")
ax.set_ylabel("metres north")
ax.set_title("USGS 3DEP lidar around the Washington Monument, coloured by class")
ax.legend(markerscale=25, loc="upper left", fontsize=8, framealpha=0.95)
fig.savefig(OUT / "01-classification.png")
plt.close(fig)

band = (np.abs(y - my) < 2.0) & keep
fig, ax = plt.subplots(figsize=(10, 4))
for label, codes, colour in GROUPS[:3]:
    m = band & np.isin(cls, codes)
    ax.scatter(x[m] - mx, z[m], s=0.6, c=colour, label=label, linewidths=0)
ax.set_xlabel("metres east of the Monument")
ax.set_ylabel("elevation (m, NAVD88)")
ax.set_title("4 m-wide east–west slice through the point cloud")
ax.legend(markerscale=10, fontsize=8, loc="upper left")
ax.grid(alpha=0.3)
fig.savefig(OUT / "02-profile.png")
plt.close(fig)

# ---------------------------------------------------------------- 6. rasterise
print("\n=== rasterise to 1 m ===")
RES = 1.0
ncol, nrow = int((xmax - xmin) / RES), int((ymax - ymin) / RES)
col = np.clip(((x - xmin) / RES).astype(int), 0, ncol - 1)
row = np.clip(((ymax - y) / RES).astype(int), 0, nrow - 1)
cell = row * ncol + col

# DSM: highest non-noise return in each cell.
dsm = np.full(nrow * ncol, -np.inf)
np.maximum.at(dsm, cell[keep], z[keep])
dsm[np.isinf(dsm)] = np.nan
dsm = dsm.reshape(nrow, ncol)

# DTM: mean of ground and water returns (water is the "ground" of a lake) in each cell.
terrain = (GROUND | WATER) & keep
gsum = np.bincount(cell[terrain], weights=z[terrain], minlength=nrow * ncol)
gcnt = np.bincount(cell[terrain], minlength=nrow * ncol)
dtm = np.where(gcnt > 0, gsum / np.maximum(gcnt, 1), np.nan).reshape(nrow, ncol)
print(f"DSM cells with no return: {np.isnan(dsm).mean():.1%}")
print(f"DTM cells with no ground point: {np.isnan(dtm).mean():.1%}")

t = time.perf_counter()
rr, cc = np.nonzero(~np.isnan(dtm))
hr, hc = np.nonzero(np.isnan(dtm))
dtm[hr, hc] = griddata((rr, cc), dtm[rr, cc], (hr, hc), method="linear")
for arr in (dtm, dsm):  # anything left (edges, outside the convex hull): nearest neighbour
    holes = np.isnan(arr)
    if holes.any():
        idx = ndimage.distance_transform_edt(holes, return_distances=False, return_indices=True)
        arr[holes] = arr[tuple(idx)][holes]
print(f"gap filling: {time.perf_counter() - t:.1f} s")

ndsm = np.clip(dsm - dtm, 0, None)
print(f"DTM {dtm.min():.1f}..{dtm.max():.1f} m | DSM {dsm.min():.1f}..{dsm.max():.1f} m | "
      f"height above ground max {ndsm.max():.1f} m")

# ---------------------------------------------------------------- 7. canopy height model
print("\n=== canopy height model ===")
height_pt = z - dtm.ravel()[cell]
tall = keep & ~GROUND & ~WATER & (height_pt > 2)
multi = np.bincount(cell[tall & (nr > 1)], minlength=nrow * ncol)
alln = np.bincount(cell[tall], minlength=nrow * ncol)
frac = np.where(alln > 0, multi / np.maximum(alln, 1), 0).reshape(nrow, ncol)
elevated = ndsm > 2
# Laser pulses split into several returns in tree crowns but not on solid roofs.
veg = ndimage.binary_opening(elevated & (frac >= 0.5), iterations=1)
struct = ndimage.binary_opening(elevated & (frac < 0.5), iterations=1)
chm = np.where(veg, ndsm, 0.0)
print(f"cells more than 2 m above ground: {elevated.mean():.1%}")
print(f"  mostly multi-return (trees): {veg.mean():.1%} | mostly single-return (buildings, structures): "
      f"{struct.mean():.1%}")
canopy = ndsm[veg]
print(f"tree canopy height: median {np.median(canopy):.1f} m, 95th percentile {np.percentile(canopy, 95):.1f} m, "
      f"max {canopy.max():.1f} m")
labels, _ = ndimage.label(veg)
sizes = np.bincount(labels.ravel())[1:]
print(f"canopy patches of 20 m² or more: {(sizes >= 20).sum()}")

transform = from_origin(xmin, ymax, RES, RES)
crs = CRS.from_wkt(header.parse_crs().to_wkt())
for name, arr in [("dtm", dtm), ("dsm", dsm), ("chm", chm)]:
    with rasterio.open(OUT / f"{name}.tif", "w", driver="GTiff", height=nrow, width=ncol, count=1,
                       dtype="float32", crs=crs, transform=transform, compress="deflate", nodata=-9999) as dst:
        dst.write(arr.astype("float32"), 1)
print("wrote output/dtm.tif, output/dsm.tif, output/chm.tif")

ls = LightSource(azdeg=315, altdeg=45)
ext = [0, xmax - xmin, 0, ymax - ymin]
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6), constrained_layout=True)
for ax, arr, title in [(axes[0], dtm, "DTM: bare earth (ground and water points)"),
                       (axes[1], dsm, "DSM: highest return in each 1 m cell")]:
    shade = ls.shade(np.clip(arr, -2, 40), cmap=plt.cm.gist_earth, vert_exag=2, blend_mode="soft",
                     vmin=-2, vmax=40)
    ax.imshow(shade, extent=ext)
    ax.set_title(title)
    ax.set_xlabel("metres east")
axes[0].set_ylabel("metres north")
fig.savefig(OUT / "03-dtm-dsm.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8.5, 6.4), constrained_layout=True)
im = ax.imshow(np.clip(ndsm, 0, 40), extent=ext, cmap="viridis")
fig.colorbar(im, ax=ax, shrink=0.85, label="height above ground (m), clipped at 40 m")
ax.annotate("Washington Monument (169 m)", xy=(mx - xmin, my - ymin), xytext=(mx - xmin - 330, my - ymin + 150),
            color="white", fontsize=9, arrowprops=dict(arrowstyle="->", color="white"))
ax.annotate("construction cranes", xy=(720, 560), xytext=(560, 700), color="white", fontsize=9,
            arrowprops=dict(arrowstyle="->", color="white"))
ax.set_title("Normalised DSM (DSM − DTM): height of everything above the ground")
ax.set_xlabel("metres east")
ax.set_ylabel("metres north")
fig.savefig(OUT / "04-ndsm.png")
plt.close(fig)

fig, axes = plt.subplots(1, 2, figsize=(11, 4.6), width_ratios=[1.6, 1], constrained_layout=True)
fig.get_layout_engine().set(wspace=0.12)
lab = np.zeros(ndsm.shape, dtype=int)
lab[struct] = 1
axes[0].imshow(lab, extent=ext, interpolation="nearest", cmap=ListedColormap(["#f2efe9", "#9a9a9a"]),
               norm=BoundaryNorm([-0.5, 0.5, 1.5], 2))
im = axes[0].imshow(np.where(veg, chm, np.nan), extent=ext, cmap="YlGn", vmin=0, vmax=30, interpolation="nearest")
fig.colorbar(im, ax=axes[0], shrink=0.85, label="canopy height (m)")
axes[0].set_title("Canopy height model (buildings and structures in grey)")
axes[0].set_xlabel("metres east")
axes[0].set_ylabel("metres north")
axes[1].hist(canopy, bins=np.arange(2, 36, 1), color="#31a354", edgecolor="white", linewidth=0.4)
axes[1].axvline(np.median(canopy), color="#222", ls="--", lw=1)
axes[1].text(np.median(canopy) + 0.6, axes[1].get_ylim()[1] * 0.92, f"median {np.median(canopy):.1f} m", fontsize=8)
axes[1].set_xlabel("canopy height (m)")
axes[1].set_ylabel("number of 1 m cells")
axes[1].set_title("Distribution of tree heights")
axes[1].grid(axis="y", alpha=0.3)
fig.savefig(OUT / "05-canopy-height-model.png")
plt.close(fig)

# ---------------------------------------------------------------- 8. measure the Monument
print("\n=== Washington Monument ===")
d = np.hypot(x - mx, y - my)
near = (d < 10) & keep
top_i = np.argmax(np.where(near, z, -np.inf))
ring = GROUND & keep & (d > 15) & (d < 30)
base = np.median(z[ring])
print(f"highest point within 10 m: z={z[top_i]:.2f} m at {d[top_i]:.1f} m from the centre "
      f"(class {cls[top_i]}, return {rn[top_i]} of {nr[top_i]})")
print(f"ground 15–30 m from the centre: median z={base:.2f} m from {ring.sum():,} ground points")
height = z[top_i] - base
print(f"lidar height: {height:.2f} m ({height / 0.3048:.1f} ft)")
OFFICIAL = 169.046  # NGS 2013–14 survey: 554 ft 7 11/32 in
print(f"official height (NGS survey): {OFFICIAL} m; difference {height - OFFICIAL:+.2f} m")