"""Writes the heightmap and the three weightmaps of L_World to RawContent/World/Heightmaps/ from the manifest
RawContent/World/World.json: 16-bit greyscale height, 8-bit greyscale weights, sized for the landscape.

Pure numpy, no engine: run it with the engine's Python (numpy lives in Scripts/Authoring/.pylib, see
bootstrap-pylib.sh), or let create_world.py call it when the PNGs are missing.

  D:/UE_5.8/Engine/Binaries/ThirdParty/Python3/Win64/python.exe Scripts/Authoring/generate_heightmap.py
  ... --seed 12                                   # another noise continent, manifest otherwise as is
  ... --source-file RawContent/World/Sources/dem.png --source-elevation 0 2400   # a real heightmap, this once

The manifest names the source. Noise builds a continent with heightmap_noise.py. A file (16-bit PNG or raw
.r16 from any DEM tool) is cropped to a square, read as its own elevation range, resampled onto the world and
re-encoded into the world's range; the paint layers are then derived from the height exactly as for noise,
so swapping to a real heightmap is a manifest edit and a rerun, nothing else. Either way a flat spawn pad is
blended into the centre so the player starts stand on level ground.
"""
import argparse
import os
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, ".pylib"))
import numpy as np  # noqa: E402

import heightmap_erosion  # noqa: E402
import heightmap_io  # noqa: E402
import heightmap_noise  # noqa: E402
from world_manifest import DERIVED_FILES, HEIGHTMAP_DIR, LAYER_FILES, MANIFEST_PATH, load_manifest  # noqa: E402

LAYER_DEFAULTS = {
    "rock_slope_start": 0.55,       # rise over run where rock starts to show through the meadow (about 29 degrees)
    "rock_slope_full": 1.05,        # and where it is all rock (about 46 degrees)
    "high_altitude_start_m": 1100,  # where the high rock layer starts
    "high_altitude_full_m": 1650,   # and where it has taken over
    "breakup_m": 18,                # noise added to the altitude before the rules, so boundaries are not contour lines
    "wear_rock_start": 0.35,        # scraped bedrock (wear map, 0..1) reads as rock from here
    "ridge_rock": 0.6,              # how much convex curvature (ridges, shoulders) adds rock
    "deposit_softens": 0.7,         # how much laid-down sediment (deposit map, 0..1) takes rock away: fans and basins are meadow
}


def source_height_metres(manifest, size):
    """The world's height in metres, size x size, from whichever source the manifest names."""
    source = manifest.source
    kind = source.get("kind", "noise")
    quad_m = manifest.quad_cm / 100.0
    if kind == "noise":
        seed = int(source.get("seed", 7))
        print(f"noise source, seed {seed}")
        return heightmap_noise.generate_metres(size, seed, quad_m, manifest.sea_level_m)
    if kind == "file":
        path = manifest.resolve(source["path"])
        print(f"file source {path}")
        values = heightmap_io.read_heightmap(path, source.get("width"))
        if source.get("flip_y", False):
            values = values[::-1]
        values = heightmap_io.center_crop_square(values)
        elevation = source.get("elevation_m", {"min": manifest.elevation_min_m, "max": manifest.elevation_max_m})
        low, high = float(elevation["min"]), float(elevation["max"])
        metres = low + values.astype(np.float32) / 65535.0 * (high - low)
        print(f"  {values.shape[1]}x{values.shape[0]} samples spanning {low:g}..{high:g} m, resampled to {size}x{size}")
        metres = heightmap_io.resample(metres, size)
        smooth = int(source.get("smooth_passes", 0))
        if smooth > 0:
            metres = heightmap_noise.box_blur(metres, smooth)
        return metres
    raise ValueError(f"{manifest.path}: unknown source kind {kind!r}; use 'noise' or 'file'")


