///////////////////////////////////////////////////
// Compute
///////////////////////////////////////////////////


struct Params { 
    nx: f32, 
    ny: f32, 
    dx: f32, 
    dt: f32, 
    g: f32,
    waves: f32,
};

@group(0) @binding(0) var<storage, read_write> hh : array<f32>;
@group(0) @binding(1) var<storage, read_write> uu : array<f32>;
@group(0) @binding(2) var<storage, read_write> vv : array<f32>;
@group(0) @binding(3) var<storage, read_write> hh2 : array<f32>;
@group(0) @binding(4) var<storage, read_write> uu2 : array<f32>;
@group(0) @binding(5) var<storage, read_write> vv2 : array<f32>;
@group(0) @binding(6) var<storage, read> dd : array<f32>;
@group(0) @binding(7) var<uniform> params : Params;
@group(0) @binding(8) var<storage, read_write> t : f32;

fn idx(i: u32, j: u32) -> u32 {
    return i * u32(params.ny) + j;
}

fn waves_bc_hh(i: f32, j: f32, t: f32) -> f32 {
    return 0.025 * sin( 0.05*i + 0.01 * j + 10*t ) * min(1.0, 0.005 * (t / params.dt));
}

@compute @workgroup_size(16, 16)
fn bcs(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    let j = gid.y;
    let nx = u32(params.nx);
    let ny = u32(params.ny);
    let id = idx(i, j);

    uu[0] = 0 * uu[0] * hh[0] * vv[0] * hh2[0] * vv2[0] * uu2[0] * dd[0] * t;

    // matched gradient
    // const a = 2;
    // const b = -1;

    // zero gradient
    const a = 1;
    const b = 0;

    let c = sqrt(params.g * dd[idx(i,j)]); // Approx; depth is changing.
    let f = (params.dt * c / params.dx);
    
    if (j == 0) {
        // hh2[id] = a * hh2[idx(i, 1)] + b * hh2[idx(i, 2)];
        // uu2[id] = a * uu2[idx(i, 1)] + b * uu2[idx(i, 2)];
        // vv2[id] = a * vv2[idx(i, 1)] + b * vv2[idx(i, 2)];
        hh2[id] = hh2[idx(i, 0)] + f * (hh2[idx(i, 1)] - hh2[idx(i, 0)]);
        uu2[id] = hh2[idx(i, 0)] + f * (hh2[idx(i, 1)] - hh2[idx(i, 0)]);
        vv2[id] = hh2[idx(i, 0)] + f * (hh2[idx(i, 1)] - hh2[idx(i, 0)]);
        // uu2[id] = -uu2[id] * params.g / c;
        // vv2[id] = -vv2[id] * params.g / c;
    }
    else if (j == ny - 1) {
        hh2[id] = a * hh2[idx(i, ny-2)] + b * hh2[idx(i, ny-3)];
        uu2[id] = a * uu2[idx(i, ny-2)] + b * uu2[idx(i, ny-3)];
        vv2[id] = a * vv2[idx(i, ny-2)] + b * vv2[idx(i, ny-3)];
    }

    if (i==0) {
        hh2[id] = a * hh2[idx(1, j)] + b * hh2[idx(2, j)];
        uu2[id] = a * uu2[idx(1, j)] + b * uu2[idx(2, j)];
        vv2[id] = a * vv2[idx(1, j)] + b * vv2[idx(2, j)];
    }
    else if (i == nx-1) {
        hh2[id] = a * hh2[idx(nx-2, j)] + b * hh2[idx(nx-3, j)];
        uu2[id] = a * uu2[idx(nx-2, j)] + b * uu2[idx(nx-3, j)];
        vv2[id] = a * vv2[idx(nx-2, j)] + b * vv2[idx(nx-3, j)];
    }


   if (j == 0) {
        // hh[id] = a * hh[idx(i, 1)] + b * hh[idx(i, 2)];
        // uu[id] = a * uu[idx(i, 1)] + b * uu[idx(i, 2)];
        // vv[id] = a * vv[idx(i, 1)] + b * vv[idx(i, 2)];
        hh[id] = (1 - f) * hh[idx(i, 0)] + f * hh[idx(i, 1)];
        uu[id] = (1 - f) * hh[idx(i, 0)] + f * hh[idx(i, 1)];
        vv[id] = (1 - f) * hh[idx(i, 0)] + f * hh[idx(i, 1)];
        // uu[id] = -uu[id] * params.g / c;
        // vv[id] = -vv[id] * params.g / c;
    }
    else if (j == ny - 1) {
        hh[id] = a * hh[idx(i, ny-2)] + b * hh[idx(i, ny-3)];
        uu[id] = a * uu[idx(i, ny-2)] + b * uu[idx(i, ny-3)];
        vv[id] = a * vv[idx(i, ny-2)] + b * vv[idx(i, ny-3)];
    }

    if (i==0) {
        hh[id] = a * hh[idx(1, j)] + b * hh[idx(2, j)];
        uu[id] = a * uu[idx(1, j)] + b * uu[idx(2, j)];
        vv[id] = a * vv[idx(1, j)] + b * vv[idx(2, j)];
    } 
    else if (i == nx-1) {
        hh[id] = a * hh[idx(nx-2, j)] + b * hh[idx(nx-3, j)];
        uu[id] = a * uu[idx(nx-2, j)] + b * uu[idx(nx-3, j)];
        vv[id] = a * vv[idx(nx-2, j)] + b * vv[idx(nx-3, j)];
    }
    

    if (j == 0 || j == ny - 1 || i == 0 || i == nx - 1) {
        hh[id] = 0;
        hh2[id] = 0;
    }

    t = t;
    if (j==0) {
        hh[id] = params.waves * waves_bc_hh(f32(i), f32(j), t);
        hh2[id] = params.waves * waves_bc_hh(f32(i),f32(j), t);
    }
}

