/**
 * WaveSurface — BufferGeometry height-map mesh creation and per-frame vertex updates.
 * ColorGradient — Maps amplitude to diverging blue-white-red colour scheme.
 * Loaded as a plain script; classes are defined on the global scope.
 */

class ColorGradient {
  /**
   * Map an amplitude value to an RGB colour.
   * Negative → blue, zero → white/light gray, positive → red.
   * @param {number} amplitude
   * @param {number} maxAmp
   * @returns {{ r: number, g: number, b: number }}
   */
  getColor(amplitude, maxAmp) {
    if (maxAmp === 0) {
      return { r: 0.15, g: 0.15, b: 0.2 };
    }

    // Normalize to [-1, 1]
    var t = amplitude / maxAmp;
    if (t < -1) t = -1;
    if (t > 1) t = 1;

    // 5-stop gradient for vivid amplitude visualization:
    //   t = -1.0 → deep blue    (0.05, 0.05, 0.6)
    //   t = -0.5 → cyan         (0.0,  0.5,  0.8)
    //   t =  0.0 → dark gray    (0.12, 0.12, 0.15)
    //   t = +0.5 → orange       (0.9,  0.5,  0.0)
    //   t = +1.0 → bright red   (1.0,  0.15, 0.1)
    var r, g, b;
    if (t < -0.5) {
      var s = (t + 1.0) / 0.5; // 0 at t=-1, 1 at t=-0.5
      r = 0.05 + s * (0.0 - 0.05);
      g = 0.05 + s * (0.5 - 0.05);
      b = 0.6  + s * (0.8 - 0.6);
    } else if (t < 0) {
      var s = (t + 0.5) / 0.5; // 0 at t=-0.5, 1 at t=0
      r = 0.0  + s * (0.12 - 0.0);
      g = 0.5  + s * (0.12 - 0.5);
      b = 0.8  + s * (0.15 - 0.8);
    } else if (t < 0.5) {
      var s = t / 0.5; // 0 at t=0, 1 at t=0.5
      r = 0.12 + s * (0.9 - 0.12);
      g = 0.12 + s * (0.5 - 0.12);
      b = 0.15 + s * (0.0 - 0.15);
    } else {
      var s = (t - 0.5) / 0.5; // 0 at t=0.5, 1 at t=1
      r = 0.9  + s * (1.0 - 0.9);
      g = 0.5  + s * (0.15 - 0.5);
      b = 0.0  + s * (0.1 - 0.0);
    }

    return { r: r, g: g, b: b };
  }
}

class WaveSurface {
  constructor() {
    this.mesh = null;
    this.geometry = null;
    this.nx = 0;
    this.ny = 0;
  }

  /**
   * Create the PlaneGeometry mesh with nx × ny grid points.
   * PlaneGeometry is created in XY plane then rotated so Y is the height axis (XZ plane).
   * @param {number} nx - grid width in cells
   * @param {number} ny - grid height in cells
   * @param {number} dx - spatial step size
   * @returns {THREE.Mesh}
   */
  init(nx, ny, dx) {
    this.nx = nx;
    this.ny = ny;

    var width = (nx - 1) * dx;
    var height = (ny - 1) * dx;

    // Create PlaneGeometry with (nx-1) x (ny-1) segments → nx*ny vertices
    this.geometry = new THREE.PlaneGeometry(width, height, nx - 1, ny - 1);

    // Rotate from XY plane to XZ plane so Y becomes the height axis
    this.geometry.rotateX(-Math.PI / 2);

    // Add vertex color attribute (Float32Array, 3 components per vertex)
    var vertexCount = this.geometry.attributes.position.count;
    var colors = new Float32Array(vertexCount * 3);
    // Initialize all vertices to white
    for (var i = 0; i < vertexCount * 3; i++) {
      colors[i] = 1.0;
    }
    this.geometry.setAttribute('color', new THREE.BufferAttribute(colors, 3));

    // MeshPhongMaterial with vertex colors and smooth shading
    var material = new THREE.MeshPhongMaterial({
      vertexColors: true,
      flatShading: false,
      side: THREE.DoubleSide
    });

    this.mesh = new THREE.Mesh(this.geometry, material);
    return this.mesh;
  }

