Water ripples

local natiny = require("natiny")
local demo = require("common.modules.demo")

local WIDTH, HEIGHT = 900, 980
local SIM_W, SIM_H = 384, 512
local STEP, MAX_STEPS = 1.0 / 240.0, 1

local BG     = {  21 / 255,  30 / 255,  44 / 255 } -- #151E2C
local OBJECT = { 169 / 255, 179 / 255, 197 / 255 } -- #A9B3C5
local LIGHT_TEXT = { 243 / 255, 245 / 255, 250 / 255 } -- #F3F5FA

local VERTEX = [[
attribute vec3 position;
attribute vec2 texcoord0;
attribute vec4 color0;

varying vec2 uv;
varying vec4 tint;

uniform mat4 mtx_worldviewproj;

void vertex()
{
    uv = texcoord0;
    tint = color0;
    POSITION = mtx_worldviewproj * vec4(position, 1.0);
}
]]

local SIM = [[
uniform sampler2D state_map;
uniform vec4 texel_size;
uniform vec4 drop;
uniform vec4 plunk;
uniform vec4 sim_params;

vec2 tap(vec2 p, vec2 d)
{
    vec2 q = p + d * texel_size.xy;
    if (q.x < 0.0) q.x = p.x + texel_size.x;
    if (q.x > 1.0) q.x = p.x - texel_size.x;
    if (q.y < 0.0) q.y = p.y + texel_size.y;
    if (q.y > 1.0) q.y = p.y - texel_size.y;
    return texture(state_map, q).rg;
}

float splash(vec2 p, vec4 s, float radius)
{
    vec2 d = p - s.xy;
    d.x = d.x * texel_size.w;
    float t = length(d) / max(radius, 0.0001);
    float f = 1.0 - smoothstep(0.0, 1.0, t);
    return f * f * cos(t * 7.4) * s.w;
}

void fragment()
{
    vec2 c = texture(state_map, uv).rg;
    vec2 a = (tap(uv, vec2(-1.0, 0.0)) + tap(uv, vec2(1.0, 0.0))
            + tap(uv, vec2(0.0, -1.0)) + tap(uv, vec2(0.0, 1.0))) * 0.2
           + (tap(uv, vec2(-1.0, -1.0)) + tap(uv, vec2(1.0, -1.0))
            + tap(uv, vec2(-1.0, 1.0)) + tap(uv, vec2(1.0, 1.0))) * 0.05;

    float v = c.r - c.g;
    float next = c.r + (v + sim_params.x * (a.r - c.r)
                          + sim_params.z * ((a.r - a.g) - v)) * sim_params.y;

    if (drop.z > 0.5) next = next + splash(uv, drop, sim_params.w);
    if (plunk.z > 0.0) next = next + splash(uv, plunk, plunk.z);

    COLOR = vec4(next, c.r, 0.0, 1.0);
}
]]

