#include <metal_stdlib>
using namespace metal;

// Seascape — Alexander Alekseev (TDM, 2014)
// MSL port of seascape.vk.glsl. iTime threaded through helpers because
// MSL does not allow file-scope mutable globals.

struct ShadertoyUBO {
    float iTime;
    float2 iResolution;
};

struct PassthroughVOut {
    float4 pos [[position]];
    float2 uv;
};

constant int NUM_STEPS = 32;
constant float PI = 3.141592;
constant float EPSILON = 1e-3;

constant int ITER_GEOMETRY = 3;
constant int ITER_FRAGMENT = 5;
constant float SEA_HEIGHT = 0.6;
constant float SEA_CHOPPY = 4.0;
constant float SEA_SPEED = 0.8;
constant float SEA_FREQ = 0.16;
constant float3 SEA_BASE = float3(0.0, 0.09, 0.18);
constant float3 SEA_WATER_COLOR = float3(0.8, 0.9, 0.6) * 0.6;
constant float2x2 octave_m = float2x2(1.6, 1.2, -1.2, 1.6);

static float3x3 fromEuler(float3 ang) {
    float2 a1 = float2(sin(ang.x), cos(ang.x));
    float2 a2 = float2(sin(ang.y), cos(ang.y));
    float2 a3 = float2(sin(ang.z), cos(ang.z));
    float3x3 m;
    m[0] = float3(a1.y * a3.y + a1.x * a2.x * a3.x, a1.y * a2.x * a3.x + a3.y * a1.x, -a2.y * a3.x);
    m[1] = float3(-a2.y * a1.x, a1.y * a2.y, a2.x);
    m[2] = float3(a3.y * a1.x * a2.x + a1.y * a3.x, a1.x * a3.x - a1.y * a3.y * a2.x, a2.y * a3.y);
    return m;
}

static float hash(float2 p) {
    float h = dot(p, float2(127.1, 311.7));
    return fract(sin(h) * 43758.5453123);
}

static float noise(float2 p) {
    float2 i = floor(p);
    float2 f = fract(p);
    float2 u = f * f * (3.0 - 2.0 * f);
    return -1.0 + 2.0 * mix(mix(hash(i + float2(0.0, 0.0)),
                                hash(i + float2(1.0, 0.0)), u.x),
                            mix(hash(i + float2(0.0, 1.0)),
                                hash(i + float2(1.0, 1.0)), u.x), u.y);
}

static float diffuseLight(float3 n, float3 l, float p) {
    return pow(dot(n, l) * 0.4 + 0.6, p);
}

static float specularLight(float3 n, float3 l, float3 e, float s) {
    float nrm = (s + 8.0) / (PI * 8.0);
    return pow(max(dot(reflect(e, n), l), 0.0), s) * nrm;
}

static float3 getSkyColor(float3 e) {
    e.y = (max(e.y, 0.0) * 0.8 + 0.2) * 0.8;
    return float3(pow(1.0 - e.y, 2.0), 1.0 - e.y, 0.6 + (1.0 - e.y) * 0.4) * 1.1;
}

static float sea_octave(float2 uv, float choppy) {
    uv += noise(uv);
    float2 wv = 1.0 - abs(sin(uv));
    float2 swv = abs(cos(uv));
    wv = mix(wv, swv, wv);
    return pow(1.0 - pow(wv.x * wv.y, 0.65), choppy);
}

static float seaMap(float3 p, float iTime, int iter) {
    float SEA_TIME = 1.0 + iTime * SEA_SPEED;
    float freq = SEA_FREQ;
    float amp = SEA_HEIGHT;
    float choppy = SEA_CHOPPY;
    float2 uv = p.xz; uv.x *= 0.75;

    float d, h = 0.0;
    for (int i = 0; i < iter; i++) {
        d = sea_octave((uv + SEA_TIME) * freq, choppy);
        d += sea_octave((uv - SEA_TIME) * freq, choppy);
        h += d * amp;
        uv = uv * octave_m;
        freq *= 1.9;
        amp *= 0.22;
        choppy = mix(choppy, 1.0, 0.2);
    }
    return p.y - h;
}

