"""Cinematic serene flyover over the eroded Judean terrain.

Usage:
  python3 flyover.py test              -> one frame -> flyover_test.png
  python3 flyover.py <start> <count>   -> render frames [start,start+count)
"""
import sys, numpy as np
import geomorph as G
from scipy.ndimage import gaussian_filter, zoom

D = np.load("baked_judean.npz")
h = gaussian_filter(D["h"].astype(np.float32), 0.9)   # tame grid-scale noise
acc = D["acc"].astype(np.float32)

SCREEN = (640, 360)
GAIN = 120.0
NFRAMES = 150

_PAL = np.array([
    [0.34, 0.24, 0.16],  # shadowed wadi floor (ochre)
    [0.49, 0.37, 0.23],  # lower slope tan
    [0.58, 0.50, 0.33],  # olive-tan scrub
    [0.67, 0.61, 0.46],  # dry grass / soil
    [0.76, 0.72, 0.60],  # pale limestone
    [0.88, 0.86, 0.80],  # bright crest rock
], np.float32)
_PT = np.linspace(0, 1, len(_PAL))


def cinematic_color(h, acc):
    p2, p98 = np.percentile(h, 2), np.percentile(h, 98)
    t = np.clip((h - p2) / (p98 - p2 + 1e-9), 0, 1)
    rgb = np.stack([np.interp(t, _PT, _PAL[:, c]) for c in range(3)], -1)
    # warm low sun from the NW — low altitude = long dramatic shadows
    hs = G.hillshade(h, az=300, alt=30, z=4.2)[..., None]
    rgb = rgb * (0.35 + 1.00 * hs)                          # strong shade contrast
    rgb += 0.05 * (1 - hs) * np.array([0.16, 0.26, 0.44])   # cool skylight fill
    rgb *= np.array([1.06, 1.0, 0.92])                      # warm golden grade
    # drainage: blue water in channels, white foam where it plunges
    a = (acc - acc.min()) / (np.ptp(acc) + 1e-9)
    river = np.clip((a - 0.60) / 0.40, 0, 1)[..., None]
    gy, gx = np.gradient(h)
    slope = np.hypot(gx, gy)
    foam = (np.clip((slope - 0.10) * 7, 0, 1) * river[..., 0])[..., None]
    rgb = rgb * (1 - river) + np.array([0.11, 0.30, 0.45]) * river
    rgb = rgb * (1 - foam) + np.array([0.90, 0.95, 1.0]) * foam
    return np.clip(rgb, 0, 1).astype(np.float32)


COLOR = cinematic_color(h, acc)
# 2x supersample the field for smoother silhouettes under the voxel march
h = zoom(h, 2, order=1)
COLOR = zoom(COLOR, (2, 2, 1), order=1)
Hh, Wh = h.shape


def camera(s):
    """s in [0,1] -> serene west->east descent toward the rift, looking east."""
    px = Wh * (0.12 + 0.74 * s)
    py = Hh * (0.5 + 0.10 * np.sin(2 * np.pi * s))
    ground = h[int(np.clip(py, 0, Hh - 1)), int(np.clip(px, 0, Wh - 1))] * GAIN
    cam_h = ground + 42 + 10 * np.sin(2 * np.pi * s * 1.5)     # gentle bob
    heading = -np.pi / 2 + 0.22 * np.sin(2 * np.pi * s)        # look ~east, slow sway
    pitch = 58.0                                              # steeper downward gaze
    return (px, py, cam_h, heading, pitch)


def render_frame(i):
    s = i / (NFRAMES - 1)
    img = G.render_voxel(h, COLOR, camera(s), screen=SCREEN, scale=320.0,
                         distance=480, horizon_frac=0.50, dz0=0.7, dzg=1.013,
                         elevation_gain=GAIN, fog=(0.70, 0.72, 0.74),
                         sky_top=(0.28, 0.44, 0.68), fog_k=0.0026, fog_max=0.32)
    img = np.clip(img, 0, 1)
    img = img * img * (3 - 2 * img)                 # smoothstep S-curve (adds contrast)
    lum = img.mean(-1, keepdims=True)               # saturation boost
    img = np.clip(lum + 1.35 * (img - lum), 0, 1)
    yy, xx = np.mgrid[0:SCREEN[1], 0:SCREEN[0]]
    q = (xx / SCREEN[0]) * (1 - xx / SCREEN[0]) * (yy / SCREEN[1]) * (1 - yy / SCREEN[1])
    img *= (0.78 + 0.22 * np.clip(16 * q, 0, 1))[..., None]
    return np.clip(img ** 0.92, 0, 1)


if __name__ == "__main__":
    if sys.argv[1] == "test":
        G.save_png(render_frame(int(0.45 * NFRAMES)), "flyover_test.png")
        print("wrote flyover_test.png")
    else:
        import os
        os.makedirs("frames_fly", exist_ok=True)
        start, count = int(sys.argv[1]), int(sys.argv[2])
        import time; t0 = time.time()
        for i in range(start, min(start + count, NFRAMES)):
            G.save_png(render_frame(i), f"frames_fly/f_{i:03d}.png")
        print(f"frames {start}..{start+count-1} in {time.time()-t0:.1f}s")