"""Generates a seeded heightmap and three weightmaps for L_World, as 16-bit and 8-bit greyscale PNGs.

Pure numpy, no engine: run it with any Python that has numpy, or let create_world.py call it. The output is
plain files under Content/World/Heightmaps/, so swapping the terrain later is dropping in a different PNG of
any resolution and rerunning create_world.py; nothing else in the project knows how the terrain was made.

  UE_5.8/Engine/Binaries/ThirdParty/Python3/Win64/python.exe Scripts/Authoring/generate_heightmap.py [--seed N] [--size 4033]

The shape, in metres, with the default scale in create_world.py (350 cm per quad, Z scale 500):
  a continent with ragged coasts and sea around it, lowland plains, rolling hills, one or two mountain ranges
  along a low-frequency band, thermal smoothing so slopes read as slopes, and a flat 200 m pad at the centre
  for the player start. Weightmaps: base (grass) everywhere, layer 2 (rock) by slope, layer 3 by altitude.
"""
import argparse
import os
import struct
import sys
import zlib

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

OUT_DIR = os.path.normpath(os.path.join(HERE, "..", "..", "Content", "World", "Heightmaps"))
SEA_LEVEL = 0.18          # fraction of the 16-bit range that is sea; create_world.py places the water plane here
PAD_RADIUS_FRACTION = 0.01  # flat spawn pad, as a fraction of the map width


def write_png(path, data):
    """Greyscale PNG, 8 or 16 bit from the array dtype. Row filter 0, one zlib stream."""
    if data.dtype == np.uint16:
        depth, payload = 16, data.astype(">u2")
    else:
        depth, payload = 8, data.astype(np.uint8)
    height, width = data.shape
    raw = b"".join(b"\x00" + payload[y].tobytes() for y in range(height))

    def chunk(kind, body):
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body) & 0xFFFFFFFF)

    ihdr = struct.pack(">IIBBBBB", width, height, depth, 0, 0, 0, 0)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw, 6)) + chunk(b"IEND", b""))


def smoothstep(t):
    return t * t * (3.0 - 2.0 * t)


def value_noise(size, cells, rng):
    """One octave: a random lattice of cells x cells, smoothly interpolated to size x size. Tileable enough."""
    lattice = rng.random((cells + 1, cells + 1), dtype=np.float32)
    coords = np.linspace(0.0, cells, size, endpoint=False, dtype=np.float32)
    i = np.floor(coords).astype(np.int32)
    t = smoothstep(coords - i)
    i1 = np.minimum(i + 1, cells)
    top = lattice[i[:, None], i[None, :]] * (1 - t[None, :]) + lattice[i[:, None], i1[None, :]] * t[None, :]
    bottom = lattice[i1[:, None], i[None, :]] * (1 - t[None, :]) + lattice[i1[:, None], i1[None, :]] * t[None, :]
    return top * (1 - t[:, None]) + bottom * t[:, None]


def fbm(size, rng, base_cells=4, octaves=8, gain=0.5, ridged=False):
    total = np.zeros((size, size), dtype=np.float32)
    amplitude, cells, norm = 1.0, base_cells, 0.0
    for _ in range(octaves):
        n = value_noise(size, cells, rng)
        if ridged:
            n = 1.0 - np.abs(n * 2.0 - 1.0)
            n = n * n
        total += n * amplitude
        norm += amplitude
        amplitude *= gain
        cells *= 2
    return total / norm


def box_blur(h, passes):
    for _ in range(passes):
        padded = np.pad(h, 1, mode="edge")
        h = (padded[:-2, 1:-1] + padded[2:, 1:-1] + padded[1:-1, :-2] + padded[1:-1, 2:] + h) / 5.0
    return h.astype(np.float32)


def thermal_smooth(h, passes, talus):
    """Cheap erosion: where a cell is much higher than a neighbour, move a little material downhill."""
    for _ in range(passes):
        padded = np.pad(h, 1, mode="edge")
        for dy, dx in ((0, 1), (0, -1), (1, 0), (-1, 0)):
            neighbour = padded[1 + dy:1 + dy + h.shape[0], 1 + dx:1 + dx + h.shape[1]]
            diff = h - neighbour
            move = np.where(diff > talus, (diff - talus) * 0.25, 0.0).astype(np.float32)
            h -= move
    return h


