
// Define variables for WebGL resources
var gl;
var program;
var positionLocation;
var texcoordLocation;
var positionBuffer;
var texcoordBuffer;
var texture;
var resolutionLocation;
var textureLocation; // Define texture location

function difftime(unix_timestamp){
    let client_unix_time = new Date().getTime()/1000.0
    var diff = client_unix_time - unix_timestamp
          return diff.toFixed(3)
}

// Function to set rectangle geometry in buffer
function setRectangle(gl, x, y, width, height) {
    var x1 = x;
    var x2 = x + width;
    var y1 = y;
    var y2 = y + height;
    gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([x1, y1, x2, y1, x1, y2, x1, y2, x2, y1, x2, y2]), gl.STATIC_DRAW);
}

// Function to setup WebGL resources
function setupWebGL(canvasId) {
    // Get the WebGL context
    var canvas = document.getElementById(canvasId);
    gl = canvas.getContext("webgl");
    if (!gl) {
        return false; // WebGL not supported
    }

    // Setup GLSL program
    program = webglUtils.createProgramFromScripts(gl, [
        "vertex-shader-2d",
        "fragment-shader-2d"
    ]);

    // Lookup attribute and uniform locations
    positionLocation = gl.getAttribLocation(program, "a_position");
    texcoordLocation = gl.getAttribLocation(program, "a_texCoord");
    resolutionLocation = gl.getUniformLocation(program, "u_resolution");
    textureLocation = gl.getUniformLocation(program, "u_texture"); // Get texture location

    // Create position buffer
    positionBuffer = gl.createBuffer();
    // Create texcoord buffer
    texcoordBuffer = gl.createBuffer();

    return true;
}

// Precompute texture data offline
var precomputedTextureData = [];

function typeArray(data, dtype) {
    var typedArray = [];
    var maxDataValue = 65535;
    if (dtype === 'float32') {
        typedArray = new Float32Array(data);
    } else if (dtype === 'int32') {
        typedArray = new Int32Array(data);
    } else if (dtype === 'int16') {
        typedArray = new Int16Array(data);
        maxDataValue = 32767;
    } else if (dtype === 'int8' || dtype === 'uint8') {
        typedArray = new Uint8Array(data);
        maxDataValue = 255;
    } else if (dtype === 'uint16') {
        typedArray = new Uint16Array(data);
    } else {
        console.error('Unsupported data type:', dtype);
        return;
    }
    return typedArray
}


function precomputeTextureData(data, width, height, dtype, is_mask=false) {
    var typedArray = typeArray(data, dtype);
    var maxDataValue = 65535
    if (is_mask) {
        for (var i = 0; i < typedArray.length; i++) {            
            var value = typedArray[i];
            precomputedTextureData.push(0,255,255, value * 0.8);
        }
    } else {
        for (var i = 0; i < typedArray.length; i++) {
            var value = typedArray[i];

            var index = Math.round((value / maxDataValue) * (viridis.length - 1));

            var rgb = viridis[index];
            precomputedTextureData.push(rgb[0], rgb[1], rgb[2], 255);
        }
    }

}

// Upload precomputed texture data to GPU
function uploadTextureDataToGPU(width, height) {
    // Create a texture
    texture = gl.createTexture();
    gl.bindTexture(gl.TEXTURE_2D, texture);
    // Set texture parameters
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);

    var type = gl.UNSIGNED_BYTE;
    var level = 0;
    var internalFormat = gl.RGBA;
    var border = 0;
    var format = gl.RGBA;

    var textureData = new Uint8Array(precomputedTextureData);

    gl.texImage2D(gl.TEXTURE_2D, level, internalFormat, width, height, border, format, type, textureData);
}

// Function to upload matrix data to texture
function uploadMatrixData(data, width, height, dtype, is_mask) {
    precomputedTextureData = [];
    precomputeTextureData(data, width, height, dtype, is_mask);
    uploadTextureDataToGPU(width, height);
}

// Function to draw matrix with WebGL, scaling each value by a custom factor and starting coordinates
function drawMatrixWithWebGL(data, canvasId, scaleFactor=1, x0=0, y0=0, is_mask=false) {
    // Check if WebGL resources need to be setup
    if (!gl) {
        if (!setupWebGL(canvasId)) {
            console.error("WebGL not supported");
            return;
        }
    }

    // Clear the canvas
    gl.clearColor(0, 0, 0, 0);
    gl.clear(gl.COLOR_BUFFER_BIT);

    // Use the shader program
    gl.useProgram(program);

    // Setup vertex position buffer with adjusted starting coordinates
    gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer);
    setRectangle(gl, x0, y0, data.shape[0] * scaleFactor + x0, data.shape[1] * scaleFactor + y0);
    gl.enableVertexAttribArray(positionLocation);
    gl.vertexAttribPointer(positionLocation, 2, gl.FLOAT, false, 0, 0);

    // Setup texture coordinate buffer with scaled values
    gl.bindBuffer(gl.ARRAY_BUFFER, texcoordBuffer);
    gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([
        0, 0,
        1, 0,
        0, 1,
        0, 1,
        1, 0,
        1, 1
    ]), gl.STATIC_DRAW);
    gl.enableVertexAttribArray(texcoordLocation);
    gl.vertexAttribPointer(texcoordLocation, 2, gl.FLOAT, false, 0, 0);

    // Set the resolution uniform
    gl.uniform2f(resolutionLocation, gl.canvas.width, gl.canvas.height);

    // Upload matrix data to texture
    uploadMatrixData(data.data, data.shape[0], data.shape[1], data.dtype, is_mask);

    // Draw the rectangle with scaled dimensions
    gl.drawArrays(gl.TRIANGLES, 0, 6);
}
