"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.getPadFunction = exports.WebGLPad = void 0;
const pad_1 = require("../../../ops/pad");
const util_1 = require("../../../util");
const glsl_source_1 = require("../glsl-source");
class WebGLPad extends pad_1.Pad {
    run(inferenceHandler, inputs) {
        return inferenceHandler.run(this, inputs);
    }
    createProgramInfo(inferenceHandler, inputs) {
        const outputShape = util_1.ShapeUtil.padShape(inputs[0].dims.slice(), this.pads);
        const rank = outputShape.length;
        const alayout = inferenceHandler.getOrCreateTextureLayout(inputs[0]);
        const padFunction = getPadFunction(glsl_source_1.getGlsl(inferenceHandler.session.backend.glContext.version), 'A', alayout, this.mode, this.pads, this.value);
        const shaderSource = `
      ${padFunction}
      float process(int[${rank}] indices) {
          return padA(indices);
      }`;
        return {
            inputLayouts: [alayout],
            outputLayout: inferenceHandler.createTextureLayoutFromShape(outputShape),
            samplers: ['A'],
            shaderSource,
        };
    }
    createRunData(inferenceHandler, programInfo, inputs) {
        const inputTDs = [inferenceHandler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0])];
        return {
            inputTextureDatas: inputTDs,
            outputTextureData: inferenceHandler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
            uniformData: {}
        };
    }
}
exports.WebGLPad = WebGLPad;
function getPadFunction(glsl, name, inputLayout, mode, pads, value) {
    switch (mode) {
        case 'constant':
            return getPadConstant(glsl, name, inputLayout.shape, inputLayout.strides, inputLayout.width, inputLayout.height, pads, value);
        case 'reflect':
            return getPadReflect(glsl, name, inputLayout.shape, inputLayout.strides, inputLayout.width, inputLayout.height, pads);
        case 'edge':
            return getPadEdge(glsl, name, inputLayout.shape, inputLayout.strides, inputLayout.width, inputLayout.height, pads);
        default:
            throw new Error('Invalid mode');
    }
}
exports.getPadFunction = getPadFunction;
function getPadConstant(glsl, name, shape, strides, width, height, pads, value) {
    const rank = shape.length;
    let block = '';
    for (let i = rank - 1; i >= 0; --i) {
        block += `
          k = m[${i}] - ${pads[i]};
          if (k < 0)  return constant;
          if (k >= ${shape[i]}) return constant;
          offset += k * ${strides[i]};
          `;
    }
    return `
        float pad${name}(int m[${rank}]) {
          const float constant = float(${value});
          int offset = 0;
          int k = 0;
          ${block}
          vec2 coords = offsetToCoords(offset, ${width}, ${height});
          float value = getColorAsFloat(${glsl.texture2D}(${name}, coords));
          return value;
        }
        `;
}
function getPadReflect(glsl, name, shape, strides, width, height, pads) {
    const rank = shape.length;
    let block = '';
    for (let i = rank - 1; i >= 0; --i) {
        block += `
        k = m[${i}] - ${pads[i]};
        if (k < 0) { k = -k; }
        {
          const int _2n_1 = ${2 * (shape[i] - 1)};
          k = int( mod( float(k), float(_2n_1) ) ) ;
          if(k >= ${shape[i]}) { k = _2n_1 - k; }
        }
        offset += k * ${strides[i]};
        `;
    }
    return `
      float pad${name}(int m[${rank}]) {
        int offset = 0;
        int k = 0;
        ${block}
        vec2 coords = offsetToCoords(offset, ${width}, ${height});
        float value = getColorAsFloat(${glsl.texture2D}(${name}, coords));
        return value;
      }
      `;
}
function getPadEdge(glsl, name, shape, strides, width, height, pads) {
    const rank = shape.length;
    let block = '';
    for (let i = rank - 1; i >= 0; --i) {
        block += `
      k = m[${i}] - ${pads[i]};
      if (k < 0)  k = 0;
      if (k >= ${shape[i]}) k = ${shape[i] - 1};
      offset += k * ${strides[i]};
      `;
    }
    return `
    float pad${name}(int m[${rank}]) {
      int offset = 0;
      int k = 0;
      ${block}
      vec2 coords = offsetToCoords(offset, ${width}, ${height});
      float value = getColorAsFloat(${glsl.texture2D}(${name}, coords));
      return value;
    }
    `;
}
//# sourceMappingURL=pad.js.map