static float3 getSeaColor(float3 p, float3 n, float3 l, float3 eye, float3 dist) {
    float fresnel = clamp(1.0 - dot(n, -eye), 0.0, 1.0);
    fresnel = min(fresnel * fresnel * fresnel, 0.5);

    float3 reflected = getSkyColor(reflect(eye, n));
    float3 refracted = SEA_BASE + diffuseLight(n, l, 80.0) * SEA_WATER_COLOR * 0.12;

    float3 color = mix(refracted, reflected, fresnel);

    float atten = max(1.0 - dot(dist, dist) * 0.001, 0.0);
    color += SEA_WATER_COLOR * (p.y - SEA_HEIGHT) * 0.18 * atten;

    color += specularLight(n, l, eye, 600.0 * rsqrt(dot(dist, dist)));

    return color;
}

static float3 getNormal(float3 p, float eps, float iTime) {
    float3 n;
    n.y = seaMap(p, iTime, ITER_FRAGMENT);
    n.x = seaMap(float3(p.x + eps, p.y, p.z), iTime, ITER_FRAGMENT) - n.y;
    n.z = seaMap(float3(p.x, p.y, p.z + eps), iTime, ITER_FRAGMENT) - n.y;
    n.y = eps;
    return normalize(n);
}

static float heightMapTracing(float3 ori, float3 dir, thread float3& p, float iTime) {
    float tm = 0.0;
    float tx = 1000.0;
    float hx = seaMap(ori + dir * tx, iTime, ITER_GEOMETRY);
    if (hx > 0.0) {
        p = ori + dir * tx;
        return tx;
    }
    float hm = seaMap(ori, iTime, ITER_GEOMETRY);
    for (int i = 0; i < NUM_STEPS; i++) {
        float tmid = mix(tm, tx, hm / (hm - hx));
        p = ori + dir * tmid;
        float hmid = seaMap(p, iTime, ITER_GEOMETRY);
        if (hmid < 0.0) {
            tx = tmid;
            hx = hmid;
        } else {
            tm = tmid;
            hm = hmid;
        }
        if (abs(hmid) < EPSILON) break;
    }
    return mix(tm, tx, hm / (hm - hx));
}

static float3 getPixel(float2 coord, float time, float2 iResolution, float iTime) {
    float2 uv = coord / iResolution;
    uv = uv * 2.0 - 1.0;
    uv.x *= iResolution.x / iResolution.y;

    float3 ang = float3(sin(time * 3.0) * 0.1, sin(time) * 0.2 + 0.3, time);
    float3 ori = float3(0.0, 3.5, time * 5.0);
    float3 dir = normalize(float3(uv.xy, -2.0));
    dir.z += length(uv) * 0.14;
    dir = normalize(dir) * fromEuler(ang);

    float3 p;
    heightMapTracing(ori, dir, p, iTime);
    float3 dist = p - ori;
    float EPSILON_NRM = 0.1 / iResolution.x;
    float3 n = getNormal(p, dot(dist, dist) * EPSILON_NRM, iTime);
    float3 light = normalize(float3(0.0, 1.0, 0.8));

    return mix(
        getSkyColor(dir),
        getSeaColor(p, n, light, dir, dist),
        pow(smoothstep(0.0, -0.02, dir.y), 0.2));
}

fragment float4 seascape_fs(PassthroughVOut in [[stage_in]],
                            constant ShadertoyUBO& U [[buffer(0)]]) {
    float2 fragCoord = in.uv * U.iResolution;
    float time = U.iTime * 0.3;
    float3 color = getPixel(fragCoord, time, U.iResolution, U.iTime);
    return float4(pow(color, float3(0.65)), 1.0);
}