"""Terrain analysis with rasterio and rioxarray - TechnicalGIS tutorial.

Streams part of a Copernicus GLO-30 DEM Cloud Optimised GeoTIFF, clips it
to the Peak District, reprojects to British National Grid and derives
hillshade, slope and aspect plus zonal statistics.
Run with:  python terrain_analysis.py
"""
from pathlib import Path

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio
import rioxarray  # noqa: F401  (registers the .rio accessor)
import xarray as xr
from matplotlib.colors import LightSource
from pyproj import Transformer
from rasterio.enums import Resampling
from rasterio.features import rasterize
from shapely.geometry import Point, box

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

DEM_URL = ("https://copernicus-dem-30m.s3.amazonaws.com/"
           "Copernicus_DSM_COG_10_N53_00_W002_00_DEM/Copernicus_DSM_COG_10_N53_00_W002_00_DEM.tif")
AOI = (-1.95, 53.30, -1.50, 53.52)  # lon/lat bounding box: Dark Peak to Wharncliffe


def section(title: str) -> None:
    print(f"\n=== {title} ===")


# --- Step 1: inspect the remote COG without downloading it ------------------
with rasterio.open(DEM_URL) as src:
    section("remote COG metadata")
    print("driver:", src.driver)
    print("size:", src.width, "x", src.height, "| bands:", src.count, "| dtype:", src.dtypes[0])
    print("crs:", src.crs)
    print("resolution (deg):", tuple(round(r, 6) for r in src.res))
    print("bounds:", src.bounds)
    print("overviews:", src.overviews(1))
    print("block shape:", src.block_shapes[0])

# --- Step 2: open lazily with rioxarray and clip to the AOI -----------------
dem = rioxarray.open_rasterio(DEM_URL, masked=True).squeeze("band", drop=True)
dem_aoi = dem.rio.clip_box(*AOI).load()
section("clipped DataArray")
print(dem_aoi)

# --- Step 3: reproject to British National Grid at 30 m ---------------------
dem_bng = dem_aoi.rio.reproject("EPSG:27700", resolution=30, resampling=Resampling.bilinear,
                                nodata=np.nan)
# A lat/lon rectangle becomes slightly skewed in BNG, so trim to the largest
# axis-aligned rectangle that is fully covered by data.
to_bng = Transformer.from_crs("EPSG:4326", "EPSG:27700", always_xy=True)
cx, cy = to_bng.transform([AOI[0], AOI[0], AOI[2], AOI[2]], [AOI[1], AOI[3], AOI[1], AOI[3]])
dem_bng = dem_bng.rio.clip_box(max(cx[0], cx[1]) + 60, max(cy[0], cy[2]) + 60,
                               min(cx[2], cx[3]) - 60, min(cy[1], cy[3]) - 60)
dem_bng.name = "elevation"
section("reprojected")
print("crs:", dem_bng.rio.crs, "| shape:", dem_bng.shape, "| res:", dem_bng.rio.resolution())
print("elevation min/mean/max (m):", *np.round([float(dem_bng.min()), float(dem_bng.mean()), float(dem_bng.max())], 1))
dem_bng.rio.to_raster(OUT / "dem_peak_district_bng.tif", compress="deflate")

# --- Step 4: slope, aspect and hillshade with NumPy ---------------------------
z = dem_bng.values.astype("float64")
res = abs(dem_bng.rio.resolution()[0])
dz_dy, dz_dx = np.gradient(z, res)          # rows increase southwards
dz_dy = -dz_dy                               # make +y point north
slope = np.degrees(np.arctan(np.hypot(dz_dx, dz_dy)))
aspect = (np.degrees(np.arctan2(-dz_dx, -dz_dy)) + 360) % 360  # 0 = north-facing, clockwise

slope_da = dem_bng.copy(data=slope.astype("float32"))
aspect_da = dem_bng.copy(data=aspect.astype("float32"))
slope_da.rio.to_raster(OUT / "slope_deg.tif", compress="deflate")
aspect_da.rio.to_raster(OUT / "aspect_deg.tif", compress="deflate")

ls = LightSource(azdeg=315, altdeg=45)
hillshade = ls.hillshade(np.nan_to_num(z, nan=np.nanmean(z)), vert_exag=2, dx=res, dy=res)

section("slope statistics (degrees)")
print(pd.Series(slope[~np.isnan(slope)]).describe(percentiles=[0.5, 0.9, 0.99]).round(2).to_string())

extent = [float(dem_bng.x.min()), float(dem_bng.x.max()), float(dem_bng.y.min()), float(dem_bng.y.max())]

fig, ax = plt.subplots(figsize=(9, 5.6))
ax.imshow(hillshade, cmap="gray", extent=extent)
im = ax.imshow(z, cmap="terrain", alpha=0.55, extent=extent, vmin=0, vmax=650)
fig.colorbar(im, ax=ax, label="Elevation (m)", shrink=0.8)
ax.set_title("Copernicus GLO-30 DEM over the Dark Peak, hillshade × elevation")
ax.set_xlabel("Easting (m, EPSG:27700)")
ax.set_ylabel("Northing (m)")
fig.savefig(OUT / "01-dem-hillshade.png")
plt.close(fig)