def generate(size, seed):
    rng = np.random.default_rng(seed)
    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / (size - 1)

    # Continent: a radial falloff with a ragged, noise-warped edge, so the coast is not a circle.
    cx, cy = 0.5 + (rng.random() - 0.5) * 0.15, 0.5 + (rng.random() - 0.5) * 0.15
    radius = np.sqrt(((x - cx) * 1.05) ** 2 + ((y - cy) * 0.95) ** 2)
    coast_warp = (fbm(size, rng, base_cells=3, octaves=5) - 0.5) * 0.35
    continent = np.clip(1.0 - (radius + coast_warp) / 0.55, 0.0, 1.0)
    continent = smoothstep(np.clip(continent * 1.6, 0.0, 1.0))

    plains = fbm(size, rng, base_cells=6, octaves=4) * 0.06
    hills = fbm(size, rng, base_cells=12, octaves=6, gain=0.5) * 0.22
    # Mountain ranges: ridged noise, masked by a low-frequency band so they come as ranges, not everywhere.
    range_band = fbm(size, rng, base_cells=3, octaves=3)
    range_mask = smoothstep(np.clip((range_band - 0.47) / 0.2, 0.0, 1.0))
    mountains = fbm(size, rng, base_cells=10, octaves=8, gain=0.5, ridged=True) * range_mask

    land = 0.05 + plains + hills * (0.4 + 0.6 * continent) + mountains * 0.75
    height = SEA_LEVEL + continent * land
    # The sea floor keeps a little shape so the shore is not a hard step.
    sea_floor = SEA_LEVEL - 0.03 - (1.0 - continent) * 0.04 + plains * 0.3
    height = np.where(continent > 0.02, height, np.maximum(sea_floor, 0.0)).astype(np.float32)
    height = np.maximum(height, sea_floor.astype(np.float32))

    height = thermal_smooth(height, passes=6, talus=0.0025)

    # A flat pad at the centre for the player start, blended into the terrain around it.
    pad_radius = PAD_RADIUS_FRACTION
    pad_dist = np.sqrt((x - 0.5) ** 2 + (y - 0.5) ** 2)
    pad_weight = smoothstep(np.clip(1.0 - (pad_dist - pad_radius) / pad_radius, 0.0, 1.0))
    pad_height = max(float(height[size // 2, size // 2]), SEA_LEVEL + 0.06)
    height = height * (1 - pad_weight) + pad_height * pad_weight

    height = np.clip(height, 0.0, 1.0)

    # Weightmaps from the finished shape. Slope is per quad in height-range units; the thresholds are guesses
    # to be tuned by eye in the editor.
    gy, gx = np.gradient(box_blur(height, 3))
    slope = np.sqrt(gx * gx + gy * gy) * size
    rock = smoothstep(np.clip((slope - 1.8) / 1.6, 0.0, 1.0))
    high = smoothstep(np.clip((height - 0.5) / 0.16, 0.0, 1.0)) * (1.0 - rock * 0.5)
    base = np.clip(1.0 - rock - high, 0.0, 1.0)
    total = base + rock + high
    weights = [np.rint(w / total * 255.0).astype(np.uint8) for w in (base, rock, high)]

    return np.rint(height * 65535.0).astype(np.uint16), weights


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--size", type=int, default=4033, help="vertices per side; 4033 fits 32x32 components of 126 quads")
    parser.add_argument("--out", default=OUT_DIR)
    args = parser.parse_args()

    os.makedirs(args.out, exist_ok=True)
    height, weights = generate(args.size, args.seed)
    write_png(os.path.join(args.out, "L_World_Height.png"), height)
    for name, data in zip(("Base_Layer", "Layer_02", "Layer_03"), weights):
        write_png(os.path.join(args.out, f"L_World_{name}.png"), data)
    land = float((height > SEA_LEVEL * 65535).mean()) * 100.0
    print(f"seed {args.seed}: {args.size}x{args.size}, {land:.0f}% land, written to {args.out}")


if __name__ == "__main__":
    main()