"""Generates Game/assets/models/nebula_planet_planet_surface.png, the sidecar
that actually feeds the live nebula_planet.glb material -- Godot's glTF
importer extracted the embedded image there at import time, so the imported
scene references that external file, not the glb's internal binary. A
2048x1024 equirectangular decorative planet with latitude bands, storm
vortices, and a lit/unlit terminator, via layered FFT/domain-warped noise in
the same style as gen_nebula_sky.py (helpers duplicated here to keep both
scripts standalone).
"""

from pathlib import Path

import numpy as np
from PIL import Image

W, H = 2048, 1024

# Direction *toward* arena_02's DirectionalLight3D in world space (its basis Z-column;
# NebulaPlanet has no rotation, so mesh-local axes equal world axes and this can be used
# directly against the UV-derived normal below). A surface texel is lit when its normal
# points roughly toward this direction.
LIGHT_DIR = np.array([0.321394, 0.766044, 0.55667])

# Sphere's exact UV convention (extracted from Game/assets/models/nebula_planet.glb's
# vertex data: +X -> uv(0.5,0.5), -X -> uv(0.0,0.5), +Z -> uv(0.25,0.5),
# -Z -> uv(0.75,0.5), +Y -> uv(*, 0.0), -Y -> uv(*, 1.0)):
#   theta = v*pi, phi = 2*pi*(u-0.5)
#   x = sin(theta)*cos(phi), y = cos(theta), z = -sin(theta)*sin(phi)


def fft_field(power, seed, remove_dc=True):
    r = np.random.default_rng(seed)
    white = r.normal(size=(H, W))
    F = np.fft.fft2(white)
    fy = np.fft.fftfreq(H)[:, None]
    fx = np.fft.fftfreq(W)[None, :]
    freq = np.sqrt(fx ** 2 + fy ** 2)
    freq[0, 0] = 1e-6
    filt = 1.0 / (freq ** power)
    if remove_dc:
        filt[0, 0] = 0.0
    field = np.fft.ifft2(F * filt).real
    field -= field.min()
    field /= (field.max() + 1e-9)
    return field


def bilinear_sample(field, xs, ys):
    x0 = np.floor(xs).astype(np.int64) % W
    x1 = (x0 + 1) % W
    y0 = np.clip(np.floor(ys).astype(np.int64), 0, H - 1)
    y1 = np.clip(y0 + 1, 0, H - 1)
    fx = xs - np.floor(xs)
    fy = ys - np.floor(ys)
    v00 = field[y0, x0]
    v10 = field[y0, x1]
    v01 = field[y1, x0]
    v11 = field[y1, x1]
    return v00 * (1 - fx) * (1 - fy) + v10 * fx * (1 - fy) + v01 * (1 - fx) * fy + v11 * fx * fy


def ramp(t, stops):
    out = np.zeros(t.shape + (3,), dtype=np.float64)
    for i in range(len(stops) - 1):
        t0, c0 = stops[i]
        t1, c1 = stops[i + 1]
        seg = (t >= t0) & (t <= t1)
        if not np.any(seg):
            continue
        local = (t[seg] - t0) / (t1 - t0)
        for ch in range(3):
            out[seg, ch] = c0[ch] + (c1[ch] - c0[ch]) * local
    return out


yy, xx = np.mgrid[0:H, 0:W].astype(np.float64)

print("generating warp fields...")
warp_x = (fft_field(3.0, seed=101) - 0.5) * 160
warp_y = (fft_field(3.0, seed=102) - 0.5) * 80

print("generating storm vortices...")
# Each vortex spirals nearby coordinates around its center before bands/turbulence
# are sampled through them, so bands visibly swirl into the storm rather than sitting
# as flat stripes underneath a cosmetic overlay.
vortex_defs = [
    {"cx": W * 0.50, "cy": H * 0.50, "radius": W * 0.09, "strength": 2.6, "dir": 1},
    {"cx": W * 0.20, "cy": H * 0.30, "radius": W * 0.05, "strength": 1.8, "dir": -1},
    {"cx": W * 0.78, "cy": H * 0.68, "radius": W * 0.045, "strength": -2.0, "dir": 1},
]

vx, vy = xx.copy(), yy.copy()
falloffs = []
for v in vortex_defs:
    dxv = xx - v["cx"]
    dxv -= W * np.round(dxv / W)  # shortest signed horizontal wrap distance
    dyv = yy - v["cy"]
    rr = np.sqrt(dxv ** 2 + dyv ** 2)
    falloff = np.exp(-((rr / v["radius"]) ** 2))
    falloffs.append(falloff)
    theta = np.arctan2(dyv, dxv) + v["dir"] * v["strength"] * falloff
    vx = np.where(falloff > 0.01, v["cx"] + rr * np.cos(theta), vx)
    vy = np.where(falloff > 0.01, v["cy"] + rr * np.sin(theta), vy)

