// PointLightedCube_perFragment.js (c) 2012 matsuda, 2022 Jonathon Doran
// BivariateFunction.js (c) 2022 John Breaux


// k: number of subdivisions per edge
// I prefer 96
var k = 64
// f(x, y): function to graph
function f(x, y) {
  var u = 80 * x - 40, v = 90 * y - 45;
  var norm = Math.sqrt(u * u + v * v);
  return (1 / 2) * Math.pow(Math.E, -0.04 * norm) * Math.cos(0.15 * norm)
}


const vertex_shader = `
  attribute vec4 a_Position;
  attribute vec4 a_Color;
  attribute vec4 a_Normal;
  uniform mat4 u_MvpMatrix;
  uniform mat4 u_ModelMatrix;    // Model matrix
  uniform mat4 u_NormalMatrix;   // Transformation matrix of the normal
  varying vec4 v_Color;
  varying vec3 v_Normal;
  varying vec3 v_Position;

  void main()
  {
    gl_Position = u_MvpMatrix * a_Position;

    // Calculate the vertex position in the world coordinate
    v_Position = vec3(u_ModelMatrix * a_Position);
    v_Normal = normalize(vec3(u_NormalMatrix * a_Normal));
    v_Color = a_Color;
  }	`;


const fragment_shader = `
  #ifdef GL_ES
	precision mediump float;
  #endif

  uniform vec3 u_LightColor;     // Light color
  uniform vec3 u_LightPosition;  // Position of the light source
  uniform vec3 u_AmbientLight;   // Ambient light color
  varying vec3 v_Normal;
  varying vec3 v_Position;
  varying vec4 v_Color;
  void main()
  {
    // Normalize the normal because it is interpolated and not 1.0 in length any more
    vec3 normal = normalize(v_Normal);

    // Calculate the light direction and make its length 1.
    vec3 lightDirection = normalize(u_LightPosition - v_Position);

    // The dot product of the light direction and the orientation of a surface (the normal)
    float nDotL = max(dot(lightDirection, normal), 0.0);

    // Calculate the final color from diffuse reflection and ambient reflection
    vec3 diffuse = u_LightColor * v_Color.rgb * nDotL;

    vec3 ambient = u_AmbientLight * v_Color.rgb;
    gl_FragColor = vec4(diffuse + ambient, v_Color.a);
  }	`;

function main() {
  // Retrieve <canvas> element
  var canvas = document.getElementById('webgl');

  // Get the rendering context for WebGL *2*
  var gl = canvas.getContext('webgl2');

  if (!gl) {
    console.log('Failed to get the rendering context for WebGL');
    return;
  }

  // Initialize shaders
  if (!initShaders(gl, vertex_shader, fragment_shader)) {
    console.log('Failed to intialize shaders.');
    return;
  }

  //
  var n = initVertexBuffers(gl);
  if (n < 0) {
    console.log('Failed to set the vertex information');
    return;
  }

  // Set the clear color and enable the depth test
  gl.clearColor(0.0, 0.0, 0.0, 1.0);
  gl.enable(gl.DEPTH_TEST);

  // Get the storage locations of uniform variables
  var u_ModelMatrix = gl.getUniformLocation(gl.program, 'u_ModelMatrix');
  var u_MvpMatrix = gl.getUniformLocation(gl.program, 'u_MvpMatrix');
  var u_NormalMatrix = gl.getUniformLocation(gl.program, 'u_NormalMatrix');
  var u_LightColor = gl.getUniformLocation(gl.program, 'u_LightColor');
  var u_LightPosition = gl.getUniformLocation(gl.program, 'u_LightPosition');
  var u_AmbientLight = gl.getUniformLocation(gl.program, 'u_AmbientLight');

  if (!u_ModelMatrix || !u_MvpMatrix || !u_NormalMatrix || !u_LightColor || !u_LightPosition || !u_AmbientLight) {
    console.log('Failed to get the storage location');
    return;
  }

  // Set the light color (white)
  gl.uniform3f(u_LightColor, 1.0, 1.0, 1.0);

  // Set the light direction (in the world coordinate)
  gl.uniform3f(u_LightPosition, 2.3, 4.0, 3.5);

  // Set the ambient light
  gl.uniform3f(u_AmbientLight, 0.2, 0.2, 0.2);

  var modelMatrix = new Matrix4();  // Model matrix
  var mvpMatrix = new Matrix4();    // Model view projection matrix
  var normalMatrix = new Matrix4(); // Transformation matrix for normals

  // Calculate the model matrix
  // Move center of model to center
  modelMatrix.setTranslate(-0.5, 0, 0.5)
  // Rotate the model upright, around X
  modelMatrix.rotate(-90, 1, 0, 0);

  // Calculate the view projection matrix
  mvpMatrix.setPerspective(30, canvas.width / canvas.height, 1, 100);
  var theta = Math.PI / 3
  mvpMatrix.lookAt(Math.tan(theta), Math.sin(theta), Math.cos(theta), // Why fiddle with multiple constants when you can just fiddle with one?
    0.0, 0.0, 0.0,
    0.0, 1.0, 0.0
  );
  mvpMatrix.multiply(modelMatrix);

  // Calculate the matrix to transform the normal based on the model matrix
  normalMatrix.setInverseOf(modelMatrix);
  normalMatrix.transpose();

  // Pass the model matrix to u_ModelMatrix
  gl.uniformMatrix4fv(u_ModelMatrix, false, modelMatrix.elements);

  // Pass the model view projection matrix to u_mvpMatrix
  gl.uniformMatrix4fv(u_MvpMatrix, false, mvpMatrix.elements);

  // Pass the transformation matrix for normals to u_NormalMatrix
  gl.uniformMatrix4fv(u_NormalMatrix, false, normalMatrix.elements);

  // Clear color and depth buffer
  gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);

  // Draw the cube
  gl.drawElements(gl.TRIANGLES, n, gl.UNSIGNED_INT, 0);
}


