"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLUint8Encode = void 0;
const util_1 = require("../../../util");
const glsl_source_1 = require("../glsl-source");
class WebGLUint8Encode {
    runInternal(inferenceHandler, input) {
        const outputShape = input.shape;
        const [width, height] = inferenceHandler.session.layoutStrategy.computeTextureWH(input.shape);
        const outputLayout = {
            width,
            height,
            channels: 4,
            shape: outputShape,
            strides: util_1.ShapeUtil.computeStrides(outputShape),
            unpackedShape: outputShape
        };
        const glsl = glsl_source_1.getGlsl(inferenceHandler.session.backend.glContext.version);
        // TODO: remove this special script. Use graph transformer instead.
        /**
         * https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/encode_float_gpu.ts
         */
        const shaderSource = `
      const float FLOAT_MAX = 1.70141184e38;
      const float FLOAT_MIN = 1.17549435e-38;

      bool isNaN(float val) {
        return (val < 1.0 || 0.0 < val || val == 0.0) ? false : true;
      }

      highp vec4 encodeAsUint8(highp float v) {
        if (isNaN(v)) {
          return vec4(255, 255, 255, 255);
        }

        highp float av = abs(v);

        if(av < FLOAT_MIN) {
          return vec4(0.0, 0.0, 0.0, 0.0);
        } else if(v > FLOAT_MAX) {
          return vec4(0.0, 0.0, 128.0, 127.0) / 255.0;
        } else if(v < -FLOAT_MAX) {
          return vec4(0.0, 0.0,  128.0, 255.0) / 255.0;
        }

        highp vec4 c = vec4(0,0,0,0);

        highp float e = floor(log2(av));
        highp float m = exp2(fract(log2(av))) - 1.0;

        c[2] = floor(128.0 * m);
        m -= c[2] / 128.0;
        c[1] = floor(32768.0 * m);
        m -= c[1] / 32768.0;
        c[0] = floor(8388608.0 * m);

        highp float ebias = e + 127.0;
        c[3] = floor(ebias / 2.0);
        ebias -= c[3] * 2.0;
        c[2] += floor(ebias) * 128.0;

        c[3] += 128.0 * step(0.0, -v);

        return c / 255.0;
      }

      void main() {
        float value = ${glsl.texture2D}(X,TexCoords).r;
        ${glsl.output} = encodeAsUint8(value);
      }`;
        const programInfo = { inputLayouts: [input], outputLayout, samplers: ['X'], shaderSource, hasMain: true };
        const artifact = inferenceHandler.session.programManager.build(programInfo);
        const encoder = inferenceHandler.session.backend.glContext.getEncoder('byte', 4);
        const texture = inferenceHandler.session.backend.glContext.allocateTexture(outputLayout.width, outputLayout.height, encoder);
        const outputTextureData = inferenceHandler.createSharedTextureData(outputLayout, 'uint8', texture);
        const runData = { inputTextureDatas: [input], outputTextureData, uniformData: {} };
        inferenceHandler.session.programManager.run(artifact, runData);
        return runData.outputTextureData;
    }
}
exports.WebGLUint8Encode = WebGLUint8Encode;
//# sourceMappingURL=uint8-encode.js.map