"""
worldgen.py  --  bake fantasy-world heightmaps from noise.

Run:  python3 worldgen.py
Out:  heightmap_1.png .. heightmap_20.png   (map.scad picks one by seed)
      worlds.png   (contact sheet -- browse, pick a seed)

The island shape lives HERE (it's instant array math). Doing it in
OpenSCAD needs CSG, which is exactly what was freezing the preview.
"""
import numpy as np
import opensimplex
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, LightSource
from PIL import Image
from pathlib import Path
OUT = Path(__file__).parent

# ---- settings ---------------------------------------------------
SEEDS   = range(1, 21)   # which worlds to bake (seed numbers)
GRID    = 200            # heightmap resolution = terrain detail
OCTAVES = 6              # layers of fractal detail
ZOOM    = 3.2            # bigger = more, smaller landmasses
ISLAND  = True           # True = island, False = mainland (land to edges)
COAST   = 0.95           # island only: bigger = smaller island, more sea

# ---- grow one world from a seed ---------------------------------
def world(seed):
    opensimplex.seed(seed)

    axis = np.linspace(0, ZOOM, GRID)
    e = np.zeros((GRID, GRID))

    freq = 1.0
    amp = 1.0
    total = 0.0

    for _ in range(OCTAVES):
        e += amp * opensimplex.noise2array(axis * freq, axis * freq)
        total += amp
        freq *= 2.0
        amp *= 0.5

    # normalize base noise to 0..1
    e = (e / total + 1) / 2
    e = (e - e.min()) / np.ptp(e)

    if ISLAND:
        # Make an island mask: center stays high, edges get pushed down hard
        gy, gx = np.mgrid[-1:1:GRID*1j, -1:1:GRID*1j]

        # distance from center: 0 in middle, ~1.4 in corners
        dist = np.sqrt(gx**2 + gy**2)

        # smooth edge falloff
        # lower START = smaller island, more ocean
        START = 0.42
        END   = 1.00

        mask = np.clip((dist - START) / (END - START), 0, 1)
        mask = mask ** 2.2

        # add slight noisy coast shape so it isn't a perfect circle
        rough = (opensimplex.noise2array(axis * 1.2, axis * 1.2) + 1) / 2
        coast_noise = 0.75 + rough * 0.5

        # sink edges
        e = e - mask * COAST * coast_noise

        # clamp ocean
        e = np.clip(e, 0, 1)

        # important: don't fully normalize from the new minimum,
        # because that can bring ocean edges back up
        e = np.clip(e, 0, 1)

    else:
        # mainland mode: lift the lowest terrain so it reaches the edges
        # this prevents the map from still becoming "ocean" after normalization
        SEA_LEVEL = 0.28

        e = np.maximum(e, SEA_LEVEL)
        e = (e - e.min()) / np.ptp(e)

        # optional: make the whole thing slightly more land-like
        e = e * 0.85 + 0.15

    return e

# ---- bake every seed + build the contact sheet ------------------
pal = ListedColormap(['#39618a', '#d8c98e', '#3f7d4e', '#8a7d6b', '#f4f4f2'])
ls  = LightSource(315, 45)
seeds = list(SEEDS)
cols, rows = 5, -(-len(seeds) // 5)
fig, ax = plt.subplots(rows, cols, figsize=(cols * 2.4, rows * 2.4))
ax = ax.ravel()
for i, s in enumerate(seeds):
    e = world(s)
    Image.fromarray((e * 255).astype(np.uint8), 'L').save(OUT / f'heightmap_{s}.png')
    b = np.zeros_like(e, int)
    b[e > 0.34] = 1; b[e > 0.48] = 2; b[e > 0.66] = 3; b[e > 0.84] = 4
    rgb = pal(b / 4.0)[:, :, :3]
    ax[i].imshow(ls.shade_rgb(rgb, e, blend_mode='soft', vert_exag=4))
    ax[i].set_title(f"seed {s}", fontsize=10)
    ax[i].axis('off')
for j in range(len(seeds), len(ax)):
    ax[j].axis('off')
plt.tight_layout()
plt.savefig(OUT / 'worlds.png', dpi=110, bbox_inches='tight')
print(f"baked {len(seeds)} worlds at {GRID}x{GRID}px. "
      f"set heightmap_px = {GRID} in map.scad.")