local DISPLAY = [[
uniform sampler2D tex0;
uniform sampler2D wave_map;
uniform vec4 wave_texel;
uniform vec4 water_look;
uniform vec4 water_chop;
uniform vec4 water_light;
uniform vec4 water_body;
uniform vec4 water_sky;

float height(vec2 p)
{
    return texture(wave_map, clamp(p, wave_texel.xy * 1.5,
                                   vec2(1.0) - wave_texel.xy * 1.5)).r;
}

vec2 chop(vec2 p, vec2 d, float frequency, float speed, float amount)
{
    return d * (amount * cos(dot(p, d) * frequency * water_chop.z
                             + water_chop.x * speed * water_chop.w));
}

void fragment()
{
    float c = height(uv);
    float w = height(uv - vec2(wave_texel.x, 0.0));
    float e = height(uv + vec2(wave_texel.x, 0.0));
    float n = height(uv - vec2(0.0, wave_texel.y));
    float s = height(uv + vec2(0.0, wave_texel.y));

    vec2 gradient = vec2(e - w, s - n) * 0.5;
    float stir = clamp(length(gradient) * 260.0, 0.0, 1.0);

    vec2 slope = gradient * water_look.x
               + (chop(uv, vec2( 0.92,  0.39), 118.0, 2.3, 1.00)
                + chop(uv, vec2(-0.42,  0.91), 187.0, 3.1, 0.66)
                + chop(uv, vec2( 0.71, -0.70), 301.0, 4.3, 0.42)
                + chop(uv, vec2(-0.98, -0.19), 467.0, 5.8, 0.24))
               * (water_chop.y * stir);

    vec2 bend = slope * water_look.y;
    vec2 lo = vec2(0.0015, 0.0015);
    vec2 hi = vec2(0.9985, 0.9985);
    vec3 color = vec3(
        texture(tex0, clamp(uv - bend * (1.0 + water_look.z), lo, hi)).r,
        texture(tex0, clamp(uv - bend, lo, hi)).g,
        texture(tex0, clamp(uv - bend * (1.0 - water_look.z), lo, hi)).b);

    color = color * clamp(1.0 + ((w + e + n + s) * 0.25 - c) * water_look.w,
                          0.55, 2.4);
    color = color * mix(vec3(1.0), water_body.rgb,
                        clamp(0.5 + c * water_body.w, 0.0, 1.0));
    color = mix(color, water_sky.rgb,
                clamp(0.5 - dot(slope, water_light.xy) * 3.2, 0.0, 1.0)
                * water_sky.w);

    vec3 light = normalize(water_light.xyz);
    vec3 half_way = normalize(vec3(light.x, light.y, light.z + 1.0));
    float lit = max(dot(normalize(vec3(-slope.x, -slope.y, 1.0)), half_way), 0.0);
    float flat_water = max(half_way.z, 0.0);

    color = color + vec3(
        max(pow(lit, water_light.w) - pow(flat_water, water_light.w), 0.0) * 0.45
      + max(pow(lit, 22.0) - pow(flat_water, 22.0), 0.0) * 0.12);

    COLOR = vec4(color, 1.0) * tint;
}
]]

local function material(fragment)
    local source = natiny.resource.from_data(VERTEX .. fragment)
    local shader = natiny.shader.load(source)
    natiny.resource.destroy(source)
    return shader, natiny.material.create(shader)
end

if not natiny.backend.init(arg and arg[1] or natiny.backend.AUTO) then
    os.exit(1)
end

local window = natiny.window.create("Natiny - water ripples", WIDTH, HEIGHT)
natiny.window.set_scale_mode(window, natiny.window.SCALE_MODE_STRETCH)

-- The picture and the face arrive together; until they do the window is the
-- shell's loading screen, and the simulation below is only set up, not run.
local image, image_w, image_h, font

local sim_shader, sim = material(SIM)
local display_shader, display = material(DISPLAY)

local state_a = natiny.surface.create(SIM_W, SIM_H, natiny.surface.RGBA16F, false)
local state_b = natiny.surface.create(SIM_W, SIM_H, natiny.surface.RGBA16F, false)
local current, spare = state_a, state_b

natiny.material.set_constant(sim, "texel_size",
    1.0 / SIM_W, 1.0 / SIM_H, SIM_W, SIM_W / SIM_H)
natiny.material.set_constant(sim, "drop", 0, 0, 0, 0)
natiny.material.set_constant(sim, "plunk", 0, 0, 0, 0)

-- speed squared (limit 2.5), velocity damping per step, viscosity,
-- radius of the ring the held pointer draws
natiny.material.set_constant(sim, "sim_params", 1.8, 0.9990, 0.055, 0.030)

natiny.material.set_constant(display, "wave_texel",
    1.0 / SIM_W, 1.0 / SIM_H, SIM_W, SIM_H)
-- slope gain, refraction, dispersion between the channels, caustics
natiny.material.set_constant(display, "water_look", 16.0, 0.042, 0.05, 3.0)
-- seconds, chop amount, chop scale, chop speed
natiny.material.set_constant(display, "water_chop", 0.0, 0.018, 1.0, 1.0)
-- towards the lamp, highlight tightness
natiny.material.set_constant(display, "water_light", -0.42, -0.56, 0.72, 180.0)
-- what the water takes out of the picture, and how fast with depth
natiny.material.set_constant(display, "water_body", 0.90, 0.975, 1.0, 1.6)
-- sheen colour and how much of it
natiny.material.set_constant(display, "water_sky", 0.62, 0.74, 0.92, 0.10)

local ready, accumulator, elapsed = false, 0.0, 0.0
local emitter, held, pending, phase = nil, false, false, 0.0
local splashes = {}

