// 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, TextureType} from '../types';
import {getCoordsDataType} from '../utils';

import {unpackFromChannel} from './packing-utils';
import {parseUpsampleAttributes, scalesValidation, UpsampleAttributes, validateInputs} from './upsample';

const resizeProgramMetadata = {
  name: 'Resize',
  inputNames: ['A'],
  inputTypes: [TextureType.packed]
};

export const resize: OperatorImplementation<UpsampleAttributes> =
    (inferenceHandler: WebGLInferenceHandler, inputs: Tensor[], attributes: UpsampleAttributes): Tensor[] => {
      validateInputs(inputs, attributes);
      const output = inferenceHandler.run(
          {
            ...resizeProgramMetadata,
            cacheHint: attributes.cacheKey,
            get: () => createPackedResizeProgramInfo(inferenceHandler, inputs, attributes)
          },
          inputs);
      return [output];
    };

export const parseResizeAttributesV10: OperatorInitialization<UpsampleAttributes> =
    (node: Graph.Node): UpsampleAttributes => parseUpsampleAttributes(node, 10);

export const parseResizeAttributesV11: OperatorInitialization<UpsampleAttributes> =
    (node: Graph.Node): UpsampleAttributes => parseUpsampleAttributes(node, 11);

const createPackedResizeProgramInfo =
    (inferenceHandler: WebGLInferenceHandler, inputs: Tensor[], attributes: UpsampleAttributes): ProgramInfo => {
      const glsl = getGlsl(inferenceHandler.session.backend.glContext.version);
      const [scales, outputShape] = prepareInputs(inputs, attributes);

      const isSame =
          scales.every((s: number) => s === 1) && attributes.coordinateTransformMode !== 'tf_crop_and_resize';
      if (isSame) {
        return {
          ...resizeProgramMetadata,
          output: {dims: outputShape, type: inputs[0].type, textureType: TextureType.packed},
          hasMain: true,
          shaderSource: `void main() {
                    vec4 v = ${glsl.texture2D}(X, TexCoords);
                    ${glsl.output} = v;
                }`
        };
      }

      const dim = outputShape.length;
      if (dim < 2) {
        throw new Error(`output dimension should be at least 2, but got ${dim}`);
      }

      const outputHeight = outputShape[dim - 2];
      const outputWidth = outputShape[dim - 1];

      const inputShape = inputs[0].dims;
      if (dim !== inputShape.length) {
        throw new Error(`output dimension should match input ${inputShape.length}, but got ${dim}`);
      }
      const inputHeight = inputShape[dim - 2];
      const inputWidth = inputShape[dim - 1];

      const scalesHeight = scales[dim - 2];
      const scalesWidth = scales[dim - 1];

      let getSourceFracIndex = '';

      if (attributes.mode !== 'linear') {
        // TODO: support other modes
        throw new Error(`resize (packed) does not support mode: '${attributes.mode}'`);
      }
      switch (attributes.coordinateTransformMode) {
        case 'asymmetric':
          getSourceFracIndex = `
                    vec4 getSourceFracIndex(ivec4 coords) {
                        return vec4(coords) / scaleWHWH;
                    }
                `;
          break;
        case 'half_pixel':
          getSourceFracIndex = `
                    vec4 getSourceFracIndex(ivec4 coords) {
                        return (vec4(coords) + 0.5) / scaleWHWH - 0.5;
                    }
                `;
          break;
        case 'pytorch_half_pixel':
          getSourceFracIndex = `
                    vec4 getSourceFracIndex(ivec4 coords) {
                        vec4 fcoords = vec4(coords);
                        return vec4(
                            ${outputWidth}.0 > 1.0 ? (fcoords.x + 0.5) / scaleWHWH.x - 0.5 : 0.0,
                            ${outputHeight}.0 > 1.0 ? (fcoords.y + 0.5) / scaleWHWH.y - 0.5 : 0.0,
                            ${outputWidth}.0 > 1.0 ? (fcoords.z + 0.5) / scaleWHWH.z - 0.5 : 0.0,
                            ${outputHeight}.0 > 1.0 ? (fcoords.w + 0.5) / scaleWHWH.w - 0.5 : 0.0
                          );
                    }
                `;
          break;
        case 'align_corners':
          getSourceFracIndex = `
                    vec4 getSourceFracIndex(ivec4 coords) {
                        vec4 resized = vec4(${outputWidth}.0 - 1.0, ${outputHeight}.0 - 1.0, ${outputWidth}.0 - 1.0,
                            ${outputHeight}.0 - 1.0);
                        vec4 original = vec4(${inputWidth}.0 - 1.0, ${inputHeight}.0 - 1.0, ${inputWidth}.0 - 1.0,
                            ${inputHeight}.0 - 1.0);
                        vec4 new_scale = original / resized;
                        return vec4(coords) * new_scale;
                    }
                `;
          break;
        default:
          // TODO:supporting other coordinateTransformModes
          throw new Error(`resize (packed) does not support coordinateTransformMode: \
                                '${attributes.coordinateTransformMode}'`);
      }

      const coordsDataType = getCoordsDataType(dim);
      const unpackChannel = unpackFromChannel();
      const shaderSource = `
            const vec2 inputWH = vec2(${inputHeight}.0, ${inputWidth}.0);
            const vec4 scaleWHWH = vec4(float(${scalesHeight}), float(${scalesWidth}), float(${scalesHeight}), float(${
          scalesWidth}));
            ${unpackChannel}
            ${getSourceFracIndex}
            float getAValue(int x10, int r, int c, int d) {
                return getChannel(getA(x10, r, c, d), vec2(c, d));
            }
            void main() {
                ${coordsDataType} rc = getOutputCoords();

                int batch = rc[0];
                int depth = rc[1];

                // retrieve the 4 coordinates that is used in the 4 packed output values.
                ivec4 coords = ivec4(rc.wz, rc.w + 1, rc.z + 1);

                // calculate the source index in fraction
                vec4 sourceFrac = getSourceFracIndex(coords);

                // get the lower and upper bound of the 4 values that will be packed into one texel.
                ivec4 x00 = ivec4(max(sourceFrac.xy, vec2(0.0)), min(inputWH - 1.0, ceil(sourceFrac.xy)));
                ivec4 x01 = ivec4(max(sourceFrac.xw, vec2(0.0)), min(inputWH - 1.0, ceil(sourceFrac.xw)));
                ivec4 x10 = ivec4(max(sourceFrac.zy, vec2(0.0)), min(inputWH - 1.0, ceil(sourceFrac.zy)));
                ivec4 x11 = ivec4(max(sourceFrac.zw, vec2(0.0)), min(inputWH - 1.0, ceil(sourceFrac.zw)));

                bool hasNextRow = rc.w < ${outputHeight - 1};
                bool hasNextCol = rc.z < ${outputWidth - 1};

                // pack x00, x01, x10, x11's top-left corner into one vec4 structure
                vec4 topLeft = vec4(
                    getAValue(batch, depth, x00.x, x00.y),
                    hasNextCol ? getAValue(batch, depth, x01.x, x01.y) : 0.0,
                    hasNextRow ? getAValue(batch, depth, x10.x, x10.y) : 0.0,
                    (hasNextRow && hasNextCol) ? getAValue(batch, depth, x11.x, x11.y) : 0.0);

                // pack x00, x01, x10, x11's top-right corner into one vec4 structure
                vec4 topRight = vec4(
                    getAValue(batch, depth, x00.x, x00.w),
                    hasNextCol ? getAValue(batch, depth, x01.x, x01.w) : 0.0,
                    hasNextRow ? getAValue(batch, depth, x10.x, x10.w) : 0.0,
                    (hasNextRow && hasNextCol) ? getAValue(batch, depth, x11.x, x11.w) : 0.0);

                // pack x00, x01, x10, x11's bottom-left corner into one vec4 structure
                vec4 bottomLeft = vec4(
                    getAValue(batch, depth, x00.z, x00.y),
                    hasNextCol ? getAValue(batch, depth, x01.z, x01.y) : 0.0,
                    hasNextRow ? getAValue(batch, depth, x10.z, x10.y) : 0.0,
                    (hasNextRow && hasNextCol) ? getAValue(batch, depth, x11.z, x11.y) : 0.0);

                // pack x00, x01, x10, x11's bottom-right corner into one vec4 structure
                vec4 bottomRight = vec4(
                    getAValue(batch, depth, x00.z, x00.w),
                    hasNextCol ? getAValue(batch, depth, x01.z, x01.w) : 0.0,
                    hasNextRow ? getAValue(batch, depth, x10.z, x10.w) : 0.0,
                    (hasNextRow && hasNextCol) ? getAValue(batch, depth, x11.z, x11.w) : 0.0);

                // calculate the interpolation fraction on u and v direction
                vec4 frac = vec4(sourceFrac) - floor(sourceFrac);
                vec4 clampFrac = clamp(frac, vec4(0.0), vec4(1.0));

                vec4 top = mix(topLeft, topRight, clampFrac.ywyw);
                vec4 bottom = mix(bottomLeft, bottomRight, clampFrac.ywyw);
                vec4 newValue = mix(top, bottom, clampFrac.xxzz);

                ${glsl.output} = vec4(newValue);
            }
        `;
      return {
        ...resizeProgramMetadata,
        output: {dims: outputShape, type: inputs[0].type, textureType: TextureType.packed},
        hasMain: true,
        shaderSource
      };
    };


