eagle0/tools/map_generator/roughen_borders.py

#!/usr/bin/env python3
"""Roughen newly-drawn borders to look natural.

Takes the fixed import (user-edited) and applies localized domain warp
only near the changed boundary pixels, then cleans up fragments and
verifies the result.
"""
import gzip
from collections import Counter
from pathlib import Path

import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter, label, binary_dilation

SCRIPT_DIR = Path(__file__).parent
UNITY_ASSETS = (
    SCRIPT_DIR.parent.parent
    / "src/main/csharp/net/eagle0/clients/unity/eagle0/Assets/Eagle"
)
MAP_WIDTH = 3786
MAP_HEIGHT = 1834
OCEAN_ID = 0
BORDER_ID = 255


def load_gz(path):
    with gzip.open(path, "rb") as f:
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(MAP_HEIGHT, MAP_WIDTH).copy()


def save_gz(arr, path):
    path = path.with_suffix(".gz.bytes")
    with gzip.open(path, "wb") as f:
        f.write(arr.astype(np.uint8).tobytes())
    print(f"Saved to {path}")


def main():
    original = load_gz(UNITY_ASSETS / "rawGray.gz.bytes")
    edited = load_gz(SCRIPT_DIR / "output/rawGray.gz.bytes")

    # Restore ocean/border from original
    ocean_or_border = (original == OCEAN_ID) | (original == BORDER_ID)
    edited[ocean_or_border] = original[ocean_or_border]

    diff_mask = original != edited
    print(f"Changed pixels: {np.sum(diff_mask):,}")

    # Find the boundary of the changed region (where edited provinces meet)
    ys, xs = np.where(diff_mask)
    pad = 30  # roughen pixels within this distance of changes
    y_min = max(0, int(ys.min()) - pad)
    y_max = min(MAP_HEIGHT, int(ys.max()) + pad)
    x_min = max(0, int(xs.min()) - pad)
    x_max = min(MAP_WIDTH, int(xs.max()) + pad)
    print(f"Roughening region: y[{y_min}:{y_max}], x[{x_min}:{x_max}]")

    # Apply domain warp only in the changed region
    rng = np.random.default_rng(42)
    h = y_max - y_min
    w = x_max - x_min
    amplitude = 3.0
    scale = 15.0

    dx_noise = rng.standard_normal((h, w))
    dy_noise = rng.standard_normal((h, w))
    dx = gaussian_filter(dx_noise, sigma=scale)
    dy = gaussian_filter(dy_noise, sigma=scale)
    dx /= np.abs(dx).max() + 1e-10
    dy /= np.abs(dy).max() + 1e-10
    dx *= amplitude
    dy *= amplitude

    yy, xx = np.mgrid[0:h, 0:w]
    src_y = np.clip(np.round(yy + dy).astype(int), 0, h - 1) + y_min
    src_x = np.clip(np.round(xx + dx).astype(int), 0, w - 1) + x_min

    # Only warp pixels near the changed boundary, not the whole region
    # Dilate the diff mask to get a band around changes
    region_diff = diff_mask[y_min:y_max, x_min:x_max]
    struct = np.ones((3, 3), dtype=int)
    warp_band = region_diff.copy()
    for _ in range(pad):
        warp_band = binary_dilation(warp_band, structure=struct)

    result = edited.copy()
    crop = edited[y_min:y_max, x_min:x_max].copy()

    warped_ids = edited[src_y, src_x]
    land_src = (warped_ids != OCEAN_ID) & (warped_ids != BORDER_ID)
    land_dst = (crop != OCEAN_ID) & (crop != BORDER_ID)
    apply = warp_band & land_src & land_dst

    crop[apply] = warped_ids[apply]
    result[y_min:y_max, x_min:x_max] = crop

    # Preserve ocean/border
    result[ocean_or_border] = original[ocean_or_border]

    warped_changed = np.sum(result != edited)
    print(f"Warp changed {warped_changed:,} additional pixels")

    # Fragment cleanup (just in case warp created disconnections)
    land_ids = sorted(set(np.unique(result)) - {OCEAN_ID, BORDER_ID})
    fragments_fixed = 0
    for pid in land_ids:
        mask = result == pid
        labeled, num = label(mask, structure=struct)
        if num <= 1:
            continue
        sizes = np.bincount(labeled.ravel())
        sizes[0] = 0
        largest = np.argmax(sizes)
        for comp in range(1, num + 1):
            if comp == largest:
                continue
            comp_mask = labeled == comp
            comp_size = int(sizes[comp])
            if comp_size < 100:  # only fix small fragments
                # Assign to nearest neighbor province
                dilated = binary_dilation(comp_mask, structure=struct, iterations=3)
                neighbors = result[dilated & ~comp_mask]
                neighbors = neighbors[(neighbors != OCEAN_ID) & (neighbors != BORDER_ID) & (neighbors != pid)]
                if len(neighbors) > 0:
                    new_id = Counter(neighbors.tolist()).most_common(1)[0][0]
                    result[comp_mask] = new_id
                    fragments_fixed += comp_size

    if fragments_fixed:
        print(f"Fixed {fragments_fixed} fragment pixels")

    # Verify
    orig_ids = set(np.unique(original)) - {OCEAN_ID, BORDER_ID}
    result_ids = set(np.unique(result)) - {OCEAN_ID, BORDER_ID}
    lost = orig_ids - result_ids
    if lost:
        print(f"WARNING: Lost province IDs: {lost}")
    else:
        print(f"OK: All {len(orig_ids)} province IDs preserved")

    total_changed = np.sum(result != original)
    print(f"Total pixels changed from original: {total_changed:,}")

    # Save
    out_dir = SCRIPT_DIR / "output"
    save_gz(result, out_dir / "rawGray_roughened")

    # Generate comparison visualization
    from rawgray_editor import build_color_table
    color_table = build_color_table()

    viz_pad = 80
    vy_min = max(0, int(ys.min()) - viz_pad)
    vy_max = min(MAP_HEIGHT, int(ys.max()) + viz_pad)
    vx_min = max(0, int(xs.min()) - viz_pad)
    vx_max = min(MAP_WIDTH, int(xs.max()) + viz_pad)

    def colorize_crop(pmap):
        crop = pmap[vy_min:vy_max, vx_min:vx_max]
        ch, cw = crop.shape
        rgb = np.zeros((ch, cw, 3), dtype=np.uint8)
        for pid, color in color_table.items():
            m = crop == pid
            if np.any(m):
                rgb[m] = color
        # Border darkening
        gx = np.abs(np.diff(crop.astype(np.int16), axis=1))
        gy = np.abs(np.diff(crop.astype(np.int16), axis=0))
        gx = np.pad(gx, ((0, 0), (0, 1)), mode="constant")
        gy = np.pad(gy, ((0, 1), (0, 0)), mode="constant")
        borders = (gx > 0) | (gy > 0)
        for c in range(3):
            rgb[:, :, c] = np.where(borders, rgb[:, :, c] // 3, rgb[:, :, c])
        return np.flipud(rgb)

    before_rgb = colorize_crop(original)
    after_rgb = colorize_crop(result)

    ch, cw, _ = before_rgb.shape
    gap = 10
    combined = np.zeros((ch, cw * 2 + gap, 3), dtype=np.uint8)
    combined[:, :cw] = before_rgb
    combined[:, cw + gap:] = after_rgb
    Image.fromarray(combined).save(out_dir / "roughened_comparison.png")
    print(f"Saved comparison to {out_dir}/roughened_comparison.png")


if __name__ == "__main__":
    main()