"""Geological passes over a heightmap in metres, numpy only: particle hydraulic erosion, thermal weathering with
an angle of repose, strata hardness, and the derivative maps (flow, wear, deposition) they leave behind.
Fractal noise gives pillowy hills; these passes give drainage, V-valleys, alluvial fans, scree aprons and rock
shelves. Applied by generate_heightmap.py to whatever the manifest's source produced, noise or file.

Units inside: heights are in cell widths (metres over the cell size), so a slope of 1.0 is 45 degrees and the
droplet constants mean the same thing at any resolution. The hydraulic pass runs twice: on a downsampled map
(coarse cells, long droplet lives) for the valleys, then at full resolution (short lives) for the gullies;
the coarse result is applied to the full map as a delta, so the fine detail underneath survives.

Droplets are simulated in vectorised batches: a batch of tens of thousands takes one step together, reading
the map as it was at the start of the step and scattering its erosion and deposits back with np.add.at. Two
droplets in the same cell in the same step do not see each other; at these densities that is invisible.
"""
import time

import numpy as np

import heightmap_io
import heightmap_noise

DEFAULTS = {
    "enabled": True,
    "coarse_factor": 4,            # the coarse pass runs on the map downsampled by this
    "coarse_droplets": 800000,
    "coarse_lifetime": 120,        # steps, one cell each: 120 coarse cells is 1.7 km of path at 4x on 3.5 m quads
    "fine_droplets": 3000000,
    "fine_lifetime": 40,
    "thermal_passes": 24,
    "talus_deg": 35.0,             # angle of repose
    "inertia": 0.1,
    "capacity": 2.0,               # sediment a droplet can carry, in cell-heights per unit of slope, speed and water
    "max_load": 2.0,               # cell-heights: the most one droplet carries, so the mound it can leave where it stops is bounded
    "min_slope": 0.01,
    "deposit_rate": 0.2,
    "erode_rate": 0.2,
    "evaporation": 0.02,
    "gravity": 4.0,
    "strata_period_m": 160.0,      # vertical period of the hard and soft bands
    "strata_contrast": 0.6,        # 0 is uniform rock, 1 is hard bands that barely erode next to soft ones that melt
    "max_change": 0.2,             # cell-heights one droplet may cut or fill in one step; batches of droplets share cells, so this is the brake
    "max_speed": 5.0,
    "min_erode_slope": 0.25,       # below this slope (about 14 degrees) water deposits but barely cuts: lowland soil holds, so meadows stay meadows
    "fine_scale": 0.5,             # the fine pass cuts at this fraction of the coarse pass's rates: gullies, not trenches, at 3.5 m cells
    "batch": 200000,               # droplets stepping together, at most one per 40 cells of the map
    "seed": 11,
}


class Hardness:
    """Rock hardness in [0, 1] as a function of position and elevation: horizontal strata with a slow tilt and a
    slow change of rock type across the map. Erosion is scaled by (1 - hardness), so hard bands hold shelves."""

    def __init__(self, size, rng, period_cells, contrast):
        self.period = max(float(period_cells), 1e-3)
        self.contrast = float(contrast)
        self.tilt = heightmap_noise.fbm(size, rng, base_cells=3, octaves=3, gain=0.5).astype(np.float32)
        self.kind = heightmap_noise.fbm(size, rng, base_cells=2, octaves=3, gain=0.5).astype(np.float32)

    def at(self, ix, iy, height):
        band = 0.5 + 0.5 * np.sin(2.0 * np.pi * (height / self.period + self.tilt[iy, ix] * 2.0))
        return np.clip(0.5 + self.contrast * (band - 0.5) * (0.4 + 0.8 * self.kind[iy, ix]), 0.05, 0.95).astype(np.float32)


BRUSH = ((0, 0, 0.36), (0, 1, 0.12), (0, -1, 0.12), (1, 0, 0.12), (-1, 0, 0.12),
         (1, 1, 0.04), (1, -1, 0.04), (-1, 1, 0.04), (-1, -1, 0.04))   # offsets (dy, dx) and weights summing to 1


