"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLInstanceNormalization = void 0;
const instance_normalization_1 = require("../../../ops/instance-normalization");
const glsl_source_1 = require("../glsl-source");
class WebGLInstanceNormalization extends instance_normalization_1.InstanceNormalization {
    run(inferenceHandler, inputs) {
        if (!this.artifacts) {
            this.artifacts = [];
            const programInfos = this.createProgramInfos(inferenceHandler, inputs);
            programInfos.forEach((pi) => {
                const artifact = inferenceHandler.session.programManager.build(pi);
                this.artifacts.push(artifact);
            });
        }
        const runDatas = this.createRunDatas(inferenceHandler, this.artifacts.map(a => a.programInfo), inputs);
        runDatas.forEach((v, i) => inferenceHandler.session.programManager.run(this.artifacts[i], v));
        return [runDatas[1].outputTextureData.tensor];
    }
    checkInputTypes(inputs) {
        if (!super.checkInputTypes(inputs)) {
            return false;
        }
        if (inputs[0].dims.length !== 4) {
            // currently webgl implementation only support 4-D input.
            return false;
        }
        return true;
    }
    createMeanAndVarianceProgramInfo(inferenceHandler, xLayout) {
        const xDims = xLayout.shape;
        const channel = xDims[1];
        const channelSize = xDims[2] * xDims[3];
        const outputShape = [xDims[0], channel];
        const outputUnpackedShape = [xDims[0], channel * 4];
        const shaderSource = `
    vec4 process(int[2] indices) {
      vec4 v = vec4(0.0);
      int a[4];
      a[0] = indices[0];
      a[1] = indices[1];
      float temp = 0.0;
      for(int a2=0; a2<${xDims[2]}; a2++) {
        a[2] = a2;
        for(int a3=0; a3<${xDims[3]}; a3++) {
          a[3] = a3;
          float x = _X(a);
          temp += x;
        }
      }
      float mean = temp / float(${channelSize});
      temp = 0.0;
      for(int a2=0; a2<${xDims[2]}; a2++) {
        a[2] = a2;
        for(int a3=0; a3<${xDims[3]}; a3++) {
          a[3] = a3;
          float x = _X(a);
          temp += (x - mean) * (x - mean);
        }
      }
      v.r = mean;
      v.g = temp / float(${channelSize});

      return v;
    }`;
        return {
            inputLayouts: [xLayout],
            outputLayout: inferenceHandler.createTextureLayoutFromShape(outputShape, 4, outputUnpackedShape),
            samplers: ['X'],
            shaderSource,
        };
    }
    createComputOutputProgramInfo(inferenceHandler, xLayout, scaleLayout, bLayout, meanAndVarianceLayout) {
        const glsl = glsl_source_1.getGlsl(inferenceHandler.session.backend.glContext.version);
        const shaderSource = `
    vec4 get_MeanAndVariance(int[2] mv) {
      int offset = indicesToOffset_MeanAndVariance(mv);
      vec2 coords = offsetToCoords(offset, ${meanAndVarianceLayout.width}, ${meanAndVarianceLayout.height});
      return ${glsl.texture2D}(MeanAndVariance, coords);
    }

    float process(int[4] indices) {

          int mv[2];
          mv[0] = indices[0];
          mv[1] = indices[1];
          vec4 mean_and_variance = get_MeanAndVariance(mv);
          float mean = mean_and_variance.r;
          float variance = mean_and_variance.g;

          int sb[1];
          sb[0] = indices[1];
          float scale = _Scale(sb);
          float b = _B(sb);

          return scale * (_X(indices) - mean) / sqrt(variance + epsilon) + b;
        }`;
        return {
            inputLayouts: [xLayout, meanAndVarianceLayout, scaleLayout, bLayout],
            outputLayout: inferenceHandler.createTextureLayoutFromShape(xLayout.shape),
            samplers: ['X', 'MeanAndVariance', 'Scale', 'B'],
            variables: [{ name: 'epsilon', type: 'float' }],
            shaderSource,
        };
    }
    createProgramInfos(inferenceHandler, inputs) {
        const xLayout = inferenceHandler.getOrCreateTextureLayout(inputs[0]);
        const scaleLayout = inferenceHandler.getOrCreateTextureLayout(inputs[1]);
        const bLayout = inferenceHandler.getOrCreateTextureLayout(inputs[2]);
        const meanAndVarianceProgramInfo = this.createMeanAndVarianceProgramInfo(inferenceHandler, xLayout);
        const computeOutputProgramInfo = this.createComputOutputProgramInfo(inferenceHandler, xLayout, scaleLayout, bLayout, meanAndVarianceProgramInfo.outputLayout);
        const programInfos = [meanAndVarianceProgramInfo, computeOutputProgramInfo];
        return programInfos;
    }
    createRunDatas(inferenceHandler, programInfos, inputs) {
        const dataType = inputs[0].type;
        const inputTD = inferenceHandler.getOrCreateTextureData(inputs[0], programInfos[0].inputLayouts[0]);
        const scaleTD = inferenceHandler.getOrCreateTextureData(inputs[1], programInfos[1].inputLayouts[2]);
        const bTD = inferenceHandler.getOrCreateTextureData(inputs[2], programInfos[1].inputLayouts[3]);
        const runDatas = [];
        runDatas.push({
            inputTextureDatas: [inputTD],
            outputTextureData: inferenceHandler.createTextureDataFromLayout(programInfos[0].outputLayout, dataType),
            uniformData: {}
        });
        runDatas.push({
            inputTextureDatas: [inputTD, runDatas[0].outputTextureData, scaleTD, bTD],
            outputTextureData: inferenceHandler.createTextureDataFromLayout(programInfos[1].outputLayout, dataType),
            uniformData: { 'epsilon': this.epsilon }
        });
        return runDatas;
    }
}
exports.WebGLInstanceNormalization = WebGLInstanceNormalization;
//# sourceMappingURL=instance-normalization.js.map