class BivariateFunction {
  vertices = [];
  vertex_colors = [];
  vertex_normals = [];
  indices = [];

  resolution = 50
  func = (x, y) => { return 0; }

  constructor({ resolution = 50, func } = { resolution: 50, func: this.func }) {
    this.resolution = resolution;
    this.func = func;
    this.construct_mesh();
  }

  generate_quad(width, index) {
    var upper = [index + 1, index, index + width + 1]
    var lower = [index + 1, index + width + 1, index + width + 2]
    return [upper, lower]
  }

  construct_mesh() {
    var vertices = [], vertex_colors = [], vertex_normals = [], triangles = [];
    // Generate a
    for (var i = 0, y = 0; y <= this.resolution; y++) {
      for (var x = 0; x <= this.resolution; x++) {
        // Create a new vertex
        vertices.push([x / this.resolution, y / this.resolution,
        this.func(x / this.resolution, y / this.resolution)]);
        // Color this vertex a very pleasing shade of velvet-red
        vertex_colors.push([0.625, 0.025, 0.075])
        //vertex_colors.push([x / width, y / width, 1]);
        // Give it a normal vector
        vertex_normals.push([0, 0, 0]);

        // If this point is the top-left corner of a quad,
        // then create a new quad
        if (x < this.resolution && y < this.resolution) {
          triangles.push(...this.generate_quad(this.resolution, i));
        }
        i++
      }
    }
    // Sum the normals around each vertex
    for (var triangle of triangles) {
      // Turn verts into vec3's
      var v = [new vec3(...vertices[triangle[0]]), new vec3(...vertices[triangle[1]]), new vec3(...vertices[triangle[2]])];
      // Calculate alpha and beta vectors
      var alpha = v[2].sub(v[0]);
      var beta = v[1].sub(v[0]);
      // Take the cross product & normalize to get the normal
      var normal = alpha.cross(beta).normalize();
      // Add the normal to each vertex
      vertex_normals[triangle[0]] = normal.add(new vec3(...vertex_normals[triangle[0]])).a();
      vertex_normals[triangle[1]] = normal.add(new vec3(...vertex_normals[triangle[1]])).a();
      vertex_normals[triangle[2]] = normal.add(new vec3(...vertex_normals[triangle[2]])).a();
    }
    // Average the normals by normalizing the sums of the normals
    // this assumes the sum of normals around a vertex is not 0
    for (var i = 0; i < vertex_normals.length; i++) {
      // taking full advantage of the garbage collector here
      vertex_normals[i] = new vec3(...vertex_normals[i]).normalize().a();
    }
    // Save the vertices and indices as flat arrays
    this.vertices = vertices.flat();
    this.vertex_colors = vertex_colors.flat();
    this.vertex_normals = vertex_normals.flat();
    this.indices = triangles.flat();
    return
  }
}


function initVertexBuffers(gl) {
  // Create a new BivariateFunction
  var bf = new BivariateFunction({
    resolution: k, func: f
  });

  var vertices = new Float32Array(bf.vertices);
  var colors = new Float32Array(bf.vertex_colors);
  var normals = new Float32Array(bf.vertex_normals);
  var indices = new Uint32Array(bf.indices);

  // Write the vertex property to buffers (coordinates, colors and normals)
  if (!initArrayBuffer(gl, 'a_Position', vertices, 3)) return -1;
  if (!initArrayBuffer(gl, 'a_Color', colors, 3)) return -1;
  if (!initArrayBuffer(gl, 'a_Normal', normals, 3)) return -1;

  // Unbind the buffer object
  gl.bindBuffer(gl.ARRAY_BUFFER, null);

  // Write the indices to the buffer object
  var indexBuffer = gl.createBuffer();
  if (!indexBuffer) {
    console.log('Failed to create the buffer object');
    return false;
  }
  gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, indexBuffer);
  gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, indices, gl.STATIC_DRAW);

  return indices.length;
}

function initArrayBuffer(gl, attribute, data, num) {
  // Create a buffer object
  var buffer = gl.createBuffer();
  if (!buffer) {
    console.log('Failed to create the buffer object');
    return false;
  }

  // Write date into the buffer object
  gl.bindBuffer(gl.ARRAY_BUFFER, buffer);
  gl.bufferData(gl.ARRAY_BUFFER, data, gl.STATIC_DRAW);

  // Assign the buffer object to the attribute variable
  var a_attribute = gl.getAttribLocation(gl.program, attribute);
  if (a_attribute < 0) {
    console.log('Failed to get the storage location of ' + attribute);
    return false;
  }

  gl.vertexAttribPointer(a_attribute, num, gl.FLOAT, false, 0, 0);

  // Enable the assignment of the buffer object to the attribute variable
  gl.enableVertexAttribArray(a_attribute);

  return true;
}