const prepareInputs = (inputs: Tensor[], attributes: UpsampleAttributes): [readonly number[], readonly number[]] => {
  const x = inputs[0];
  const xDims = x.dims;

  let scales = attributes.scales;
  let outputSizes: number[]|undefined;
  if (scales.length === 0) {
    const scalesTensor = inputs[attributes.scalesInputIdx];
    if (scalesTensor && scalesTensor.size !== 0) {
      if (inputs[attributes.sizesInputIdx]) {
        throw new Error('Only one of scales or sizes must be provided as input.');
      }
      scales = parseScalesData(scalesTensor, attributes.mode, attributes.isResize);
    } else {
      const sizesTensor = inputs[attributes.sizesInputIdx];
      if (!sizesTensor || sizesTensor.size === 0) {
        throw new Error('Either scales or sizes MUST be provided as input.');
      }

      outputSizes = Array.from(sizesTensor.integerData);
      scales = parseScalesDataFromOutputSize(outputSizes, xDims, attributes.mode, attributes.isResize);
    }
  } else {
    if (inputs[attributes.sizesInputIdx]) {
      throw new Error('Only one of scales or sizes must be provided as input.');
    }
  }

  const yDims = outputSizes || (xDims.map((dim, i) => Math.floor(dim * scales[i])));

  return [scales, yDims];
};

