// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import { Tensor } from '../../../tensor';
import { ShapeUtil } from '../../../util';
import { getGlsl } from '../glsl-source';
import { WebGLInferenceHandler } from '../inference-handler';
import { ProgramInfo, ProgramInfoLoader, ProgramMetadata, TextureType } from '../types';

import { getActivationSnippet, InternalActivationAttributes } from './fuse-utils';
import { calculateIm2ColDims } from './im2col';

const createDotProductProgramMetadata = (hasBias: boolean, attributes: InternalActivationAttributes) => ({
  name: 'ConvDotProduct',
  inputNames: hasBias ? ['Im2Col', 'K', 'B'] : ['Im2Col', 'K'],
  inputTypes: hasBias
    ? [TextureType.unpacked, TextureType.packedLastDimension, TextureType.unpacked]
    : [TextureType.unpacked, TextureType.packedLastDimension],
  cacheKey: attributes.activationCacheKey,
});

const createDotProductProgramInfo = (
  inferenceHandler: WebGLInferenceHandler,
  metadata: ProgramMetadata,
  inputs: readonly Tensor[],
  outputShape: number[],
  attributes: InternalActivationAttributes,
): ProgramInfo => {
  const xshape = inputs[0].dims;
  const kshape = inputs[1].dims;
  const adjustedKernelShape = [kshape[0], Math.ceil((xshape[1] * kshape[2] * kshape[3]) / 4)];
  const im2colShape = calculateIm2ColDims(xshape, kshape, outputShape);
  const [kWidth, kHeight] = inferenceHandler.calculateTextureWidthAndHeight(
    adjustedKernelShape,
    TextureType.packedLastDimension,
  );

  const im2colStrides = ShapeUtil.computeStrides(im2colShape);
  const [im2colWidth, im2colHeight] = inferenceHandler.calculateTextureWidthAndHeight(
    im2colShape,
    TextureType.packedLastDimension,
  );
  const rank = outputShape.length;

  const initValue = inputs.length < 3 ? '0.0' : '_B(b)';
  const sharedDim = Math.ceil((xshape[1] * kshape[2] * kshape[3]) / 4);
  const { activationFunction, applyActivation } = getActivationSnippet(attributes);
  const glsl = getGlsl(inferenceHandler.session.backend.glContext.version);
  const shaderSource = `
${activationFunction}
float process(int indices[${rank}]) {
  int b[1];
  b[0] = indices[1];
  int im2col[4];
  im2col[0] = indices[0];
  im2col[1] = indices[2];
  im2col[2] = indices[3];
  int im2colOffset = im2col[0] * ${im2colStrides[0]} + im2col[1] * ${im2colStrides[1]} + im2col[2] * ${
    im2colStrides[2]
  };
  int kernelOffset = indices[1] * ${adjustedKernelShape[1]};
  float value = ${initValue};
  for (int i = 0; i < ${sharedDim}; ++i) {
    vec2 im2colCoords = offsetToCoords(im2colOffset, ${im2colWidth}, ${im2colHeight});
    vec2 kernelCoords = offsetToCoords(kernelOffset, ${kWidth}, ${kHeight});
    value += dot(${glsl.texture2D}(Im2Col, im2colCoords), ${glsl.texture2D}(K, kernelCoords));
    ++im2colOffset;
    ++kernelOffset;
  }
  ${applyActivation}
  return value;
}`;
  return {
    ...metadata,
    output: { dims: outputShape, type: inputs[0].type, textureType: TextureType.unpacked },
    shaderSource,
  };
};

export const createDotProductProgramInfoLoader = (
  inferenceHandler: WebGLInferenceHandler,
  inputs: readonly Tensor[],
  outputShape: number[],
  attributes: InternalActivationAttributes,
): ProgramInfoLoader => {
  const metadata = createDotProductProgramMetadata(inputs.length > 2, attributes);
  return {
    ...metadata,
    get: () => createDotProductProgramInfo(inferenceHandler, metadata, inputs, outputShape, attributes),
  };
};