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

import { Graph } from '../../../graph';
import { OperatorImplementation, OperatorInitialization } from '../../../operators';
import { Tensor } from '../../../tensor';
import { getGlsl } from '../glsl-source';
import { WebGLInferenceHandler } from '../inference-handler';
import { ProgramInfo, ProgramInfoLoader, ProgramMetadata, TextureType } from '../types';

export const instanceNormalization: OperatorImplementation<number> = (
  inferenceHandler: WebGLInferenceHandler,
  inputs: Tensor[],
  epsilon: number,
): Tensor[] => {
  validateInputs(inputs);

  const meanAndVariance = inferenceHandler.run(createMeanAndVarianceProgramInfoLoader(inputs[0]), inputs);
  const output = inferenceHandler.run(
    createComputeOutputProgramInfoLoader(inferenceHandler, inputs[0], epsilon, meanAndVariance.dims),
    [inputs[0], meanAndVariance, inputs[1], inputs[2]],
  );
  return [output];
};

export const parseInstanceNormalizationAttributes: OperatorInitialization<number> = (node: Graph.Node): number =>
  node.attributes.getFloat('epsilon', 1e-5);

const meanAndVarianceProgramMetadata = {
  name: 'InstanceNormalization_MeanAndVariance',
  inputNames: ['X'],
  inputTypes: [TextureType.unpacked],
};

const createMeanAndVarianceProgramInfo = (metadata: ProgramMetadata, input: Tensor): ProgramInfo => {
  const xDims = input.dims.slice();
  const channel = xDims[1];
  const channelSize = xDims[2] * xDims[3];
  const outputShape = [xDims[0], channel];

  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 {
    ...metadata,
    output: { dims: outputShape, type: input.type, textureType: TextureType.packedLastDimension },
    shaderSource,
  };
};

const createMeanAndVarianceProgramInfoLoader = (input: Tensor): ProgramInfoLoader => ({
  ...meanAndVarianceProgramMetadata,
  get: () => createMeanAndVarianceProgramInfo(meanAndVarianceProgramMetadata, input),
});

const computeOutputProgramMetadata = {
  name: 'InstanceNormalization_ComputeOutput',
  inputNames: ['X', 'MeanAndVariance', 'Scale', 'B'],
  inputTypes: [TextureType.unpacked, TextureType.packedLastDimension, TextureType.unpacked, TextureType.unpacked],
};

const createComputeOutputProgramInfo = (
  inferenceHandler: WebGLInferenceHandler,
  metadata: ProgramMetadata,
  input: Tensor,
  epsilon: number,
  meanAndVarianceShape: readonly number[],
): ProgramInfo => {
  const glsl = getGlsl(inferenceHandler.session.backend.glContext.version);
  const [textureWidth, textureHeight] = inferenceHandler.calculateTextureWidthAndHeight(
    meanAndVarianceShape,
    TextureType.packedLastDimension,
  );
  const [meanAndVarianceWidth, meanAndVarianceHeight] = [textureWidth / 4, textureHeight];
  const shaderSource = `
      vec4 get_MeanAndVariance(int[2] mv) {
        int offset = indicesToOffset_MeanAndVariance(mv);
        vec2 coords = offsetToCoords(offset, ${meanAndVarianceWidth}, ${meanAndVarianceHeight});
        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 {
    ...metadata,
    output: { dims: input.dims, type: input.type, textureType: TextureType.unpacked },
    variables: [{ name: 'epsilon', type: 'float', data: epsilon }],
    shaderSource,
  };
};

const createComputeOutputProgramInfoLoader = (
  inferenceHandler: WebGLInferenceHandler,
  input: Tensor,
  epsilon: number,
  meanAndVarianceShape: readonly number[],
): ProgramInfoLoader => {
  const metadata = { ...computeOutputProgramMetadata, cacheHint: `${epsilon}` };
  return {
    ...metadata,
    get: () => createComputeOutputProgramInfo(inferenceHandler, metadata, input, epsilon, meanAndVarianceShape),
  };
};

const validateInputs = (inputs: Tensor[]): void => {
  if (!inputs || inputs.length !== 3) {
    throw new Error('InstanceNormalization requires 3 inputs.');
  }

  const X = inputs[0];
  const scale = inputs[1];
  const B = inputs[2];

  // input should at least have three dimensions - N,C,dim1,...,dimn
  // other inputs can have only one dimensions
  if (X.dims.length < 3 || scale.dims.length !== 1 || B.dims.length !== 1) {
    throw new Error('Invalid input shape.');
  }
  if (scale.dims[0] !== X.dims[1] || B.dims[0] !== X.dims[1]) {
    throw new Error('Input shapes are mismatched.');
  }
  if (
    (X.type !== 'float32' && X.type !== 'float64') ||
    (scale.type !== 'float32' && scale.type !== 'float64') ||
    (B.type !== 'float32' && B.type !== 'float64')
  ) {
    throw new Error('Invalid input type.');
  }
  if (inputs[0].dims.length !== 4) {
    throw new Error('Only support 4-D input shape.');
  }
};