const parseScalesData = (scale: Tensor, mode: string, isResize: boolean): number[] => {
  const scales = Array.from(scale.floatData);
  scalesValidation(scales, mode, isResize);
  return scales;
};

const parseScalesDataFromOutputSize =
    (yDims: readonly number[], xDims: readonly number[], mode: string, isResize: boolean): number[] => {
      const length = xDims.length;
      const scales = new Array<number>(length);

      for (let i = 0, end = length; i < end; i++) {
        if (xDims[i] === 0) {
          if (yDims[i] !== 0) {
            throw new Error('Input dim is zero but required output dim is non-zero.');
          }
          scales[i] = 1;
        } else {
          scales[i] = yDims[i] / xDims[i];
        }
      }
      scalesValidation(scales, mode, isResize);
      return scales;
    };

// roi data is not used yet. but leave here for future usage.
// const getRoi = (inputs: Tensor[], attributes: UpsampleAttributes) : number[] => {
//     let roi: number[] = [];
//     if (attributes.needRoiInput) {
//         if (attributes.roiInputIdx <= 0) {
//             throw new Error('Invalid roi input index.');
//         }
//         const roiTensor = inputs[attributes.roiInputIdx];
//         roi = roiTensor.size > 0 ? Array.from(roiTensor.floatData) : [];
//     } else {
//         roi = new Array(inputs[0].dims.length * 2).fill(0);
//     }
//     return roi;
// };