# Draws the ground stain shapes for ImpactFx3D (Sheet row 66).
# White shapes on transparent background, tinted in game.
# Usage: python tools/dev/draw_impact_stains.py assets/textures/impact
# Rendered at 4x then downsampled for clean edges.
import math
import random
import sys
from PIL import Image, ImageDraw, ImageFilter

OUT = sys.argv[1]
SIZE = 256
SCALE = 4
BIG = SIZE * SCALE


def polar_blob(draw, cx, cy, radius, lobes, rng, spikes=0, spike_len=0.0, points=360):
    # Irregular closed shape: low-frequency wobble plus optional thin spikes.
    phases = [(k, rng.uniform(0, math.tau), rng.uniform(0.02, 0.10) / (1 + k * 0.35)) for k in range(2, lobes + 2)]
    spike_angles = [rng.uniform(0, math.tau) for _ in range(spikes)]
    spike_sizes = [rng.uniform(0.2, 1.0) for _ in range(spikes)]
    spike_widths = [rng.uniform(0.035, 0.08) for _ in range(spikes)]
    polygon = []
    for i in range(points):
        a = i / points * math.tau
        r = 1.0 + sum(amp * math.sin(k * a + ph) for k, ph, amp in phases)
        for sa, ss, sw in zip(spike_angles, spike_sizes, spike_widths):
            d = math.atan2(math.sin(a - sa), math.cos(a - sa))
            r += spike_len * ss * math.exp(-abs(d / sw) ** 1.3)
        polygon.append((cx + math.cos(a) * radius * r, cy + math.sin(a) * radius * r))
    draw.polygon(polygon, fill=255)
    return spike_angles, spike_sizes


def ellipse_along(draw, cx, cy, along, across, angle):
    # Ellipse stretched along `angle` (radians), drawn as a polygon.
    polygon = []
    for i in range(48):
        t = i / 48 * math.tau
        x, y = math.cos(t) * along, math.sin(t) * across
        polygon.append((cx + x * math.cos(angle) - y * math.sin(angle), cy + x * math.sin(angle) + y * math.cos(angle)))
    draw.polygon(polygon, fill=255)


def melt(mask, radius):
    # Blur then threshold: shapes that touch flow into each other like liquid.
    return mask.filter(ImageFilter.GaussianBlur(radius * SCALE)).point(lambda v: 255 if v >= 128 else 0)


def finish(mask, name):
    # Downsample, then shade: a slightly darker rim reads as a liquid edge.
    small = mask.resize((SIZE, SIZE), Image.LANCZOS)
    inner = small.filter(ImageFilter.GaussianBlur(3))
    image = Image.new('RGBA', (SIZE, SIZE))
    pixels = image.load()
    a_px = small.load()
    i_px = inner.load()
    for y in range(SIZE):
        for x in range(SIZE):
            alpha = a_px[x, y]
            if alpha == 0:
                pixels[x, y] = (255, 255, 255, 0)
                continue
            depth = i_px[x, y] / 255.0
            tone = int(255 * (0.80 + 0.20 * depth))
            pixels[x, y] = (tone, tone, tone, alpha)
    image.save(f'{OUT}/{name}.png')


def splat(name, seed):
    rng = random.Random(seed)
    mask = Image.new('L', (BIG, BIG), 0)
    draw = ImageDraw.Draw(mask)
    c = BIG / 2
    radius = BIG * 0.20
    angles, sizes = polar_blob(draw, c, c, radius, 5, rng, spikes=rng.randint(7, 11), spike_len=0.55)
    # Droplets flung past the spikes, stretched outward.
    for angle, size in zip(angles, sizes):
        if rng.random() < 0.75:
            dist = radius * (1.75 + size * rng.uniform(0.2, 0.9))
            r = BIG * rng.uniform(0.012, 0.028)
            ellipse_along(draw, c + math.cos(angle) * dist, c + math.sin(angle) * dist, r * 1.8, r, angle)
    for _ in range(rng.randint(10, 16)):
        angle = rng.uniform(0, math.tau)
        dist = radius * rng.uniform(1.35, 2.25)
        r = BIG * rng.uniform(0.006, 0.018)
        ellipse_along(draw, c + math.cos(angle) * dist, c + math.sin(angle) * dist, r * 1.4, r, angle)
    finish(melt(mask, 1.5), name)


def streak(name, seed):
    # A drop that hit the ground moving toward the top of the image: round head
    # at the bottom, tail thinning toward the travel direction.
    rng = random.Random(seed)
    mask = Image.new('L', (BIG, BIG), 0)
    draw = ImageDraw.Draw(mask)
    cx = BIG / 2
    head_y = BIG * 0.68
    head_r = BIG * 0.095
    polygon = []
    for i in range(240):
        t = i / 240 * math.tau
        # Half circle below, long taper above.
        x = math.sin(t) * head_r
        y = math.cos(t)
        if y >= 0:
            py = head_y + y * head_r * 1.2
        else:
            py = head_y + y * BIG * 0.46
            x *= (1 + y * 0.88) ** 2.4 + 0.015
        x *= 1 + 0.05 * math.sin(5 * t + seed)
        polygon.append((cx + x, py))
    draw.polygon(polygon, fill=255)
    # A few satellite droplets ahead of the tail.
    for k in range(3):
        y = BIG * (0.15 - k * 0.045) + rng.uniform(-6, 6) * SCALE
        r = BIG * (0.022 - k * 0.005)
        ellipse_along(draw, cx + rng.uniform(-8, 8) * SCALE, y, r * 1.6, r, -math.pi / 2)
    finish(melt(mask, 1.0), name)


def pool(name, seed):
    rng = random.Random(seed)
    mask = Image.new('L', (BIG, BIG), 0)
    draw = ImageDraw.Draw(mask)
    c = BIG / 2
    polar_blob(draw, c, c, BIG * 0.36, 6, rng)
    # A couple of side lobes where the liquid ran.
    for _ in range(3):
        angle = rng.uniform(0, math.tau)
        dist = BIG * rng.uniform(0.22, 0.29)
        r = BIG * rng.uniform(0.08, 0.12)
        draw.ellipse((c + math.cos(angle) * dist - r, c + math.sin(angle) * dist - r,
                      c + math.cos(angle) * dist + r, c + math.sin(angle) * dist + r), fill=255)
    finish(melt(mask, 7.0), name)


splat('drawn_splat_a', 11)
splat('drawn_splat_b', 23)
splat('drawn_splat_c', 37)
streak('drawn_streak', 5)
pool('drawn_pool', 8)
# The streak is cropped to its shape: the decal maps the whole image.
streak_image = Image.open(f'{OUT}/drawn_streak.png')
box = streak_image.getbbox()
streak_image.crop((max(0, box[0] - 6), max(0, box[1] - 6), min(SIZE, box[2] + 6), min(SIZE, box[3] + 6))).save(f'{OUT}/drawn_streak.png')
print('ok')