@compute @workgroup_size(16, 16)
fn halfstep(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    let j = gid.y;
    let nx = u32(params.nx);
    let ny = u32(params.ny);
    let id = idx(i, j);

    // Skip edges
    if (i < 1u || i >= nx-1u || j < 1u || j >= ny-1u) { return; }
    
    // Central differences for gradients
    let uudd_x = (uu[idx(i+1u,j)] * dd[idx(i+1u,j)] - uu[idx(i-1u,j)] * dd[idx(i-1u,j)]) / 2.0;
    let vvdd_y = (vv[idx(i,j+1u)] * dd[idx(i,j+1u)] - vv[idx(i,j-1u)] * dd[idx(i,j-1u)]) / 2.0;
    let dhh = -(uudd_x + vvdd_y) / params.dx;
    let duu = -params.g * ((hh[idx(i+1u,j)] - hh[idx(i-1u,j)]) / 2.0) / params.dx;
    let dvv = -params.g * ((hh[idx(i,j+1u)] - hh[idx(i,j-1u)]) / 2.0) / params.dx;

    // Half step 
    hh2[id] = hh[id] + dhh * params.dt;
    uu2[id] = uu[id] + duu * params.dt;
    vv2[id] = vv[id] + dvv * params.dt;  

    // Noop so we have t in the bindGroup and it matches fullstep.
    t = t;
}

@compute @workgroup_size(16, 16)
fn fullstep(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    let j = gid.y;
    let nx = u32(params.nx);
    let ny = u32(params.ny);
    let id = idx(i, j);

    if (j==0 && i==0) {
        t = t + params.dt;
    }

    // Skip edges
    if (i < 1u || i >= nx-1u || j < 1u || j >= ny-1u) { return; }

    // Recalc grads
    let uu2dd_x = (uu2[idx(i+1u,j)] * dd[idx(i+1u,j)] - uu2[idx(i-1u,j)] * dd[idx(i-1u,j)]) / 2.0;
    let vv2dd_y = (vv2[idx(i,j+1u)] * dd[idx(i,j+1u)] - vv2[idx(i,j-1u)] * dd[idx(i,j-1u)]) / 2.0;
    let dhh2 = -(uu2dd_x + vv2dd_y) / params.dx;
    let duu2 = -params.g * ((hh2[idx(i+1u,j)] - hh2[idx(i-1u,j)]) / 2.0) / params.dx;
    let dvv2 = -params.g * ((hh2[idx(i,j+1u)] - hh2[idx(i,j-1u)]) / 2.0) / params.dx;

    // Full step
    hh[id] = 0.5 * (hh[id] + hh2[id] + dhh2 * params.dt);
    uu[id] = 0.5 * (uu[id] + uu2[id] + duu2 * params.dt);
    vv[id] = 0.5 * (vv[id] + vv2[id] + dvv2 * params.dt);

    // Damping
    hh[id] = hh[id] * 0.9995;
    uu[id] = uu[id] * 0.9995;
    vv[id] = vv[id] * 0.9995;
}

///////////////////////////////////////////////////
// Rendering
///////////////////////////////////////////////////

@vertex
fn vs_main(@builtin(vertex_index) idx: u32) -> @builtin(position) vec4f {
    var pos = array<vec2f, 3>(
        // https://webgpufundamentals.org/webgpu/lessons/webgpu-large-triangle-to-cover-clip-space.html
        vec2f(-1.0, -1.0),
        vec2f( 3.0, -1.0),
        vec2f(-1.0,  3.0),
    );
    return vec4f(pos[idx], 0.0, 1.0);
}

