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

import { Tensor } from '../../../tensor';
import { WebGLInferenceHandler } from '../inference-handler';

import { calculateOutputShape, ConvAttributes } from './conv';
import { createPackedIm2ColProgramInfoLoader } from './im2col-pack';
import { createPackedMatmulProgramInfoLoader } from './matmul-pack';

export const conv2DPackedPointwise = (
  inferenceHandler: WebGLInferenceHandler,
  inputs: readonly Tensor[],
  attributes: ConvAttributes,
): Tensor => {
  const xshape = inputs[0].dims;
  const kshape = inputs[1].dims;
  const outputShape = calculateOutputShape(xshape, kshape, attributes.dilations, attributes.pads, attributes.strides);
  const reshapedX = inferenceHandler.reshapePacked(inputs[0], [xshape[1], xshape[2] * xshape[3]]);
  const reshapedK = inferenceHandler.reshapePacked(inputs[1], [kshape[0], kshape[1]]);

  const matmulInputs = inputs.length > 2 ? [reshapedK, reshapedX, inputs[2]] : [reshapedK, reshapedX];
  const matmulOutput = inferenceHandler.run(
    createPackedMatmulProgramInfoLoader(inferenceHandler, matmulInputs, attributes),
    matmulInputs,
  );
  return inferenceHandler.reshapePacked(matmulOutput, outputShape);
};

export const conv2DPacked = (
  inferenceHandler: WebGLInferenceHandler,
  inputs: readonly Tensor[],
  attributes: ConvAttributes,
): Tensor => {
  const xshape = inputs[0].dims;
  const kshape = inputs[1].dims;
  const outputShape = calculateOutputShape(xshape, kshape, attributes.dilations, attributes.pads, attributes.strides);

  // run im2col
  const im2colOutput = inferenceHandler.run(
    createPackedIm2ColProgramInfoLoader(inferenceHandler, inputs[0], inputs[1], outputShape, attributes),
    [inputs[0]],
  );

  // reshape kernel
  const kernelReshaped = inferenceHandler.reshapePacked(inputs[1], [kshape[0], kshape[1] * kshape[2] * kshape[3]]);

  // run matmul
  const matmulInputs = inputs.length === 3 ? [kernelReshaped, im2colOutput, inputs[2]] : [kernelReshaped, im2colOutput];
  const matmulOutput = inferenceHandler.run(
    createPackedMatmulProgramInfoLoader(inferenceHandler, matmulInputs, attributes),
    matmulInputs,
  );

  // reshape output
  const outputReshaped = inferenceHandler.reshapePacked(matmulOutput, outputShape);
  return outputReshaped;
};