def apply_spawn_pad(metres, manifest):
    """A flat disc at the centre for the player starts, blended into the terrain over a second radius."""
    if manifest.spawn_pad_m <= 0:
        return metres
    size = metres.shape[0]
    radius = manifest.spawn_pad_m / manifest.side_m
    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / (size - 1)
    dist = np.sqrt((x - 0.5) ** 2 + (y - 0.5) ** 2)
    weight = heightmap_noise.smoothstep(np.clip(1.0 - (dist - radius) / radius, 0.0, 1.0))
    centre = float(metres[size // 2, size // 2])
    pad_height = max(centre, manifest.sea_level_m + 150.0)   # never a pad in the sea
    return (metres * (1 - weight) + pad_height * weight).astype(np.float32)


def derived_maps(metres, maps, manifest):
    """The erosion's leftovers squashed to [0, 1]: flow (log scaled), wear, deposit, and a curvature map with
    0.5 flat, convex above, concave below."""
    curve = heightmap_erosion.curvature(metres, manifest.quad_cm / 100.0)
    scale = max(float(np.percentile(np.abs(curve), 99.0)), 1e-6)
    return {
        "flow": heightmap_erosion.to_unit(maps["flow"], 99.5, log_scale=True),
        "wear": heightmap_erosion.to_unit(maps["wear"], 99.0),
        "deposit": heightmap_erosion.to_unit(maps["deposit"], 99.0),
        "curvature": np.clip(0.5 + curve / scale * 0.5, 0.0, 1.0).astype(np.float32),
    }


def derive_layers(metres, derived, manifest, rng):
    """Rocky Meadows' three paint layers from the finished height and what the erosion left: meadow everywhere,
    rock by slope, on scraped bedrock and on convex ridges, less rock where sediment was laid down, high rock
    by altitude. Returns {layer name: uint8 weightmap}, the three summing to 255."""
    rules = {**LAYER_DEFAULTS, **manifest.layers}
    quad_m = manifest.quad_cm / 100.0
    size = metres.shape[0]

    breakup = (heightmap_noise.fbm(size, rng, base_cells=24, octaves=4) - 0.5) * 2.0 * float(rules["breakup_m"])
    gy, gx = np.gradient(heightmap_noise.box_blur(metres, 2), quad_m)
    slope = np.sqrt(gx * gx + gy * gy)
    slope_breakup = breakup / float(rules["breakup_m"]) * 0.12 if rules["breakup_m"] else 0.0

    rock = heightmap_noise.smoothstep(np.clip(
        (slope + slope_breakup - rules["rock_slope_start"]) / (rules["rock_slope_full"] - rules["rock_slope_start"]), 0.0, 1.0))
    scraped = heightmap_noise.smoothstep(np.clip((derived["wear"] - rules["wear_rock_start"]) / (1.0 - rules["wear_rock_start"]), 0.0, 1.0))
    convex = np.clip((derived["curvature"] - 0.5) * 2.0, 0.0, 1.0) * np.clip(slope / rules["rock_slope_start"], 0.0, 1.0)
    rock = np.maximum(rock, np.maximum(scraped * 0.9, convex * rules["ridge_rock"]))
    rock = rock * (1.0 - rules["deposit_softens"] * derived["deposit"] * (slope < rules["rock_slope_full"]))
    high = heightmap_noise.smoothstep(np.clip(
        (metres + breakup - rules["high_altitude_start_m"]) / (rules["high_altitude_full_m"] - rules["high_altitude_start_m"]), 0.0, 1.0))
    high = high * (1.0 - rock * 0.5)
    meadow = np.clip(1.0 - rock - high, 0.0, 1.0)
    total = np.maximum(meadow + rock + high, 1e-6)
    # The pack's names are not what they sound like: Base_Layer is its rock, Layer_02 its grass, Layer_03 its
    # high rock (read from the textures its layer functions sample).
    weights = {"Base_Layer": rock, "Layer_02": meadow, "Layer_03": high}
    return {name: np.rint(w / total * 255.0).astype(np.uint8) for name, w in weights.items()}


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--manifest", default=MANIFEST_PATH)
    parser.add_argument("--seed", type=int, help="override the noise seed for this run")
    parser.add_argument("--source-file", help="use this heightmap file for this run instead of the manifest's source")
    parser.add_argument("--source-elevation", nargs=2, type=float, metavar=("MIN_M", "MAX_M"),
                        help="what 0 and 65535 mean in the source file, in metres")
    parser.add_argument("--out", default=HEIGHTMAP_DIR)
    args = parser.parse_args(argv)

    manifest = load_manifest(args.manifest)
    if args.source_file:
        manifest.source = {"kind": "file", "path": args.source_file}
        if args.source_elevation:
            manifest.source["elevation_m"] = {"min": args.source_elevation[0], "max": args.source_elevation[1]}
    elif args.seed is not None:
        manifest.source = {"kind": "noise", "seed": args.seed}
    print(manifest.describe())

    size = manifest.vertices_per_side
    started = time.time()
    metres = source_height_metres(manifest, size)
    print(f"uplift {metres.min():.0f}..{metres.max():.0f} m in {time.time() - started:.0f} s; eroding")
    metres, maps = heightmap_erosion.erode(metres, manifest.quad_cm / 100.0, manifest.sea_level_m, manifest.erosion)
    metres = apply_spawn_pad(metres, manifest)

    clipped = float(((metres < manifest.elevation_min_m) | (metres > manifest.elevation_max_m)).mean()) * 100.0
    metres = np.clip(metres, manifest.elevation_min_m, manifest.elevation_max_m)
    height = np.rint(manifest.metres_to_value(metres)).clip(0, 65535).astype(np.uint16)
    derived = derived_maps(metres, maps, manifest)
    layers = derive_layers(metres, derived, manifest, np.random.default_rng(int(manifest.source.get("seed", 0)) + 1))

    os.makedirs(args.out, exist_ok=True)
    heightmap_io.write_png(os.path.join(args.out, "L_Canvas_Proto_Height.png"), height)
    for name, data in layers.items():
        heightmap_io.write_png(os.path.join(args.out, LAYER_FILES[name]), data)
    for name, data in derived.items():
        heightmap_io.write_png(os.path.join(args.out, DERIVED_FILES[name]), np.rint(data * 255.0).astype(np.uint8))

    land = float((metres > manifest.sea_level_m).mean()) * 100.0
    print(f"height {metres.min():.0f}..{metres.max():.0f} m, {land:.0f}% above sea level, {clipped:.2f}% clipped to the range; "
          f"layers meadow {layers['Layer_02'].mean() / 255 * 100:.0f}% rock {layers['Base_Layer'].mean() / 255 * 100:.0f}% "
          f"high {layers['Layer_03'].mean() / 255 * 100:.0f}%; {time.time() - started:.0f} s; written to {args.out}")


if __name__ == "__main__":
    main()