"""Generates Game/assets/textures/sky_nebula.png: a 4096x2048 equirectangular
nebula skybox via layered FFT/domain-warped noise, color-graded, with drawn
hero stars (bright stars get a diffraction-spike stamp).
"""

import math
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFilter

W, H = 4096, 2048


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


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

print("generating warp fields...")
warp_x = (fft_field(3.2, seed=1) - 0.5) * 320
warp_y = (fft_field(3.2, seed=2) - 0.5) * 160

print("generating density octaves...")
oct_a = bilinear_sample(fft_field(2.6, seed=11), xx + warp_x, yy + warp_y)
oct_b = bilinear_sample(fft_field(1.9, seed=12), xx + warp_x * 0.6, yy + warp_y * 0.6)
oct_c = bilinear_sample(fft_field(1.3, seed=13), xx + warp_x * 0.3, yy + warp_y * 0.3)
density = oct_a * 0.55 + oct_b * 0.30 + oct_c * 0.15
density = (density - density.min()) / (density.max() - density.min())

print("generating haze + dust lanes...")
haze = bilinear_sample(fft_field(3.0, seed=21), xx + warp_x * 1.3, yy + warp_y * 1.3)
ridged = 1.0 - np.abs(2.0 * bilinear_sample(fft_field(1.6, seed=31), xx + warp_x * 0.8, yy + warp_y * 0.8) - 1.0)
dust = np.clip((ridged - 0.72) / 0.28, 0.0, 1.0) ** 1.5

# Two more dust-lane layers at scales straddling the primary one above, so the sky
# reads as several overlapping filament systems instead of one repeated pattern.
ridged_fine = 1.0 - np.abs(2.0 * bilinear_sample(fft_field(1.0, seed=41), xx + warp_x * 1.6, yy + warp_y * 1.6) - 1.0)
dust_fine = np.clip((ridged_fine - 0.78) / 0.22, 0.0, 1.0) ** 1.2  # delicate, sparse filaments

ridged_wide = 1.0 - np.abs(2.0 * bilinear_sample(fft_field(2.2, seed=51), xx + warp_x * 0.4, yy + warp_y * 0.4) - 1.0)
dust_wide = np.clip((ridged_wide - 0.65) / 0.35, 0.0, 1.0) ** 1.5  # broad, soft secondary lane

# Confine the bright nebula body to an off-center round-ish region (like the reference's
# bright core), fading into the darker starfield toward the poles/edges.
cx, cy = W * 0.42, H * 0.46
dx = (xx - cx)
dx = np.minimum(np.abs(dx), W - np.abs(dx))  # wrap horizontally
dy = (yy - cy)
r_warp = 1.0 + (oct_b - 0.5) * 0.55 + (oct_c - 0.5) * 0.35
r = np.sqrt((dx / (W * 0.30)) ** 2 + (dy / (H * 0.34)) ** 2) * r_warp
mask = np.clip(1.0 - r, 0.0, 1.0) ** 1.4

brightness_bias = mask ** 1.3 * 0.88
density = np.clip(density * (0.45 + 0.25 * mask) + brightness_bias, 0.0, 1.0)

print("color grading...")
# Colour ramp control points (t, RGB) sampling the Orion-nebula palette: dark void ->
# blue-violet haze -> magenta wisps -> vivid pink body -> bright pink-white core.
stops = [
    (0.00, (0.010, 0.008, 0.018)),
    (0.18, (0.045, 0.035, 0.095)),
    (0.35, (0.16, 0.09, 0.28)),
    (0.55, (0.42, 0.14, 0.42)),
    (0.72, (0.75, 0.22, 0.48)),
    (0.87, (0.93, 0.45, 0.60)),
    (1.00, (0.99, 0.86, 0.88)),
]


def ramp(t):
    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


color = ramp(density)

print("adding core color variation...")
# Real emission nebulae mix H-alpha (magenta/red) and OIII (faint teal/cyan) regions
# rather than being one flat hue at a given brightness -- nudge hue in slow-varying
# patches, confined to the bright body, on top of the brightness-only ramp above.
hue_field = bilinear_sample(fft_field(2.4, seed=61), xx + warp_x * 0.9, yy + warp_y * 0.9)
hue_field = (hue_field - 0.5) * 2.0  # -1..1
warm_tint = np.array([0.25, 0.02, -0.05])
cool_tint = np.array([-0.12, 0.10, 0.06])
hue_variation = np.where(
    hue_field[..., None] > 0,
    hue_field[..., None] * warm_tint[None, None, :],
    -hue_field[..., None] * cool_tint[None, None, :],
)
core_variation_strength = (mask ** 0.8) * 0.22
color = np.clip(color + hue_variation * core_variation_strength[..., None], 0.0, 1.0)

# Blue-violet outer haze wash, additive, strongest away from the bright core.
haze_color = np.array([0.10, 0.11, 0.30])
haze_amount = (haze * 0.5 + 0.5) * (1.0 - mask) * 0.35
color += haze_amount[..., None] * haze_color[None, None, :]

