struct RU {
  mvp: mat4x4f,
  nx: u32,
  ny: u32,
  hx: f32,
  hy: f32,
  wScale: f32,     // vertical exaggeration in the 3D view
  colorRange: f32, // |w| mapped to full color [m]
  aspect: f32,     // canvas width / height
  mode: u32,       // 0: 2D, 1: 3D
  holeX: f32,
  holeY: f32,
  pushX: f32,
  pushY: f32,
  markerR: f32,
  lightX: f32,
  lightY: f32,
  lightZ: f32,
};

@group(0) @binding(0) var<uniform> R: RU;
@group(0) @binding(1) var<storage, read> w: array<f32>;

fn colormap(t: f32) -> vec3f {
  let s = clamp(t, -1.0, 1.0);
  let mid = vec3f(0.93, 0.93, 0.93);
  let blue = vec3f(0.10, 0.35, 0.85);
  let red = vec3f(0.85, 0.20, 0.12);
  if (s < 0.0) { return mix(mid, blue, -s); }
  return mix(mid, red, s);
}

fn wAt(i: i32, j: i32) -> f32 {
  let ii = clamp(i, 0, i32(R.nx) - 1);
  let jj = clamp(j, 0, i32(R.ny) - 1);
  return w[u32(jj) * R.nx + u32(ii)];
}

// ---------- 2D view: full-screen triangle, bilinear lookup in the fragment shader ----------

struct V2 { @builtin(position) pos: vec4f, @location(0) uv: vec2f };

@vertex
fn vs2d(@builtin(vertex_index) vi: u32) -> V2 {
  var p = array<vec2f, 3>(vec2f(-1.0, -1.0), vec2f(3.0, -1.0), vec2f(-1.0, 3.0));
  var o: V2;
  o.pos = vec4f(p[vi], 0.0, 1.0);
  o.uv = p[vi];
  return o;
}

@fragment
fn fs2d(in: V2) -> @location(0) vec4f {
  let Lx = f32(R.nx - 1u) * R.hx;
  let Ly = f32(R.ny - 1u) * R.hy;
  // Fit the plate into the canvas with a margin, preserving its aspect ratio.
  let margin = 0.92;
  var sx = margin; var sy = margin;
  let plateAspect = Lx / Ly;
  if (plateAspect > R.aspect) { sy = margin * R.aspect / plateAspect; } else { sx = margin * plateAspect / R.aspect; }
  let px = in.uv.x / sx * 0.5 * Lx;
  let py = in.uv.y / sy * 0.5 * Ly;
  let bg = vec4f(0.11, 0.12, 0.14, 1.0);
  if (abs(px) > 0.5 * Lx || abs(py) > 0.5 * Ly) { return bg; }
  let fi = (px + 0.5 * Lx) / R.hx;
  let fj = (py + 0.5 * Ly) / R.hy;
  let i0 = i32(floor(fi)); let j0 = i32(floor(fj));
  let a = fi - f32(i0); let b = fj - f32(j0);
  let val = mix(mix(wAt(i0, j0), wAt(i0 + 1, j0), a), mix(wAt(i0, j0 + 1), wAt(i0 + 1, j0 + 1), a), b);
  var col = colormap(val / R.colorRange);
  // Markers: hole (ring) and pushed patch (dashed-looking ring).
  let dh = distance(vec2f(px, py), vec2f(R.holeX, R.holeY));
  if (abs(dh - R.markerR) < 0.35 * R.markerR) { col = vec3f(0.1, 0.1, 0.1); }
  if (dh < 0.4 * R.markerR) { col = vec3f(0.11, 0.12, 0.14); }
  let dp = distance(vec2f(px, py), vec2f(R.pushX, R.pushY));
  if (abs(dp - R.markerR * 1.6) < 0.25 * R.markerR) { col = vec3f(0.95, 0.75, 0.1); }
  return vec4f(col, 1.0);
}

// ---------- 3D view: height field ----------

struct V3 { @builtin(position) pos: vec4f, @location(0) val: f32, @location(1) normal: vec3f };

@vertex
fn vs3d(@builtin(vertex_index) vi: u32) -> V3 {
  let i = i32(vi % R.nx); let j = i32(vi / R.nx);
  let x = (f32(i) - 0.5 * f32(R.nx - 1u)) * R.hx;
  let y = (f32(j) - 0.5 * f32(R.ny - 1u)) * R.hy;
  let val = wAt(i, j);
  let z = val * R.wScale;
  let dzdx = (wAt(i + 1, j) - wAt(i - 1, j)) * R.wScale / (2.0 * R.hx);
  let dzdy = (wAt(i, j + 1) - wAt(i, j - 1)) * R.wScale / (2.0 * R.hy);
  var o: V3;
  o.pos = R.mvp * vec4f(x, y, z, 1.0);
  o.val = val;
  o.normal = normalize(vec3f(-dzdx, -dzdy, 1.0));
  return o;
}

@fragment
fn fs3d(in: V3) -> @location(0) vec4f {
  let l = normalize(vec3f(R.lightX, R.lightY, R.lightZ));
  let n = normalize(in.normal);
  let diff = abs(dot(n, l));
  let base = colormap(in.val / R.colorRange);
  let col = base * (0.35 + 0.65 * diff) + vec3f(0.15) * pow(diff, 24.0);
  return vec4f(col, 1.0);
}