fig, axes = plt.subplots(1, 2, figsize=(11, 3.9))
im0 = axes[0].imshow(slope, cmap="magma_r", vmin=0, vmax=35, extent=extent)
fig.colorbar(im0, ax=axes[0], label="Slope (°)", shrink=0.85)
axes[0].set_title("Slope")
im1 = axes[1].imshow(aspect, cmap="twilight", vmin=0, vmax=360, extent=extent)
fig.colorbar(im1, ax=axes[1], label="Aspect (° clockwise from north)", shrink=0.85)
axes[1].set_title("Aspect")
for ax in axes:
    ax.set_xticks([])
    ax.set_yticks([])
fig.savefig(OUT / "02-slope-aspect.png")
plt.close(fig)

# --- Step 5: zonal statistics on a 2 km grid -----------------------------------
x0, y0, x1, y1 = dem_bng.rio.bounds()
size = 2000
cells = [box(x, y, x + size, y + size)
         for x in np.arange(np.ceil(x0 / size) * size, x1 - size, size)
         for y in np.arange(np.ceil(y0 / size) * size, y1 - size, size)]
grid = gpd.GeoDataFrame({"cell_id": range(1, len(cells) + 1)}, geometry=cells, crs="EPSG:27700")

zones = rasterize(zip(grid.geometry, grid.cell_id), out_shape=z.shape,
                  transform=dem_bng.rio.transform(), fill=0, dtype="int32")
df = pd.DataFrame({"cell_id": zones.ravel(), "elev": z.ravel(), "slope": slope.ravel()})
df = df[(df.cell_id > 0) & df.elev.notna()]
stats = df.groupby("cell_id").agg(elev_mean=("elev", "mean"), elev_max=("elev", "max"),
                                  slope_mean=("slope", "mean"), steep_pct=("slope", lambda s: (s > 20).mean() * 100),
                                  n_pixels=("elev", "size"))
grid = grid.merge(stats.reset_index(), on="cell_id")
section("zonal statistics (first rows)")
print(grid.drop(columns="geometry").head().round(1).to_string(index=False))
section("steepest cells")
print(grid.nlargest(5, "slope_mean").drop(columns="geometry").round(1).to_string(index=False))
grid.to_file(OUT / "zonal_stats_2km.gpkg", driver="GPKG")

fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
for ax, col, cmap, label in [(axes[0], "elev_mean", "YlGn", "Mean elevation (m)"),
                             (axes[1], "steep_pct", "OrRd", "% of cell steeper than 20°")]:
    ax.imshow(hillshade, cmap="gray", extent=extent, alpha=0.6)
    grid.plot(column=col, cmap=cmap, alpha=0.75, edgecolor="white", linewidth=0.3, ax=ax, legend=True,
              legend_kwds={"label": label, "shrink": 0.8})
    ax.set_xticks([])
    ax.set_yticks([])
axes[0].set_title("Zonal mean elevation, 2 km grid")
axes[1].set_title("Zonal share of steep ground, 2 km grid")
fig.savefig(OUT / "03-zonal-stats.png")
plt.close(fig)

# --- Step 6: sample summits ---------------------------------------------------
summits = gpd.GeoDataFrame(
    {"name": ["Kinder Scout", "Bleaklow", "Mam Tor", "Win Hill", "Stanage Edge (High Neb)"],
     "published_m": [636, 633, 517, 462, 458]},
    geometry=[Point(-1.8770, 53.3846), Point(-1.8584, 53.4605), Point(-1.8095, 53.3491),
              Point(-1.7215, 53.3622), Point(-1.6489, 53.3654)], crs="EPSG:4326").to_crs("EPSG:27700")

point_vals, max_vals = [], []
for geom in summits.geometry:
    point_vals.append(float(dem_bng.sel(x=geom.x, y=geom.y, method="nearest")))
    window = dem_bng.rio.clip([geom.buffer(300)], crs="EPSG:27700")
    max_vals.append(float(window.max()))
summits["dem_at_point"] = np.round(point_vals, 1)
summits["dem_max_300m"] = np.round(max_vals, 1)
section("summit check")
print(summits.drop(columns="geometry").to_string(index=False))

xs = dem_bng.sel(y=summits.geometry.iloc[2].y, method="nearest")  # west-east profile through Mam Tor
fig, ax = plt.subplots(figsize=(9, 2.8))
ax.fill_between((xs.x - xs.x.min()) / 1000, xs.values, color="#99c5b5", alpha=0.6)
ax.plot((xs.x - xs.x.min()) / 1000, xs.values, color="#0f766e", lw=1)
ax.set_xlabel("Distance east (km)")
ax.set_ylabel("Elevation (m)")
ax.set_title("West–east profile through Mam Tor and the Hope Valley")
ax.grid(alpha=0.3)
fig.savefig(OUT / "04-profile.png")
plt.close(fig)