  /**
   * Set each vertex y-position to amplitude[index] * heightScale.
   * After updating, recompute normals for correct lighting.
   * @param {Float32Array} amplitude
   * @param {number} heightScale
   */
  updateHeights(amplitude, heightScale) {
    if (!this.geometry) return;

    var position = this.geometry.attributes.position;
    var count = position.count;
    var len = Math.min(count, amplitude.length);

    for (var i = 0; i < len; i++) {
      // After rotateX(-PI/2), the original Y axis maps to the new Y axis (height).
      // position.setY sets the height component.
      position.setY(i, amplitude[i] * heightScale);
    }

    position.needsUpdate = true;
    this.geometry.computeVertexNormals();
  }

  /**
   * Set vertex colours using the diverging colour gradient.
   * @param {Float32Array} amplitude
   * @param {ColorGradient} colorGradient
   */
  updateColors(amplitude, colorGradient) {
    if (!this.geometry) return;

    var colorAttr = this.geometry.attributes.color;
    var count = colorAttr.count;

    // Find max absolute amplitude for normalization
    var maxAmp = 0;
    var len = Math.min(count, amplitude.length);
    for (var i = 0; i < len; i++) {
      var absVal = Math.abs(amplitude[i]);
      if (absVal > maxAmp) maxAmp = absVal;
    }
    // Avoid division by zero
    if (maxAmp === 0) maxAmp = 1;

    for (var i = 0; i < len; i++) {
      var color = colorGradient.getColor(amplitude[i], maxAmp);
      colorAttr.setXYZ(i, color.r, color.g, color.b);
    }

    colorAttr.needsUpdate = true;
  }

  /**
   * Set vertex colours for orthogonal polarization mode.
   * Slit 1 wave → yellow, Slit 2 wave → green, overlap → blended.
   * @param {Float32Array} amp1 - amplitude from slit 1
   * @param {Float32Array} amp2 - amplitude from slit 2
   */
  updateColorsOrthogonal(amp1, amp2) {
    if (!this.geometry) return;

    var colorAttr = this.geometry.attributes.color;
    var count = colorAttr.count;
    var len = Math.min(count, amp1.length, amp2.length);

    // Find max absolute amplitude across both for normalization
    var maxAmp = 0;
    for (var i = 0; i < len; i++) {
      var a = Math.abs(amp1[i]);
      var b = Math.abs(amp2[i]);
      if (a > maxAmp) maxAmp = a;
      if (b > maxAmp) maxAmp = b;
    }
    if (maxAmp === 0) maxAmp = 1;

    for (var i = 0; i < len; i++) {
      var v1 = amp1[i] / maxAmp; // [-1, 1]
      var v2 = amp2[i] / maxAmp;
      if (v1 < -1) v1 = -1; if (v1 > 1) v1 = 1;
      if (v2 < -1) v2 = -1; if (v2 > 1) v2 = 1;

      // Slit 1 (warm): negative=purple, zero=dark, positive=yellow/orange
      var r1 = 0.12 + Math.max(0, v1) * 0.88;
      var g1 = 0.12 + Math.max(0, v1) * 0.68;
      var b1 = 0.15 + Math.max(0, -v1) * 0.65;

      // Slit 2 (cool): negative=dark blue, zero=dark, positive=cyan/green
      var r2 = 0.12 + Math.max(0, -v2) * 0.1;
      var g2 = 0.12 + Math.max(0, v2) * 0.78;
      var b2 = 0.15 + Math.max(0, v2) * 0.65 + Math.max(0, -v2) * 0.55;

      // Blend: additive mix
      var r = Math.min(1, r1 * 0.5 + r2 * 0.5 + Math.abs(v1) * 0.2);
      var g = Math.min(1, g1 * 0.5 + g2 * 0.5);
      var b = Math.min(1, b1 * 0.5 + b2 * 0.5 + Math.abs(v2) * 0.15);

      colorAttr.setXYZ(i, r, g, b);
    }

    colorAttr.needsUpdate = true;
  }
}

window.ColorGradient = ColorGradient;
window.WaveSurface = WaveSurface;