"""Draws an overview of the region straight from the tiles, so there is something to navigate by that does not
depend on the editor.

  python Scripts/Authoring/region_overview.py                 # height, shaded
  python Scripts/Authoring/region_overview.py --mode paint    # the three paint layers as colour
  python Scripts/Authoring/region_overview.py --mode both     # writes both

Writes RawContent/World/Region_Overview.png (and Region_Paint.png for the paint mode).

**Orientation.** A tile's PNG has world X across its columns and world Y down its rows - verified, not assumed:
tile (0,0)'s last column is tile (1,0)'s first column, and its last row is tile (0,1)'s first row. The mosaic is
therefore laid out rows = ty, columns = tx, with no transpose anywhere, which also means the result is the same
way up as the crop it came from in the source planet map. An earlier version indexed the mosaic rows by tx and
columns by ty, which transposed every tile inside its own square and scrambled the map.

**Scale.** The hypsometric ramp runs to `--top-m`, an absolute ceiling in metres, not to the map's own maximum.
Normalising by the maximum is how a 30 km window whose median ground is 151 m comes out uniformly dark green
with the whole ramp spent on one peak - the same trap Docs record for the generator's preview.png. Every run
prints the ceiling it used.
"""
import argparse
import os
import struct
import sys
import zlib

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_io  # noqa: E402
from region_manifest import MANIFEST_PATH, load_manifest  # noqa: E402

# Height ramp, as fractions of the ceiling: water, shore, lowland, upland, rock, snow.
HEIGHT_RAMP = [
    (0.00, (58, 84, 120)),
    (0.04, (128, 150, 96)),
    (0.22, (108, 138, 80)),
    (0.45, (150, 142, 96)),
    (0.70, (146, 132, 120)),
    (0.88, (200, 198, 196)),
    (1.00, (248, 248, 250)),
]
# What each paint layer is drawn as. The pack's names mislead: Base_Layer is rock, Layer_02 meadow grass,
# Layer_03 high rock.
PAINT_COLOURS = {
    "Base_Layer": (150, 128, 108),
    "Layer_02": (96, 138, 74),
    "Layer_03": (214, 214, 218),
}


def write_rgb_png(path, rgb):
    """An 8-bit colour PNG; heightmap_io only writes the greyscale the landscape wants."""
    height, width, _ = rgb.shape
    raw = bytearray()
    for row in rgb:
        raw.append(0)
        raw.extend(row.tobytes())

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

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


def ramp_colour(t, ramp):
    stops = np.array([s for s, _ in ramp], dtype=np.float32)
    cols = np.array([c for _, c in ramp], dtype=np.float32)
    out = np.empty(t.shape + (3,), dtype=np.float32)
    for channel in range(3):
        out[..., channel] = np.interp(t, stops, cols[:, channel])
    return out


