export async function extractFaceTensors(
  input: tf.Tensor | NetInput | TNetInput,
  detections: Array<FaceDetection|Rect>
): Promise<tf.Tensor4D[]> {

  const image = input instanceof tf.Tensor
    ? input
    : await toNetInput(input)

  return tf.tidy(() => {
    const imgTensor = getImageTensor(image)

    // TODO handle batches
    const [batchSize, imgHeight, imgWidth, numChannels] = imgTensor.shape

    const boxes = detections.map(
      det => det instanceof FaceDetection
        ? det.forSize(imgWidth, imgHeight).getBox().floor()
        : det
    )
    const faceTensors = boxes.map(({ x, y, width, height }) =>
      tf.slice(imgTensor, [0, y, x, 0], [1, height, width, numChannels])
    )

    return faceTensors
  })
}

export function mobileNetV1(x: tf.Tensor4D, params: MobileNetV1.Params) {
  return tf.tidy(() => {

    let conv11 = null
    let out = pointwiseConvLayer(x, params.conv_0_params, [2, 2])

    params.conv_pair_params.forEach((param, i) => {
      const layerIdx = i + 1
      const depthwiseConvStrides = getStridesForLayerIdx(layerIdx)
      out = depthwiseConvLayer(out, param.depthwise_conv_params, depthwiseConvStrides)
      out = pointwiseConvLayer(out, param.pointwise_conv_params, [1, 1])
      if (layerIdx === 11) {
        conv11 = out
      }
    })

    if (conv11 === null) {
      throw new Error('mobileNetV1 - output of conv layer 11 is null')
    }

    return {
      out,
      conv11: conv11 as any
    }

  })
}

export function resizeLayer(x: tf.Tensor4D) {
  return tf.tidy(() => {

    const resized = tf.image.resizeBilinear(x, resizedImageSize, false)
    return tf.sub(tf.mul(resized, weight), bias)

  })
}

  public async forward(input: tf.Tensor | NetInput | TNetInput) {
    const netInput = input instanceof tf.Tensor
      ? input
      : await toNetInput(input)

    return tf.tidy(() =>
      this.forwardTensor(padToSquare(getImageTensor(netInput)))
    )
  }

export function normalize(x: tf.Tensor4D): tf.Tensor4D {
  return tf.tidy(() => {
    const avg_r = tf.fill([1, 150, 150, 1], 122.782);
    const avg_g = tf.fill([1, 150, 150, 1], 117.001);
    const avg_b = tf.fill([1, 150, 150, 1], 104.298);
    const avg_rgb = tf.concat([avg_r, avg_g, avg_b], 3)

    return tf.div(tf.sub(x, avg_rgb), tf.scalar(256))
  })
}

export function fullyConnectedLayer(
  x: tf.Tensor2D,
  params: FCParams
): tf.Tensor2D {
  return tf.tidy(() =>
    tf.add(
      tf.matMul(x, params.weights),
      params.bias
    )
  )
}

  public async detectLandmarks(input: tf.Tensor | NetInput | TNetInput) {
    if (!this._params) {
      throw new Error('FaceLandmarkNet - load model before inference')
    }

    const netInput = input instanceof tf.Tensor
      ? input
      : await toNetInput(input)

    let imageDimensions: Dimensions | undefined

    const outTensor = tf.tidy(() => {
      const params = this._params

      let imgTensor = getImageTensor(netInput)
      const [height, width] = imgTensor.shape.slice(1)
      imageDimensions = { width, height }


      // work with 128 x 128 sized face images
      if (imgTensor.shape[1] !== 128 || imgTensor.shape[2] !== 128) {
        imgTensor = tf.image.resizeBilinear(imgTensor, [128, 128])
      }

      let out = conv(imgTensor, params.conv0_params)
      out = maxPool(out)
      out = conv(out, params.conv1_params)
      out = conv(out, params.conv2_params)
      out = maxPool(out)
      out = conv(out, params.conv3_params)
      out = conv(out, params.conv4_params)
      out = maxPool(out)
      out = conv(out, params.conv5_params)
      out = conv(out, params.conv6_params)
      out = maxPool(out, [1, 1])
      out = conv(out, params.conv7_params)
      const fc0 = tf.relu(fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc0_params))
      const fc1 = fullyConnectedLayer(fc0, params.fc1_params)

      return fc1
    })

    const faceLandmarksArray = Array.from(await outTensor.data())
    outTensor.dispose()

    const xCoords = faceLandmarksArray.filter((c, i) => (i - 1) % 2)
    const yCoords = faceLandmarksArray.filter((c, i) => i % 2)

    return new FaceLandmarks(
      Array(68).fill(0).map((_, i) => new Point(xCoords[i], yCoords[i])),
      imageDimensions as Dimensions
    )
  }

    it('is padded to square by 2 columns', () => tf.tidy(() => {
      const imgTensor = tf.tensor4d(Array(24).fill(1), [1, 4, 2, 3])
      const result = padToSquare(imgTensor)

      expect(result.shape).toEqual([1, 4, 4, 3])

      const paddedCols = tf.unstack(result, 2)
      expect(paddedCols.length).toEqual(4)
      expect(paddedCols[0].dataSync()).toEqual(ones(12))
      expect(paddedCols[1].dataSync()).toEqual(ones(12))
      expect(paddedCols[2].dataSync()).toEqual(zeros(12))
      expect(paddedCols[3].dataSync()).toEqual(zeros(12))
    }))

export function pointwiseConvLayer(
  x: tf.Tensor4D,
  params: PointwiseConvParams,
  strides: [number, number]
) {
  return tf.tidy(() => {

    let out = tf.conv2d(x, params.filters, strides, 'same')
    out = tf.add(out, params.batch_norm_offset)
    return tf.clipByValue(out, 0, 6)

  })
}