print("generating latitude bands...")
band_count = 8
lat_norm = vy / H
lat_wobble = bilinear_sample(fft_field(2.2, seed=111), vx * 0.5 + warp_x * 1.4, vy + warp_y * 1.4)
warped_lat = lat_norm + (lat_wobble - 0.5) * 0.10
band_id = np.floor(warped_lat * band_count)
band_frac = warped_lat * band_count - band_id
# Wide, soft transition zone (most of each band's width) so bands blend into each
# other like atmospheric flow rather than reading as flat-shaded stripes.
band_edge = np.clip(np.minimum(band_frac, 1 - band_frac) / 0.42, 0.0, 1.0)
band_edge = band_edge * band_edge * (3 - 2 * band_edge)  # smoothstep

rng_bands = np.random.default_rng(112)
band_shade = rng_bands.uniform(-1, 1, band_count + 2)
band_shade_field = band_shade[np.clip(band_id.astype(np.int64), 0, band_count + 1)]

turbulence = bilinear_sample(fft_field(1.4, seed=121), vx + warp_x * 0.6, vy + warp_y * 0.6)
value = np.clip(0.5 + band_shade_field * 0.13 * (0.4 + 0.6 * band_edge) + (turbulence - 0.5) * 0.20, 0.0, 1.0)

print("color grading...")
# Palette-matched to the current bake (dark navy/purple -> dusty pink/mauve) and to
# gen_nebula_sky.py's Orion-esque family, for thematic consistency between the two.
stops = [
    (0.00, (0.05, 0.03, 0.12)),
    (0.30, (0.16, 0.07, 0.24)),
    (0.55, (0.38, 0.14, 0.36)),
    (0.75, (0.62, 0.28, 0.46)),
    (1.00, (0.86, 0.62, 0.72)),
]
color = ramp(value, stops)

print("adding hue variation...")
hue_field = bilinear_sample(fft_field(2.6, seed=131), vx * 0.7 + warp_x, vy * 0.7 + warp_y)
hue_field = (hue_field - 0.5) * 2.0
warm_tint = np.array([0.22, 0.03, -0.04])
cool_tint = np.array([-0.10, 0.08, 0.05])
hue_variation = np.where(
    hue_field[..., None] > 0,
    hue_field[..., None] * warm_tint[None, None, :],
    -hue_field[..., None] * cool_tint[None, None, :],
)
color = np.clip(color + hue_variation * 0.16, 0.0, 1.0)

print("adding storm highlights...")
for v, falloff in zip(vortex_defs, falloffs):
    hi = np.clip((falloff - 0.55) / 0.45, 0.0, 1.0) ** 1.3
    tint = np.array([0.35, 0.22, 0.18]) * (0.6 if v["strength"] < 0 else 1.0)
    color += hi[..., None] * tint[None, None, :]
color = np.clip(color, 0.0, 1.0)

print("baking terminator...")
u = xx / W
v = yy / H
theta = v * np.pi
phi = 2.0 * np.pi * (u - 0.5)
nx = np.sin(theta) * np.cos(phi)
ny = np.cos(theta)
nz = -np.sin(theta) * np.sin(phi)
lit_raw = nx * LIGHT_DIR[0] + ny * LIGHT_DIR[1] + nz * LIGHT_DIR[2]  # -1..1

# Soft-edged but asymmetric terminator band (favors more of the sphere reading lit,
# since a fully half-dark planet reads badly from most camera angles as a background
# decoration). Smoothstep between two dot-product thresholds.
edge0, edge1 = -0.45, 0.35
t = np.clip((lit_raw - edge0) / (edge1 - edge0), 0.0, 1.0)
lit = t * t * (3.0 - 2.0 * t)

# Pushed harder than physically correct: Planet_Surface's emissiveFactor is flat and
# uniform, which washes out real-time per-pixel lighting almost entirely in-engine, so
# the day/night contrast has to be baked directly into the diffuse texture instead.
color = color * (0.28 + 0.85 * lit)[..., None]
color = color + (lit[..., None] ** 2) * np.array([0.07, 0.03, -0.02])[None, None, :] * 0.6
color = np.clip(color, 0.0, 1.0)

print("saving planet surface...")
root = Path(__file__).resolve().parents[2]
out_path = root / "Game" / "assets" / "models" / "nebula_planet_planet_surface.png"
img = Image.fromarray((color * 255).astype(np.uint8), "RGB")
img.save(out_path)
print("saved:", out_path)
print("done")