def mosaic(manifest, per_tile, reader):
    """Every tile laid out rows = ty, columns = tx, which is the orientation the tiles are actually in."""
    wide, high = per_tile * manifest.tiles_x, per_tile * manifest.tiles_y
    out = None
    for tx, ty in manifest.tiles():
        block = reader(tx, ty)
        if block is None:
            continue
        step = max(1, block.shape[0] // per_tile)
        small = block[::step, ::step][:per_tile, :per_tile]
        if out is None:
            out = np.zeros((high, wide) + small.shape[2:], dtype=np.float32)
        out[ty * per_tile:(ty + 1) * per_tile, tx * per_tile:(tx + 1) * per_tile] = small
    return out


def hillshade(metres, metres_per_px):
    """A proper hillshade: the cosine between the surface normal and a light from the north-west, 45 degrees
    up. The previous version multiplied by the raw gradient, which clipped to black and white wherever the
    ground was steep."""
    gy, gx = np.gradient(metres.astype(np.float32), metres_per_px)
    nz = 1.0 / np.sqrt(gx * gx + gy * gy + 1.0)
    nx, ny = -gx * nz, -gy * nz
    lx, ly, lz = -0.5, -0.5, 0.7071
    norm = np.sqrt(lx * lx + ly * ly + lz * lz)
    shade = (nx * lx + ny * ly + nz * lz) / norm
    return np.clip(0.45 + 0.85 * shade, 0.25, 1.35)[..., None]


def draw_line(img, axis, index, colour, width=1):
    lo = max(0, index - width // 2)
    hi = min(img.shape[axis], lo + width)
    if lo >= hi:
        return
    if axis == 0:
        img[lo:hi, :, :] = colour
    else:
        img[:, lo:hi, :] = colour


def decorate(img, manifest, per_tile):
    high, wide = img.shape[0], img.shape[1]
    px_per_km_x = wide / (manifest.width_m / 1000.0)
    px_per_km_y = high / (manifest.height_m / 1000.0)
    for km in range(1, int(manifest.width_m / 1000.0) + 1):
        draw_line(img, 1, int(km * px_per_km_x), (110, 110, 120))
    for km in range(1, int(manifest.height_m / 1000.0) + 1):
        draw_line(img, 0, int(km * px_per_km_y), (110, 110, 120))
    for i in range(manifest.tiles_x + 1):
        draw_line(img, 1, min(wide - 1, i * per_tile), (250, 214, 84), 2)
    for j in range(manifest.tiles_y + 1):
        draw_line(img, 0, min(high - 1, j * per_tile), (250, 214, 84), 2)
    draw_line(img, 1, wide // 2, (255, 86, 86), 3)
    draw_line(img, 0, high // 2, (255, 86, 86), 3)


def read_metres(manifest, tx, ty):
    path = manifest.height_path(tx, ty)
    if not os.path.isfile(path):
        return None
    v = heightmap_io.read_png(path).astype(np.float32)
    return manifest.elevation_min_m + v / 65535.0 * manifest.elevation_span_m


def colour_height(manifest, metres, land, top_m):
    img = ramp_colour(np.clip(metres / max(top_m, 1e-6), 0.0, 1.0), HEIGHT_RAMP)
    img[~land] = (46, 62, 88)
    return np.clip(img * hillshade(metres, manifest.width_m / img.shape[1]), 0, 255)


def paint_map(manifest, per_tile):
    """The three weightmaps composited by weight: what the landscape is painted with, whatever the material
    ends up rendering."""
    layers = {}
    for name in (layer.name for layer in manifest.enabled_layers):
        def read(tx, ty, name=name):
            path = manifest.weight_path(tx, ty, name)
            if not os.path.isfile(path):
                return None
            return heightmap_io.read_png(path).astype(np.float32) / 255.0
        layers[name] = mosaic(manifest, per_tile, read)

    total = np.zeros_like(next(iter(layers.values())))
    img = np.zeros(total.shape + (3,), dtype=np.float32)
    for name, weight in layers.items():
        img += weight[..., None] * np.array(PAINT_COLOURS[name], dtype=np.float32)
        total += weight
    img /= np.maximum(total, 1e-6)[..., None]
    return np.clip(img, 0, 255), {k: float(v.mean()) for k, v in layers.items()}


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--manifest", default=MANIFEST_PATH)
    parser.add_argument("--pixels", type=int, default=2048)
    parser.add_argument("--mode", choices=("height", "paint", "both"), default="height")
    parser.add_argument("--top-m", type=float, help="ceiling of the height ramp in metres; default is the "
                                                    "99th percentile of the land, printed either way")
    args = parser.parse_args(argv)

    manifest = load_manifest(args.manifest)
    folder = os.path.dirname(manifest.path)
    per_tile = max(1, args.pixels // max(manifest.tiles_x, manifest.tiles_y))
    wide, high = per_tile * manifest.tiles_x, per_tile * manifest.tiles_y
    hx, hy = manifest.width_m / 2.0, manifest.height_m / 2.0

    print(f"{wide}x{high} px, {manifest.width_m / 1000:.2f} x {manifest.height_m / 1000:.2f} km, "
          f"{wide / manifest.width_m * 1000:.1f} px/km")
    print(f"  columns are world X {-hx:+.0f}..{+hx:+.0f} m, rows are world Y {-hy:+.0f}..{+hy:+.0f} m")
    print(f"  world X cm = col * {manifest.width_m * 100 / wide:.2f} - {hx * 100:.0f}, "
          f"Y cm = row * {manifest.height_m * 100 / high:.2f} - {hy * 100:.0f}")
    print("  yellow lines are tile edges, grey is a kilometre, red is the origin")

    if args.mode in ("height", "both"):
        metres = mosaic(manifest, per_tile, lambda tx, ty: read_metres(manifest, tx, ty))
        land = metres > manifest.sea_level_m
        top = args.top_m if args.top_m else float(np.percentile(metres[land], 99.0))
        img = colour_height(manifest, metres, land, top)
        decorate(img, manifest, per_tile)
        out = os.path.join(folder, "Region_Overview.png")
        write_rgb_png(out, np.ascontiguousarray(np.rint(img).astype(np.uint8)))
        print(f"  height -> {out}")
        print(f"    ramp ceiling {top:.0f} m; land {land.mean() * 100:.1f}%, "
              f"height {metres[land].min():.0f}..{metres.max():.0f} m, median {np.median(metres[land]):.0f} m")

    if args.mode in ("paint", "both"):
        img, means = paint_map(manifest, per_tile)
        decorate(img, manifest, per_tile)
        out = os.path.join(folder, "Region_Paint.png")
        write_rgb_png(out, np.ascontiguousarray(np.rint(img).astype(np.uint8)))
        print(f"  paint  -> {out}")
        for name, mean in means.items():
            print(f"    {name:11s} {mean * 100:5.1f}%  drawn as {PAINT_COLOURS[name]}")


if __name__ == "__main__":
    main()