fn lerp(x: f32, xmin: f32, xmax: f32) -> f32 {
    return clamp((x-xmin) / (xmax - xmin), 0, 1);
}

fn gradx(i: u32, j: u32) -> f32 {
    const w = 1.0;
    const u = 0.0;
    // return (hh[idx(i+1, j)] - hh[idx(i-1, j)]) / (2 * params.dx);
    return (
          (w*(hh[idx(i+2, j)] + hh[idx(i, j)] + hh[idx(i+1, j+1)] + hh[idx(i+1, j-1)]) + u*hh[idx(i+1, j)]) 
        - (w*(hh[idx(i, j)] + hh[idx(i-2, j)] + hh[idx(i-1, j+1)] + hh[idx(i-1, j-1)]) + u*hh[idx(i-1, j)])
    ) / ((u + 4*w) * (2 * params.dx));
}

fn grady(i: u32, j: u32) -> f32 {
    const w = 1.0;
    const u = 0.0;
    // return (hh[idx(i, j+1)] - hh[idx(i, j-1)]) / (2 * params.dx);
    return (
          (w*(hh[idx(i, j+2)] + hh[idx(i, j)] + hh[idx(i+1, j+1)] + hh[idx(i-1, j+1)]) + hh[idx(i, j+1)]) 
        - (w*(hh[idx(i, j)] + hh[idx(i, j-2)] + hh[idx(i+1, j-1)] + hh[idx(i-1, j-1)]) + hh[idx(i, j-1)])
    ) / ((u + 4*w) * (2 * params.dx));
}

@fragment
fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
    let i = pos.y / 500; // x frac
    let j = pos.x / 500; // y frac
    let nx = params.ny;
    let ny = params.nx;

    // Clamp to grid bounds (optional, for safety)
    let ii = u32(clamp(i * nx, 0, nx - 1));
    let jj = u32(clamp(j * ny, 0, ny - 1));

    if (ii == 0 || ii == u32(nx) - 1 || jj == 0 || jj == u32(nx) - 1) {
        return vec4f(0, 0, 0, 0);
    }

    let id = idx(ii, jj);

    // Base color: height to color map
    const vmax = 0.5;
    var col = vec3f(
        lerp(hh[id], -vmax, vmax) * 0.3,
        lerp(hh[id], -vmax, vmax) * 0.7,
        lerp(hh[id], -vmax, vmax) * 1.0,
    ) * 0.9;

    // To short-circuit the lighting return here:
    // return vec4f(col, 1.0);

    // Phong lighting
    let l = normalize(vec3f(nx*5, ny*2*0, 0.3*nx));
    let n = normalize(vec3f(
        -gradx(ii, jj) / params.dx,
        -grady(ii, jj) / params.dx,
        0.25,
    ));

    // To render normals return here:
    // return vec4f(n, 1.0);

    let ldotn = dot(l, n);

    let r = 2 * ldotn * n - l;
    let rdotv = r.z; // By construction v = (0, 0, 1).

    let diffuse = lerp(ldotn, 0, 1);
    let specular = lerp(rdotv, 0, 1);

    let col_spec = vec3f(0.8, 0.7, 0.5) * specular * 0.3;
    let col_diff = vec3f(0.3, 0.7, 1.0) * diffuse * 0.3;

    // To render _only_ the Phong lighting return here:
    // return vec4f(col_spec + col_diff, 1);

    // To render the Phong lighting and base colour, but not the sand, return here:
    // return vec4f(col + col_spec + col_diff, 1.0);

    // blend sand

    let depth = 4 * (hh[id] + dd[id]) * (hh[id] + dd[id]);
    let wr = 1 - exp(-1.0 * depth / 0.05);
    let wg = 1 - exp(-1.0 * depth / 0.05);
    let wb = 1 - exp(-0.5 * depth / 0.05);
    // let w = tanh(depth / 0.25);

    let w = vec3f(wr, wg, wb);
 
    col = col * w + vec3f(190./255., 221./255., 0.0) * (1 - w);
    
    // To render just the blended colour return here
    // return vec4f(col, 1);

    // To render the full blend + lighting
    col = col + col_spec + col_diff;

    // DEBUG: this just outputs the normal vector
    // col = vec3f(
    //     lerp(n.x, -1, 1),
    //     lerp(n.y, -1, 1),
    //     lerp(n.z, -1, 1),
    // );

    return vec4f(col, 1.0);
}