"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLUnpackedConcat = exports.WebGLConcat = void 0;
const concat_1 = require("../../../ops/concat");
const concat_packed_1 = require("./concat_packed");
// We provide a wrapper class so that the kernel can switch between packed and unpacked depending on the inputs on the
// fly.
class WebGLConcat extends concat_1.Concat {
    constructor() {
        super();
        this.unpackedImpl = new WebGLUnpackedConcat();
        this.packedImpl = new concat_packed_1.WebGLPackedConcat();
    }
    // No need to call super since this class only serves as a wrapper.
    initialize(attributes) {
        this.unpackedImpl.initialize(attributes);
        this.packedImpl.initialize(attributes);
    }
    run(inferenceHandler, inputs) {
        if (inferenceHandler.session.pack && inputs[0].dims.length > 1) {
            return inferenceHandler.run(this.packedImpl, inputs);
        }
        else {
            return inferenceHandler.run(this.unpackedImpl, inputs);
        }
    }
    createProgramInfo(handler, inputs) {
        if (handler.session.pack && inputs[0].dims.length > 1) {
            return this.packedImpl.createProgramInfo(handler, inputs);
        }
        else {
            return this.unpackedImpl.createProgramInfo(handler, inputs);
        }
    }
    createRunData(handler, programInfo, inputs) {
        if (handler.session.pack && inputs[0].dims.length > 1) {
            return this.packedImpl.createRunData(handler, programInfo, inputs);
        }
        else {
            return this.unpackedImpl.createRunData(handler, programInfo, inputs);
        }
    }
}
exports.WebGLConcat = WebGLConcat;
class WebGLUnpackedConcat extends concat_1.Concat {
    run(inferenceHandler, inputs) {
        return inferenceHandler.run(this, inputs);
    }
    createProgramInfo(handler, inputs) {
        const inputShape = inputs[0].dims.slice();
        if (this.axis >= inputShape.length || this.axis < (-1 * inputShape.length)) {
            throw new Error('axis specified for concat doesn\'t match input dimensionality');
        }
        if (this.axis < 0) {
            this.axis = inputShape.length + this.axis;
        }
        // ensure all of the non-concatenated axes match each other
        // calculate the shape of the output tensor while we do that
        const outputShape = inputShape.slice(0);
        for (let i = 1; i < inputs.length; i++) {
            const dataNShape = inputs[i].dims.slice();
            for (let axisIndex = 0; axisIndex < inputShape.length; axisIndex++) {
                // add to the placeholder for computing output shape
                if (axisIndex === this.axis) {
                    outputShape[this.axis] += dataNShape[axisIndex];
                }
                // ensure all non-cancatenated axes match each other
                else if (inputShape[axisIndex] !== dataNShape[axisIndex]) {
                    throw new Error('non concat dimensions must match');
                }
            }
        }
        const rank = outputShape.length;
        let getTextureIndexWhereDataResidesMethod = '';
        // in most cases linear search is sufficient, as in most scenarios, only 2 tensors are concatenated
        if (inputs.length < 5) {
            getTextureIndexWhereDataResidesMethod = this.getTextureIndexWhereDataResidesLinearSearch(inputs.length);
        }
        else {
            getTextureIndexWhereDataResidesMethod = this.getTextureIndexWhereDataResidesBinarySearch(inputs.length);
        }
        const fetchDataFromCorrectTextureMethod = this.fetchDataFromCorrectTextureMethod(inputs.length, rank);
        const getValueFromArrayIndexMethod = this.getValueFromArrayIndexMethod(inputs.length);
        const samplers = inputs.map((v, i) => `X${i}`);
        const shaderSource = `
      ${fetchDataFromCorrectTextureMethod}
      ${getValueFromArrayIndexMethod}
      ${getTextureIndexWhereDataResidesMethod}
      float process(int indices[${rank}]) {
        int textureIndex = getTextureWhereDataResides (indices[${this.axis}]);

        if(textureIndex != 0) {
          indices[${this.axis}] = indices[${this.axis}] - int(getValueFromArrayIndex(sizeInConcatAxis, textureIndex-int(1)));
        }

        return fetchDataFromCorrectTexture(textureIndex, indices);
      }`;
        return {
            inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t)),
            outputLayout: handler.createTextureLayoutFromShape(outputShape),
            samplers,
            variables: [{ name: 'sizeInConcatAxis', type: 'int', arrayLength: inputs.length }],
            shaderSource,
        };
    }
    createRunData(handler, programInfo, inputs) {
        const inputTDs = inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i]));
        const sizeInConcatAxis = new Array(programInfo.inputLayouts.length);
        let previousSum = 0;
        for (let i = 0; i < programInfo.inputLayouts.length; ++i) {
            previousSum += programInfo.inputLayouts[i].shape[this.axis];
            sizeInConcatAxis[i] = previousSum;
        }
        const uniformData = { sizeInConcatAxis };
        return {
            inputTextureDatas: inputTDs,
            outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
            uniformData
        };
    }
    getTextureIndexWhereDataResidesLinearSearch(numberOfTensors) {
        return `int getTextureWhereDataResides(int index) {
      for(int i=0; i<${numberOfTensors}; i++) {
          if(index < int(sizeInConcatAxis[i])){
              return i;
          }
        }
      }`;
    }
    // TODO: Implement BinarySearch in GLSL
    getTextureIndexWhereDataResidesBinarySearch(numberOfTensors) {
        return this.getTextureIndexWhereDataResidesLinearSearch(numberOfTensors);
    }
    fetchDataFromCorrectTextureMethod(numberOfTensors, tensorRank) {
        const codeLines = [`float fetchDataFromCorrectTexture(int textureIndex, int indices[${tensorRank}]) {`];
        for (let i = 0; i < numberOfTensors; ++i) {
            if (i === 0) {
                codeLines.push('\t' +
                    `if (textureIndex == ${i}) { return _X${i}(indices); }`);
            }
            else if (i === numberOfTensors - 1) {
                codeLines.push('\t' +
                    `else { return _X${i}(indices); }`);
            }
            else {
                codeLines.push('\t' +
                    `else if (textureIndex == ${i}) { return _X${i}(indices); }`);
            }
        }
        codeLines.push('\t' +
            '}');
        return codeLines.join('\n');
    }
    getValueFromArrayIndexMethod(arrayRank) {
        const codeLines = [`int getValueFromArrayIndex(int arr[${arrayRank}], int index) {`];
        for (let i = 0; i < arrayRank; ++i) {
            if (i === 0) {
                codeLines.push('\t' +
                    `if (index == ${i}) { return arr[${i}]; }`);
            }
            else if (i === arrayRank - 1) {
                codeLines.push('\t' +
                    `else { return arr[${i}]; }`);
            }
            else {
                codeLines.push('\t' +
                    `else if (index == ${i}) { return arr[${i}]; }`);
            }
        }
        codeLines.push('\t' +
            '}');
        return codeLines.join('\n');
    }
}
exports.WebGLUnpackedConcat = WebGLUnpackedConcat;
//# sourceMappingURL=concat.js.map