local function rectangle()
    local ww = natiny.window.get_width(window)
    local wh = natiny.window.get_height(window)
    local h = math.min(wh * 0.84, math.max(1, ww - 64) * image_h / image_w)
    local w = h * image_w / image_h
    return (ww - w) * 0.5, (wh - h) * 0.5 + 12, w, h, ww, wh
end

local function calm()
    for _, state in ipairs({ state_a, state_b }) do
        natiny.surface.bind(state, function() natiny.render.clear(0, 0, 0, 1) end)
    end
    current, spare, splashes, ready = state_a, state_b, {}, true
end

local function simulate()
    natiny.material.set_texture(sim, "state_map", natiny.surface.get_texture(current),
        natiny.material.FILTER_NEAREST, natiny.material.FILTER_NEAREST)

    if emitter and (held or pending) then
        natiny.material.set_constant(sim, "drop", emitter.x, emitter.y, 1,
            math.cos(phase * math.pi * 6.0) * 0.055)
        phase = phase + STEP
    else
        natiny.material.set_constant(sim, "drop", 0, 0, 0, 0)
    end

    local splash = table.remove(splashes, 1) or { 0, 0, 0, 0 }
    natiny.material.set_constant(sim, "plunk",
        splash[1], splash[2], splash[3], splash[4])

    natiny.surface.bind(spare, function()
        natiny.render.set_color(1, 1, 1, 1)
        natiny.material.bind(sim, function()
            natiny.render.surface(current, 0, 0, SIM_W, SIM_H)
        end)
    end)

    current, spare, pending = spare, current, false
end

local function label(text, ww, y, r, g, b)
    natiny.render.set_color(r, g, b, 1)
    natiny.render.text(font, (ww - natiny.font.measure_text(font, text)) * 0.5, y, text)
end

local function frame(dt)
    natiny.window.bind(window, function()
        local x, y, w, h, ww = rectangle()

        elapsed = elapsed + dt
        natiny.material.set_constant(display, "water_chop", elapsed, 0.018, 1.0, 1.0)

        if natiny.input.get_key_pressed(natiny.input.KEY_R) then ready = false end

        held = false
        if natiny.input.get_mouse_button_down(natiny.input.MOUSE_BUTTON_LEFT) then
            local mx, my = natiny.input.get_mouse_x(), natiny.input.get_mouse_y()
            if mx >= x and mx <= x + w and my >= y and my <= y + h then
                emitter = { x = (mx - x) / w, y = (my - y) / h }
                held, pending = true, true
                if natiny.input.get_mouse_button_pressed(natiny.input.MOUSE_BUTTON_LEFT) then
                    phase = 0.0
                    splashes[#splashes + 1] = { emitter.x, emitter.y, 0.055, -0.18 }
                end
            end
        end

        accumulator = math.min(accumulator + math.min(dt, 0.1), STEP * MAX_STEPS)

        if not ready then calm() end

        local steps = 0
        while accumulator >= STEP and steps < MAX_STEPS do
            simulate()
            accumulator, steps = accumulator - STEP, steps + 1
        end

        natiny.render.clear(BG[1], BG[2], BG[3], 1)

        natiny.material.set_texture(display, "wave_map",
            natiny.surface.get_texture(current),
            natiny.material.FILTER_LINEAR, natiny.material.FILTER_LINEAR)

        natiny.render.set_color(OBJECT[1], OBJECT[2], OBJECT[3], 0.45)
        natiny.render.rounded_rectangle(x - 10, y - 10, w + 20, h + 20, 12)

        natiny.render.set_color(1, 1, 1, 1)
        natiny.material.bind(display, function()
            natiny.render.texture(image, x, y, w, h)
        end)

        label("2D WATER / GPU HEIGHT FIELD", ww, math.max(10, y - 44),
              LIGHT_TEXT[1], LIGHT_TEXT[2], LIGHT_TEXT[3])
    end)
end

demo:run({
    window = window,
    title  = "Water ripples",
    text   = "Click and hold the picture - R calms it",
    resources = { "data/common/textures/saikosha.jpg", "data/common/fonts/roboto.ttf" },

    on_loaded = function()
        image = natiny.texture.create(demo:resource("data/common/textures/saikosha.jpg"))
        image_w = natiny.texture.get_width(image)
        image_h = natiny.texture.get_height(image)
        font = natiny.font.create(demo:resource("data/common/fonts/roboto.ttf"), 18)
    end,

    loop = frame,
})

natiny.backend.shutdown()