def sample(h, px, py):
    """Bilinear height and gradient at float positions; the caller keeps px, py inside [0, size - 2]."""
    x0 = px.astype(np.int32)
    y0 = py.astype(np.int32)
    fx = px - x0
    fy = py - y0
    h00 = h[y0, x0]
    h10 = h[y0, x0 + 1]
    h01 = h[y0 + 1, x0]
    h11 = h[y0 + 1, x0 + 1]
    gx = (h10 - h00) * (1 - fy) + (h11 - h01) * fy
    gy = (h01 - h00) * (1 - fx) + (h11 - h10) * fx
    hc = h00 * (1 - fx) * (1 - fy) + h10 * fx * (1 - fy) + h01 * (1 - fx) * fy + h11 * fx * fy
    return hc, gx, gy, x0, y0, fx, fy


def hydraulic(h, rng, droplets, lifetime, cfg, hardness, spawn_mask, maps, sea_cells=-1e9):
    """Particle erosion in place on h (cell units). maps: flow, wear, deposit arrays of h's shape, accumulated.
    A droplet that reaches water below `sea_cells` drops its whole load there and ends: the sea is a sink,
    and river mouths get their fans."""
    size = h.shape[0]
    ys, xs = np.nonzero(spawn_mask)
    if xs.size == 0:
        return
    inertia, capacity_factor = cfg["inertia"], cfg["capacity"]
    min_slope, deposit_rate, erode_rate = cfg["min_slope"], cfg["deposit_rate"], cfg["erode_rate"]
    evaporation, gravity = cfg["evaporation"], cfg["gravity"]
    max_change, max_speed, max_load = float(cfg["max_change"]), float(cfg["max_speed"]), float(cfg["max_load"])
    min_erode_slope = max(float(cfg["min_erode_slope"]), 1e-6)
    batch = max(min(int(cfg["batch"]), size * size // 40), 1000)
    limit = size - 2.001
    done = 0
    while done < droplets:
        n = min(batch, droplets - done)
        done += n
        pick = rng.integers(0, xs.size, n)
        px = np.clip(xs[pick] + rng.random(n, dtype=np.float32), 1.0, limit).astype(np.float32)
        py = np.clip(ys[pick] + rng.random(n, dtype=np.float32), 1.0, limit).astype(np.float32)
        dx = np.zeros(n, dtype=np.float32)
        dy = np.zeros(n, dtype=np.float32)
        speed = np.ones(n, dtype=np.float32)
        water = np.ones(n, dtype=np.float32)
        sediment = np.zeros(n, dtype=np.float32)

        for _ in range(lifetime):
            if px.size == 0:
                break
            hc, gx, gy, x0, y0, fx, fy = sample(h, px, py)
            dx = dx * inertia - gx * (1 - inertia)
            dy = dy * inertia - gy * (1 - inertia)
            length = np.hypot(dx, dy)
            moving = length > 1e-9
            safe = np.where(moving, length, 1.0)
            dx = np.where(moving, dx / safe, 0.0).astype(np.float32)
            dy = np.where(moving, dy / safe, 0.0).astype(np.float32)
            nx = px + dx
            ny = py + dy
            inside = moving & (nx >= 1.0) & (nx <= limit) & (ny >= 1.0) & (ny <= limit)
            hn = sample(h, np.clip(nx, 1.0, limit), np.clip(ny, 1.0, limit))[0]
            dh = np.where(inside, hn - hc, 0.0).astype(np.float32)

            slope = np.maximum(-dh, min_slope)
            capacity = np.minimum(slope * speed * water * capacity_factor, max_load)
            hard = hardness.at(x0, y0, hc) if hardness is not None else 0.0
            holds = np.clip(np.hypot(gx, gy) / min_erode_slope, 0.0, 1.0) ** 2   # flat ground resists cutting
            deposit = np.where(dh > 0.0, np.minimum(dh, sediment),
                               np.where(sediment > capacity, (sediment - capacity) * deposit_rate, 0.0))
            erode = np.where((dh <= 0.0) & (sediment <= capacity),
                             np.minimum((capacity - sediment) * erode_rate, -dh) * (1.0 - hard) * holds, 0.0)
            into_sea = inside & (hn < sea_cells)
            deposit = np.where(into_sea, sediment, np.minimum(deposit, max_change)).astype(np.float32)
            erode = np.where(into_sea, 0.0, np.minimum(erode, max_change)).astype(np.float32)
            # Cuts go through a 3x3 brush: a one-cell footprint leaves every path as a rill one cell wide, which
            # reads as brush strokes. Deposits land on the droplet's own bilinear cell: spread through the brush,
            # a pit's rim rises faster than its floor, the pit never fills, and every droplet that drains into it
            # adds to the rim until there is a mound.
            for oy, ox, weight in BRUSH:
                np.add.at(h, (y0 + oy, x0 + ox), -erode * weight)
            np.add.at(h, (y0, x0), deposit * (1 - fx) * (1 - fy))
            np.add.at(h, (y0, x0 + 1), deposit * fx * (1 - fy))
            np.add.at(h, (y0 + 1, x0), deposit * (1 - fx) * fy)
            np.add.at(h, (y0 + 1, x0 + 1), deposit * fx * fy)
            np.add.at(maps["flow"], (y0, x0), water)
            np.add.at(maps["wear"], (y0, x0), erode)
            np.add.at(maps["deposit"], (y0, x0), deposit)

            sediment = sediment + erode - deposit
            speed = np.minimum(np.sqrt(np.maximum(0.0, speed * speed - dh * gravity)), max_speed).astype(np.float32)   # downhill is faster
            water = water * (1.0 - evaporation)
            alive = inside & ~into_sea & (water > 0.001)
            px, py, dx, dy = nx[alive], ny[alive], dx[alive], dy[alive]
            speed, water, sediment = speed[alive], water[alive], sediment[alive]


DIRECTIONS = ((0, 1, 1.0), (0, -1, 1.0), (1, 0, 1.0), (-1, 0, 1.0),
              (1, 1, np.sqrt(2.0)), (1, -1, np.sqrt(2.0)), (-1, 1, np.sqrt(2.0)), (-1, -1, np.sqrt(2.0)))


def thermal(h, passes, talus):
    """Mass-conserving thermal weathering: where a cell stands above a neighbour by more than the angle of
    repose allows, half the excess slides down, shared among the lower neighbours. Cliffs keep a face, and
    scree builds at their feet. h in cell units, talus is tan(angle of repose)."""
    size = h.shape[0]

    def neighbour(padded, dy, dx):
        return padded[1 + dy:1 + dy + size, 1 + dx:1 + dx + size]

    for _ in range(passes):
        start = h.copy()
        padded = np.pad(start, 1, mode="edge")
        worst = np.zeros(h.shape, dtype=np.float32)
        total = np.zeros(h.shape, dtype=np.float32)
        for dy, dx, dist in DIRECTIONS:
            excess = np.maximum(start - neighbour(padded, dy, dx) - talus * dist, 0.0)
            worst = np.maximum(worst, excess)
            total += excess
        # A cell sheds half of its largest excess per pass, split among its lower neighbours in proportion to
        # how far each is below the angle of repose. Never more than half, so slopes settle without inverting.
        scale = np.where(total > 0.0, 0.5 * worst / np.maximum(total, 1e-9), 0.0).astype(np.float32)
        for dy, dx, dist in DIRECTIONS:
            move = np.maximum(start - neighbour(padded, dy, dx) - talus * dist, 0.0) * scale
            h -= move
            h[max(dy, 0):size + min(dy, 0), max(dx, 0):size + min(dx, 0)] += move[max(-dy, 0):size + min(-dy, 0), max(-dx, 0):size + min(-dx, 0)]
    return h


def erode(metres, quad_m, sea_level_m, settings, log=print):
    """The whole sequence on a map in metres. Returns (metres, maps) where maps holds flow, wear and deposit at
    the map's resolution, in cell-height units accumulated over both passes."""
    cfg = {**DEFAULTS, **(settings or {})}
    size = metres.shape[0]
    maps = {name: np.zeros((size, size), dtype=np.float32) for name in ("flow", "wear", "deposit")}
    if not cfg["enabled"]:
        return metres, maps
    rng = np.random.default_rng(int(cfg["seed"]))
    talus = float(np.tan(np.radians(cfg["talus_deg"])))
    started = time.time()

    factor = int(cfg["coarse_factor"])
    if factor > 1 and cfg["coarse_droplets"] > 0:
        coarse = heightmap_io.block_mean(metres, factor)
        cell_m = quad_m * factor
        hc = (coarse / cell_m).astype(np.float32)
        hardness = Hardness(hc.shape[0], rng, cfg["strata_period_m"] / cell_m, cfg["strata_contrast"])
        coarse_maps = {name: np.zeros(hc.shape, dtype=np.float32) for name in maps}
        thermal(hc, max(cfg["thermal_passes"] // 4, 1), talus)
        hydraulic(hc, rng, int(cfg["coarse_droplets"]), int(cfg["coarse_lifetime"]), cfg, hardness, coarse > sea_level_m + 2.0, coarse_maps,
                  sea_cells=sea_level_m / cell_m)
        thermal(hc, max(cfg["thermal_passes"] // 2, 1), talus)
        delta = hc * cell_m - coarse
        metres = (metres + heightmap_io.resample(delta, size)).astype(np.float32)
        for name in maps:
            maps[name] += heightmap_io.resample(coarse_maps[name], size) * factor
        log(f"  coarse erosion at {hc.shape[0]}x{hc.shape[0]}: {cfg['coarse_droplets']} droplets, "
            f"largest cut {-delta.min():.0f} m, largest fill {delta.max():.0f} m, {time.time() - started:.0f} s")

    hf = (metres / quad_m).astype(np.float32)
    before = hf.copy()
    hardness = Hardness(size, rng, cfg["strata_period_m"] / quad_m, cfg["strata_contrast"])
    fine_cfg = {**cfg, "erode_rate": cfg["erode_rate"] * cfg["fine_scale"], "max_change": cfg["max_change"] * cfg["fine_scale"]}
    hydraulic(hf, rng, int(cfg["fine_droplets"]), int(cfg["fine_lifetime"]), fine_cfg, hardness, metres > sea_level_m + 2.0, maps,
              sea_cells=sea_level_m / quad_m)
    thermal(hf, int(cfg["thermal_passes"]), talus)
    delta = (hf - before) * quad_m
    log(f"  fine erosion at {size}x{size}: {cfg['fine_droplets']} droplets, {cfg['thermal_passes']} thermal passes at "
        f"{cfg['talus_deg']:g} deg, largest cut {-delta.min():.0f} m, largest fill {delta.max():.0f} m, {time.time() - started:.0f} s total")
    return (hf * quad_m).astype(np.float32), maps


def curvature(metres, quad_m):
    """Laplacian of the lightly blurred height, in metres per cell squared: positive on ridges and convex
    shoulders, negative in gullies and sediment traps."""
    h = heightmap_noise.box_blur(metres, 2)
    padded = np.pad(h, 1, mode="edge")
    lap = (padded[:-2, 1:-1] + padded[2:, 1:-1] + padded[1:-1, :-2] + padded[1:-1, 2:] - 4.0 * h)
    return (lap / quad_m).astype(np.float32)


def to_unit(values, percentile=99.0, log_scale=False):
    """A map squashed into [0, 1] for painting and for an 8-bit PNG."""
    v = np.log1p(np.maximum(values, 0.0)) if log_scale else np.maximum(values, 0.0)
    top = float(np.percentile(v, percentile))
    return np.clip(v / max(top, 1e-6), 0.0, 1.0).astype(np.float32)