# Dust lanes: dark silhouette streaks multiplied into the bright body only.
dust_strength = (dust * mask * 0.85)[..., None]
color *= (1.0 - dust_strength * 0.9)

# Two more dust-lane layers at different scales, weaker than the primary lane. The
# wide lane is allowed to bleed slightly past the core edge (mask relaxed) since real
# dust lanes don't stop precisely at a nebula's bright-body boundary.
dust_fine_strength = (dust_fine * mask * 0.55)[..., None]
color *= (1.0 - dust_fine_strength * 0.55)

dust_wide_strength = (dust_wide * (mask * 0.7 + 0.15) * 0.5)[..., None]
color *= (1.0 - dust_wide_strength * 0.45)

color = np.clip(color, 0.0, 1.0)

print("adding stars...")
rng_stars = np.random.default_rng(777)
n_stars = 5500
star_x = rng_stars.uniform(0, W, n_stars)
star_y_lat = rng_stars.uniform(-1, 1, n_stars)
star_y = (np.arcsin(star_y_lat) / np.pi + 0.5) * H  # denser near equator, thin near poles
brightness = rng_stars.power(2.2, n_stars)  # mostly dim, few bright

star_layer = np.zeros((H, W), dtype=np.float64)
xi = star_x.astype(np.int64) % W
yi = np.clip(star_y.astype(np.int64), 0, H - 1)
star_layer[yi, xi] = np.maximum(star_layer[yi, xi], brightness)

star_img = Image.fromarray((np.clip(star_layer, 0, 1) * 255).astype(np.uint8), "L")
star_img = star_img.filter(ImageFilter.GaussianBlur(0.6))

star_rgb = np.stack([np.array(star_img)] * 3, axis=-1).astype(np.float64) / 255.0
warm = np.array([1.0, 0.97, 0.90])
cool = np.array([0.85, 0.92, 1.0])
tint = rng_stars.uniform(0, 1, (H, W))
tint_color = warm[None, None, :] * (1 - tint[..., None]) + cool[None, None, :] * tint[..., None]
color = np.clip(color + star_rgb * tint_color, 0.0, 1.0)

print("drawing hero stars with diffraction spikes...")
draw_img = Image.fromarray((color * 255).astype(np.uint8), "RGB").convert("RGBA")
spike_layer = Image.new("RGBA", (W, H), (0, 0, 0, 0))
draw = ImageDraw.Draw(spike_layer)
hero_idx = np.argsort(-brightness)[:14]
# Each hero star gets its own randomized diffraction-spike stamp -- rotation, spike
# length, and brightness all jittered per star -- instead of one fixed shape reused at
# every location (that repetition was the single most obvious "amateur" tell up close).
rng_spike = np.random.default_rng(888)
for i in hero_idx:
    sx, sy = float(star_x[i]), float(star_y[i])
    b = float(brightness[i])
    n_spikes = int(rng_spike.choice([2, 4]))  # 4-arm cross or 8-arm starburst
    base_angle = rng_spike.uniform(0.0, math.pi)
    length = (26 + b * 70) * rng_spike.uniform(0.7, 1.3)
    core_r = (2.5 + b * 4.5) * rng_spike.uniform(0.85, 1.15)
    alpha = int(np.clip((180 + 70 * b) * rng_spike.uniform(0.8, 1.15), 0, 255))
    col = (255, 250, 245, alpha)
    for k in range(n_spikes):
        angle = base_angle + k * (math.pi / n_spikes)
        ax, ay = math.cos(angle) * length, math.sin(angle) * length
        draw.line([(sx - ax, sy - ay), (sx + ax, sy + ay)], fill=col, width=1)
    draw.ellipse([sx - core_r, sy - core_r, sx + core_r, sy + core_r], fill=(255, 253, 250, 255))
spike_layer = spike_layer.filter(ImageFilter.GaussianBlur(0.8))
draw_img = Image.alpha_composite(draw_img, spike_layer).convert("RGB")

print("baking bloom...")
arr = np.array(draw_img).astype(np.float64) / 255.0
luma = arr[..., 0] * 0.3 + arr[..., 1] * 0.59 + arr[..., 2] * 0.11
bright_mask = np.clip((luma - 0.55) / 0.45, 0, 1) ** 1.5
bright_img = Image.fromarray((arr * bright_mask[..., None] * 255).astype(np.uint8), "RGB")
bloom = bright_img.filter(ImageFilter.GaussianBlur(10))
bloom_arr = np.array(bloom).astype(np.float64) / 255.0
final = 1.0 - (1.0 - arr) * (1.0 - bloom_arr * 0.9)  # screen blend
final = np.clip(final, 0.0, 1.0)

print("saving final nebula sky...")
out_path = Path(__file__).resolve().parents[2] / "Game" / "assets" / "textures" / "sky_nebula.png"
Image.fromarray((final * 255).astype(np.uint8), "RGB").save(out_path)
print("done:", out_path)