"""DEM hydrology: streams follow a synthetic valley; the outlet's catchment is the whole valley."""
import numpy as np
import pytest
import rasterio
from pyproj import Transformer
from rasterio.transform import from_origin
from shapely.geometry import shape

from backend import dem as dem_store

EPSG = "EPSG:5186"
X0, Y1, CELL, W, H = 200_000.0, 550_000.0, 10.0, 81, 120  # valley axis at column 40


@pytest.fixture(autouse=True)
def store(tmp_path, monkeypatch):
    monkeypatch.setattr(dem_store, "DEM_DIR", tmp_path / "dem")


def valley(tmp_path):
    xs = X0 + (np.arange(W) + 0.5) * CELL
    ys = Y1 - (np.arange(H) + 0.5) * CELL
    xx, yy = np.meshgrid(xs, ys)
    axis = X0 + 40.5 * CELL
    z = 100 + 0.05 * np.abs(xx - axis) + 0.01 * (yy - (Y1 - H * CELL))  # V valley draining south
    z[30, 40] -= 3  # a pit on the axis must not stop the flow
    p = tmp_path / "v.tif"
    with rasterio.open(p, "w", driver="GTiff", width=W, height=H, count=1, dtype="float32", crs=EPSG,
                       transform=from_origin(X0, Y1, CELL, CELL)) as d:
        d.write(z.astype("float32"), 1)
    up = tmp_path / "up.bin"
    up.write_bytes(p.read_bytes())
    return dem_store.import_dem(str(up), "v.tif")


def test_streams_follow_the_axis_through_a_pit(tmp_path):
    info = valley(tmp_path)
    fc = dem_store.streams(info["id"], 0.05)  # 0.05 km2 = 500 cells
    assert fc["features"]
    to_xy = Transformer.from_crs("EPSG:4326", EPSG, always_xy=True)
    axis = X0 + 40.5 * CELL
    longest = max(fc["features"], key=lambda f: f["properties"]["length_m"])
    xs = [to_xy.transform(*c)[0] for c in longest["geometry"]["coordinates"]]
    assert max(abs(x - axis) for x in xs) < CELL * 1.5
    assert longest["properties"]["length_m"] > 0.5 * H * CELL
    assert longest["properties"]["area_km2"] == pytest.approx(W * H * CELL * CELL / 1e6, rel=0.1)


def test_watershed_of_the_outlet_is_the_whole_valley(tmp_path):
    info = valley(tmp_path)
    outlet = [X0 + 40.5 * CELL, Y1 - (H - 1) * CELL]
    fc = dem_store.watershed(info["id"], [outlet], EPSG, snap_m=30)
    (f,) = fc["features"]
    assert f["properties"]["area_km2"] == pytest.approx(W * H * CELL * CELL / 1e6, rel=0.05)
    assert shape(f["geometry"]).is_valid
    # A point halfway up drains only the upper part.
    mid = [X0 + 40.5 * CELL, Y1 - (H / 2) * CELL]
    half = dem_store.watershed(info["id"], [mid], EPSG, snap_m=30)["features"][0]["properties"]["area_km2"]
    assert 0.3 * f["properties"]["area_km2"] < half < 0.7 * f["properties"]["area_km2"]