/**
 * @license
 * Copyright 2018 Google LLC. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */

import * as tf from '@tensorflow/tfjs';
import {backend_util, BackendTimingInfo, DataId, DataType, KernelBackend, ModelTensorInfo, Rank, Scalar, scalar, ScalarLike, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, TensorInfo, tidy, util} from '@tensorflow/tfjs';
import {isArray, isNullOrUndefined} from 'util';

import {encodeInt32ArrayAsInt64, Int64Scalar} from './int64_tensors';
import {TensorMetadata, TFEOpAttr, TFJSBinding} from './tfjs_binding';

// tslint:disable-next-line:no-require-imports
const messages = require('./proto/api_pb');

type TensorData = {
  shape: number[],
  dtype: number,
  values: backend_util.BackendValues,
  id: number,
  refCount: number;
};

export class NodeJSKernelBackend extends KernelBackend {
  binding: TFJSBinding;
  isGPUPackage: boolean;
  isUsingGpuDevice: boolean;
  private tensorMap: tf.DataStorage<TensorData>;

  constructor(binding: TFJSBinding, packageName: string) {
    super();
    this.binding = binding;
    this.isGPUPackage = packageName === '@tensorflow/tfjs-node-gpu';
    this.isUsingGpuDevice = this.binding.isUsingGpuDevice();
    this.tensorMap = new tf.DataStorage<TensorData>(this, tf.engine());
  }

  getDTypeInteger(dtype: DataType): number {
    switch (dtype) {
      case 'float32':
        return this.binding.TF_FLOAT;
      case 'int32':
        return this.binding.TF_INT32;
      case 'bool':
        return this.binding.TF_BOOL;
      case 'complex64':
        return this.binding.TF_COMPLEX64;
      case 'string':
        return this.binding.TF_STRING;
      default:
        throw new Error(`Unsupported DType: ${dtype}`);
    }
  }

  private typeAttributeFromTensor(value: Tensor): number {
    return this.getDTypeInteger(value.dtype);
  }

  // Creates a new Tensor and maps the dataId to the passed in ID.
  private createOutputTensor(metadata: TensorMetadata): Tensor {
    const newId = {};

    this.tensorMap.set(newId, {
      shape: metadata.shape,
      dtype: metadata.dtype,
      id: metadata.id,
      values: null,
      refCount: 1
    });

    let dtype: DataType;
    switch (metadata.dtype) {
      case this.binding.TF_FLOAT:
        dtype = 'float32';
        break;
      case this.binding.TF_INT32:
        dtype = 'int32';
        break;
      case this.binding.TF_INT64:
        console.warn('INT64 output tensor will be stored as BigInt64Array.');
        // INT64 is not supported in TFJS yet, cast it to int32.
        dtype = 'int32';
        break;
      case this.binding.TF_BOOL:
        dtype = 'bool';
        break;
      case this.binding.TF_COMPLEX64:
        dtype = 'complex64';
        break;
      case this.binding.TF_STRING:
        dtype = 'string';
        break;
      case this.binding.TF_RESOURCE:
        // NOTE(cais): We currently represent resource-type Tensors
        // as string of ubytes.
        dtype = 'string';
        break;
      case this.binding.TF_UINT8:
        // TensorFlow uses UINT8 as dtype for image tensor. UINT8 is not
        // supported in TFJS yet, cast it to int32.
        dtype = 'int32';
        break;
      default:
        throw new Error(`Unknown dtype enum ${metadata.dtype}`);
    }

    // TODO(yassogba) Enable this once all the kernels are removed from backend.
    // We can then change the return type from Tensor to TensorInfo.
    // return {dataId: newId, shape: metadata.shape, dtype};

    const tensorInfo:
        TensorInfo = {dataId: newId, shape: metadata.shape, dtype};
    return tf.engine().makeTensorFromTensorInfo(tensorInfo);
  }

  // Prepares Tensor instances for Op execution.
  private getInputTensorIds(tensors: Array<TensorInfo|Int64Scalar>): number[] {
    const ids: number[] = [];
    for (let i = 0; i < tensors.length; i++) {
      if (tensors[i] instanceof Int64Scalar) {
        // Then `tensors[i]` is a Int64Scalar, which we currently represent
        // using an `Int32Array`.
        const value = (tensors[i] as Int64Scalar).valueArray;
        const id = this.binding.createTensor([], this.binding.TF_INT64, value);
        ids.push(id);
      } else {
        const info = this.tensorMap.get((tensors[i] as TensorInfo).dataId);
        // TODO - what about ID in this case? Handle in write()??
        if (info.values != null) {
          // Values were delayed to write into the TensorHandle. Do that before
          // Op execution and clear stored values.
          info.id =
              this.binding.createTensor(info.shape, info.dtype, info.values);
          info.values = null;
        }
        ids.push(info.id);
      }
    }
    return ids;
  }

  createReductionOpAttrs(tensor: TensorInfo, keepDims = false): TFEOpAttr[] {
    return [
      {name: 'keep_dims', type: this.binding.TF_ATTR_BOOL, value: keepDims},
      createTensorsTypeOpAttr('T', tensor.dtype),
      createTensorsTypeOpAttr('Tidx', 'int32')
    ];
  }

  floatPrecision(): 16|32 {
    return 32;
  }

  epsilon(): number {
    return super.epsilon();
  }

  /**
   * Executes an op that has a single input and output.
   *
   * Helper function to wrap executeSingleOutput in a particular case.
   * @param name The name of the Op to execute.
   * @param input The input Tensor for the Op.
   */
  executeSingleInput(name: string, input: TensorInfo): Tensor {
    const opAttrs = [createTensorsTypeOpAttr('T', input.dtype)];
    return this.executeSingleOutput(name, opAttrs, [input]);
  }

  /**
   * Executes a TensorFlow Eager Op that provides one output Tensor.
   * @param name The name of the Op to execute.
   * @param opAttrs The list of Op attributes required to execute.
   * @param inputs The list of input Tensors for the Op.
   * @return A resulting Tensor from Op execution.
   */
  executeSingleOutput(name: string, opAttrs: TFEOpAttr[], inputs: TensorInfo[]):
      Tensor {
    const outputMetadata = this.binding.executeOp(
        name, opAttrs, this.getInputTensorIds(inputs), 1);
    return this.createOutputTensor(outputMetadata[0]);
  }

  /**
   * Executes a TensorFlow Eager Op that provides multiple output Tensors.
   * @param name The name of the Op to execute.
   * @param opAttrs The list of Op attributes required to execute.
   * @param inputs The list of input Tensors for the Op.
   * @param numOutputs The number of output Tensors for Op execution.
   * @return A resulting Tensor array from Op execution.
   */
  executeMultipleOutputs(
      name: string, opAttrs: TFEOpAttr[], inputs: TensorInfo[],
      numOutputs: number): Tensor[] {
    const outputMetadata = this.binding.executeOp(
        name, opAttrs, this.getInputTensorIds(inputs), numOutputs);
    return outputMetadata.map(m => this.createOutputTensor(m));
  }

  numDataIds(): number {
    return this.tensorMap.numDataIds();
  }

  dispose(): void {}

  async read(dataId: DataId): Promise<backend_util.BackendValues> {
    return this.readSync(dataId);
  }

  readSync(dataId: DataId): backend_util.BackendValues {
    if (!this.tensorMap.has(dataId)) {
      throw new Error(`Tensor ${dataId} was not registered!`);
    }
    const info = this.tensorMap.get(dataId);
    if (info.values != null) {
      return info.values;
    } else {
      return this.binding.tensorDataSync(info.id);
    }
  }

  /**
   * Dispose the memory if the dataId has 0 refCount. Return true if the memory
   * is released, false otherwise.
   * @param dataId
   * @oaram force Optional, remove the data regardless of refCount
   */
  disposeData(dataId: DataId, force = false): boolean {
    // No-op if already disposed.
    if (this.tensorMap.has(dataId)) {
      const id = this.tensorMap.get(dataId).id;
      this.tensorMap.get(dataId).refCount--;
      if (!force && this.tensorMap.get(dataId).refCount > 0) {
        return false;
      }

      if (id != null && id >= 0) {
        this.binding.deleteTensor(id);
      }
      this.tensorMap.delete(dataId);
    }
    return true;
  }
  /** Return refCount of a `TensorData`. */
  refCount(dataId: DataId): number {
    if (this.tensorMap.has(dataId)) {
      const tensorData = this.tensorMap.get(dataId);
      return tensorData.refCount;
    }
    return 0;
  }

  incRef(dataId: DataId) {
    this.tensorMap.get(dataId).refCount++;
  }

  move(
      dataId: DataId, values: backend_util.BackendValues, shape: number[],
      dtype: DataType, refCount: number): void {
    this.tensorMap.set(
        dataId, {shape, dtype: getTFDType(dtype), values, id: -1, refCount});
  }

  write(values: backend_util.BackendValues, shape: number[], dtype: DataType):
      DataId {
    const dataId = {};
    this.move(dataId, values, shape, dtype, 1);
    return dataId;
  }

  applyActivation<T extends Tensor>(
      input: T, activation: string, preluActivationWeights?: Tensor,
      leakyreluAlpha?: number): T {
    let result = input;
    if (activation != null) {
      if (activation === 'linear') {
        // No-op
      } else if (activation === 'relu') {
        result = tf.relu(result);
      } else if (activation === 'prelu') {
        result = tf.prelu(result, preluActivationWeights) as T;
      } else if (activation === 'leakyrelu') {
        result = tf.leakyRelu(result, leakyreluAlpha);
      } else if (activation === 'elu') {
        result = tf.elu(result);
      } else if (activation === 'relu6') {
        result = tf.relu6(result);
      } else if (activation === 'sigmoid') {
        result = tf.sigmoid(result);
      } else {
        throw new Error(`Activation: ${
            activation} has not been implemented for the Node.js backend`);
      }
    }
    return result;
  }

  divide(a: Tensor, b: Tensor): Tensor {
    const opAttrs = [createTensorsTypeOpAttr(
        'T', backend_util.upcastType(a.dtype, b.dtype))];
    return this.executeSingleOutput('Div', opAttrs, [a, b]);
  }

  divNoNan(a: Tensor, b: Tensor): Tensor {
    const opAttrs = [createTensorsTypeOpAttr(
        'T', backend_util.upcastType(a.dtype, b.dtype))];
    return this.executeSingleOutput('DivNoNan', opAttrs, [a, b]);
  }

  where(condition: Tensor): Tensor2D {
    return this.executeSingleOutput('Where', [], [condition]) as Tensor2D;
  }

  topKValues<T extends Tensor>(x: T, k: number): Tensor1D {
    throw new Error('Method not implemented.');
  }

  topKIndices(x: Tensor, k: number): Tensor1D {
    throw new Error('Method not implemented.');
  }

  int<T extends Tensor>(x: T): T {
    throw new Error('Method not implemented.');
  }

  decodeJpeg(
      contents: Uint8Array, channels: number, ratio: number,
      fancyUpscaling: boolean, tryRecoverTruncated: boolean,
      acceptableFraction: number, dctMethod: string): Tensor3D {
    const opAttrs = [
      {name: 'channels', type: this.binding.TF_ATTR_INT, value: channels},
      {name: 'ratio', type: this.binding.TF_ATTR_INT, value: ratio}, {
        name: 'fancy_upscaling',
        type: this.binding.TF_ATTR_BOOL,
        value: fancyUpscaling
      },
      {
        name: 'try_recover_truncated',
        type: this.binding.TF_ATTR_BOOL,
        value: tryRecoverTruncated
      },
      {
        name: 'acceptable_fraction',
        type: this.binding.TF_ATTR_FLOAT,
        value: acceptableFraction
      },
      {name: 'dct_method', type: this.binding.TF_ATTR_STRING, value: dctMethod}
    ];
    const inputArgs = [scalar(contents, 'string')];
    return this.executeSingleOutput('DecodeJpeg', opAttrs, inputArgs) as
        Tensor<Rank.R3>;
  }

  decodePng(contents: Uint8Array, channels: number): Tensor3D {
    const opAttrs =
        [{name: 'channels', type: this.binding.TF_ATTR_INT, value: channels}];
    const inputArgs = [scalar(contents, 'string')];

    return this.executeSingleOutput('DecodePng', opAttrs, inputArgs) as
        Tensor<Rank.R3>;
  }

  decodeBmp(contents: Uint8Array, channels: number): Tensor3D {
    const opAttrs =
        [{name: 'channels', type: this.binding.TF_ATTR_INT, value: channels}];
    const inputArgs = [scalar(contents, 'string')];
    return this.executeSingleOutput('DecodeBmp', opAttrs, inputArgs) as
        Tensor<Rank.R3>;
  }

  decodeGif(contents: Uint8Array): Tensor4D {
    const inputArgs = [scalar(contents, 'string')];
    return this.executeSingleOutput('DecodeGif', [], inputArgs) as
        Tensor<Rank.R4>;
  }

  executeEncodeImageOp(
      name: string, opAttrs: TFEOpAttr[], imageData: Uint8Array,
      imageShape: number[]): Tensor {
    const inputTensorId =
        this.binding.createTensor(imageShape, this.binding.TF_UINT8, imageData);
    const outputMetadata =
        this.binding.executeOp(name, opAttrs, [inputTensorId], 1);
    this.binding.deleteTensor(inputTensorId);
    const outputTensorInfo = outputMetadata[0];
    // prevent the tensor data from being converted to a UTF8 string, since
    // the encoded data is not valid UTF8
    outputTensorInfo.dtype = this.binding.TF_UINT8;
    return this.createOutputTensor(outputTensorInfo);
  }

  encodeJpeg(
      imageData: Uint8Array, imageShape: number[], format: ''|'grayscale'|'rgb',
      quality: number, progressive: boolean, optimizeSize: boolean,
      chromaDownsampling: boolean, densityUnit: 'in'|'cm', xDensity: number,
      yDensity: number, xmpMetadata: string): Tensor {
    const opAttrs = [
      {name: 'format', type: this.binding.TF_ATTR_STRING, value: format},
      {name: 'quality', type: this.binding.TF_ATTR_INT, value: quality}, {
        name: 'progressive',
        type: this.binding.TF_ATTR_BOOL,
        value: progressive
      },
      {
        name: 'optimize_size',
        type: this.binding.TF_ATTR_BOOL,
        value: optimizeSize
      },
      {
        name: 'chroma_downsampling',
        type: this.binding.TF_ATTR_BOOL,
        value: chromaDownsampling
      },
      {
        name: 'density_unit',
        type: this.binding.TF_ATTR_STRING,
        value: densityUnit
      },
      {name: 'x_density', type: this.binding.TF_ATTR_INT, value: xDensity},
      {name: 'y_density', type: this.binding.TF_ATTR_INT, value: yDensity}, {
        name: 'xmp_metadata',
        type: this.binding.TF_ATTR_STRING,
        value: xmpMetadata
      }
    ];
    return this.executeEncodeImageOp(
        'EncodeJpeg', opAttrs, imageData, imageShape);
  }

  encodePng(imageData: Uint8Array, imageShape: number[], compression: number):
      Tensor {
    const opAttrs = [
      {name: 'compression', type: this.binding.TF_ATTR_INT, value: compression}
    ];
    return this.executeEncodeImageOp(
        'EncodePng', opAttrs, imageData, imageShape);
  }

  deleteSavedModel(id: number): void {
    this.binding.deleteSavedModel(id);
  }

  loadSavedModelMetaGraph(path: string, tags: string): number {
    return this.binding.loadSavedModel(path, tags);
  }

  private getMappedInputTensorIds(
      inputs: Tensor[], inputTensorInfos: ModelTensorInfo[]) {
    const tensorIds = this.getInputTensorIds(inputs);
    const newTensors = [];
    for (let i = 0; i < inputs.length; i++) {
      if (inputTensorInfos[i] != null) {
        if (inputTensorInfos[i].tfDtype === 'DT_UINT8') {
          const data = Uint8Array.from(inputs[i].dataSync());
          const inputTensorId = this.binding.createTensor(
              inputs[i].shape, this.binding.TF_UINT8, data);
          tensorIds[i] = inputTensorId;
          newTensors.push(i);
        } else if (inputTensorInfos[i].tfDtype === 'DT_INT64') {
          const data =
              encodeInt32ArrayAsInt64(inputs[i].dataSync() as Int32Array);
          const inputTensorId = this.binding.createTensor(
              inputs[i].shape, this.binding.TF_INT64, data);
          tensorIds[i] = inputTensorId;
          newTensors.push(i);
        }
      }
    }
    return {tensorIds, newTensors};
  }

  runSavedModel(
      id: number, inputs: Tensor[], inputTensorInfos: ModelTensorInfo[],
      outputOpNames: string[]): Tensor[] {
    const {tensorIds, newTensors} =
        this.getMappedInputTensorIds(inputs, inputTensorInfos);
    const outputMetadata = this.binding.runSavedModel(
        id, tensorIds, inputTensorInfos.map(info => info.name).join(','),
        outputOpNames.join(','));
    for (let i = 0; i < tensorIds.length; i++) {
      if (newTensors.includes(i)) {
        this.binding.deleteTensor(tensorIds[i]);
      }
    }
    return outputMetadata.map(m => this.createOutputTensor(m));
  }

  // ------------------------------------------------------------
  // TensorBoard-related (tfjs-node-specific) backend kernels.

  summaryWriter(logdir: string): Tensor1D {
    const opAttrs = [
      {
        name: 'shared_name',
        type: this.binding.TF_ATTR_STRING,
        value: `logdir:${logdir}`
      },
      {name: 'container', type: this.binding.TF_ATTR_STRING, value: ''}
    ];
    const writerResource =
        this.executeSingleOutput('SummaryWriter', opAttrs, []);
    return writerResource as Tensor1D;
  }

  createSummaryFileWriter(
      resourceHandle: Tensor, logdir: string, maxQueue?: number,
      flushMillis?: number, filenameSuffix?: string): void {
    const inputArgs = [
      resourceHandle, scalar(logdir),
      scalar(maxQueue == null ? 10 : maxQueue, 'int32'),
      scalar(flushMillis == null ? 2 * 60 * 1000 : flushMillis, 'int32'),
      scalar(filenameSuffix == null ? '.v2' : filenameSuffix)
    ];
    this.executeMultipleOutputs('CreateSummaryFileWriter', [], inputArgs, 0);
  }

  writeScalarSummary(
      resourceHandle: Tensor, step: number, name: string,
      value: Scalar|number): void {
    tidy(() => {
      util.assert(
          Number.isInteger(step),
          () => `step is expected to be an integer, but is instead ${step}`);
      const inputArgs: Array<Tensor|Int64Scalar> =
          [resourceHandle, new Int64Scalar(step), scalar(name, 'string')];

      let typeAttr: number;
      if (typeof value === 'number') {
        inputArgs.push(scalar(value));
        typeAttr = this.binding.TF_FLOAT;
      } else {
        // `value` is a Scalar.
        util.assert(
            value.rank === 0,
            () => `A non-scalar tensor (rank ${value.rank}) is passed to ` +
                `writeScalarSummary()`);
        inputArgs.push(value);
        typeAttr = this.typeAttributeFromTensor(value);
      }
      const opAttrs: TFEOpAttr[] =
          [{name: 'T', type: this.binding.TF_ATTR_TYPE, value: typeAttr}];
      const ids = this.getInputTensorIds(inputArgs);
      this.binding.executeOp('WriteScalarSummary', opAttrs, ids, 0);
      // release the tensorflow tensor for Int64Scalar value of step
      this.binding.deleteTensor(ids[1]);
    });
  }

  writeHistogramSummary(
      resourceHandle: Tensor, step: number, name: string, data: Tensor,
      bucketCount: number|undefined, description: string|undefined): void {
    tidy(() => {
      util.assert(
          Number.isInteger(step),
          () => `step is expected to be an integer, but is instead ${step}`);

      // We use the WriteSummary op, and not WriteHistogramSummary. The
      // difference is that WriteHistogramSummary takes a tensor of any shape,
      // and places the values in 30 buckets, while WriteSummary expects a
      // tensor which already describes the bucket widths and counts.
      //
      // If we were to use WriteHistogramSummary, we wouldn't have to
      // implement the "bucketization" of the input tensor, but we also
      // wouldn't have control over the number of buckets, or the description
      // of the graph.
      //
      // Therefore, we instead use WriteSummary, which makes it possible to
      // support these features. However, the trade-off is that we have to
      // implement our own "bucketization", and have to write the summary as a
      // protobuf message.
      const content = new messages.HistogramPluginData().setVersion(0);
      const pluginData = new messages.SummaryMetadata.PluginData()
                             .setPluginName('histograms')
                             .setContent(content.serializeBinary());
      const summary = new messages.SummaryMetadata()
                          .setPluginData(pluginData)
                          .setDisplayName(null)
                          .setSummaryDescription(description);
      const summaryTensor = scalar(summary.serializeBinary(), 'string');
      const nameTensor = scalar(name, 'string');
      const stepScalar = new Int64Scalar(step);
      const buckets = this.buckets(data, bucketCount);
      util.assert(
          buckets.rank === 2 && buckets.shape[1] === 3,
          () => `Expected buckets to have shape [k, 3], but they had shape ${
              buckets.shape}`);
      util.assert(
          buckets.dtype === 'float32',
          () => `Expected buckets to have dtype float32, but they had dtype ${
              buckets.dtype}`);
      const inputArgs: Array<Tensor|Int64Scalar> =
          [resourceHandle, stepScalar, buckets, nameTensor, summaryTensor];
      const typeAttr = this.typeAttributeFromTensor(buckets);
      const opAttrs: TFEOpAttr[] =
          [{name: 'T', type: this.binding.TF_ATTR_TYPE, value: typeAttr}];
      const ids = this.getInputTensorIds(inputArgs);
      this.binding.executeOp('WriteSummary', opAttrs, ids, 0);
      // release the tensorflow tensor for Int64Scalar value of step
      this.binding.deleteTensor(ids[1]);
    });
  }

  flushSummaryWriter(resourceHandle: Tensor): void {
    const inputArgs: Tensor[] = [resourceHandle];
    this.executeMultipleOutputs('FlushSummaryWriter', [], inputArgs, 0);
  }

  /**
   * Group data into histogram buckets.
   *
   * @param data A `Tensor` of any shape. Must be castable to `float32`
   * @param bucketCount Optional positive `number`
   * @returns A `Tensor` of shape `[k, 3]` and type `float32`. The `i`th row
   *     is
   *   a triple `[leftEdge, rightEdge, count]` for a single bucket. The value
   * of `k` is either `bucketCount`, `1` or `0`.
   */
  private buckets(data: Tensor, bucketCount?: number): Tensor<tf.Rank> {
    if (data.size === 0) {
      return tf.tensor([], [0, 3], 'float32');
    }

    // 30 is the default number of buckets in the TensorFlow Python
    // implementation. See
    // https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/histogram/summary_v2.py
    bucketCount = bucketCount !== undefined ? bucketCount : 30;
    util.assert(
        Number.isInteger(bucketCount) && bucketCount > 0,
        () =>
            `Expected bucket count to be a strictly positive integer, but it was ` +
            `${bucketCount}`);
    data = data.flatten();
    data = data.cast('float32');
    const min: Scalar = data.min();
    const max: Scalar = data.max();
    const range: Scalar = max.sub(min);
    const isSingular = range.equal(0).arraySync() !== 0;

    if (isSingular) {
      const center = min;
      const bucketStart: Scalar = center.sub(0.5);
      const bucketEnd: Scalar = center.add(0.5);
      const bucketCounts = tf.scalar(data.size, 'float32');
      return tf.concat([bucketStart, bucketEnd, bucketCounts]).reshape([1, 3]);
    }

    const bucketWidth = range.div(bucketCount);
    const offsets = data.sub(min);
    const bucketIndices = offsets.floorDiv(bucketWidth).cast('int32');
    const clampedIndices =
        tf.minimum(bucketIndices, bucketCount - 1).cast('int32');
    const oneHots = tf.oneHot(clampedIndices, bucketCount);
    const bucketCounts = oneHots.sum(0).cast('int32');
    let edges = tf.linspace(min.arraySync(), max.arraySync(), bucketCount + 1);
    // Ensure last value in edges is max (TF's linspace op doesn't do this)
    edges = tf.concat([edges.slice(0, bucketCount), max.reshape([1])], 0) as
        tf.Tensor1D;
    const leftEdges = edges.slice(0, bucketCount);
    const rightEdges = edges.slice(1, bucketCount);
    return tf.stack([leftEdges, rightEdges, bucketCounts.cast('float32')])
        .transpose();
  }

  // ~ TensorBoard-related (tfjs-node-specific) backend kernels.
  // ------------------------------------------------------------

  memory() {
    // Due to automatic garbage collection, the numbers are unreliable.
    // TODO(kreeger): Since there is finalization in C, count the true
    // number of undisposed tensors.
    return {unreliable: true};
  }

  async time(f: () => void): Promise<BackendTimingInfo> {
    const start = process.hrtime();
    f();
    // hrtime() returns tuple of [seconds, nanoseconds], and we need to return
    // milliseconds.
    const elapsed = process.hrtime(start);
    return {kernelMs: elapsed[0] * 1000 + elapsed[1] / 1000000};
  }

  getNumOfSavedModels() {
    return this.binding.getNumOfSavedModels();
  }

  getNumOfTFTensors() {
    return this.binding.getNumOfTensors();
  }
}

/** Returns an instance of the Node.js backend. */
export function nodeBackend(): NodeJSKernelBackend {
  return tf.findBackend('tensorflow') as NodeJSKernelBackend;
}

/** Returns the TF dtype for a given DataType. */
export function getTFDType(dataType: tf.DataType): number {
  const binding = nodeBackend().binding;
  switch (dataType) {
    case 'float32':
      return binding.TF_FLOAT;
    case 'int32':
      return binding.TF_INT32;
    case 'bool':
      return binding.TF_BOOL;
    case 'complex64':
      return binding.TF_COMPLEX64;
    case 'string':
      return binding.TF_STRING;
    // tslint:disable-next-line:no-any
    case 'int64' as any:
      // int64 is not a generally supported dtype in TensorFlow.js
      // (tfjs-core). However, it needs to be included here for the purpose of
      // writing the `step` value to TensorBoard via WriteScalarSummary and
      // other op kernels.
      return binding.TF_INT64;
    default:
      const errorMessage = `Unknown dtype: ${dataType}`;
      throw new Error(errorMessage);
  }
}

/**
 * Creates a TFEOpAttr for a 'type' OpDef attribute from a Tensor or list of
 * Tensors.
 */
export function createTensorsTypeOpAttr(
    attrName: string,
    tensorsOrDtype: tf.Tensor|tf.Tensor[]|tf.DataType): TFEOpAttr {
  if (isNullOrUndefined(tensorsOrDtype)) {
    throw new Error('Invalid input tensors value.');
  }
  return {
    name: attrName,
    type: nodeBackend().binding.TF_ATTR_TYPE,
    value:
        (tensorsOrDtype instanceof tf.Tensor || Array.isArray(tensorsOrDtype)) ?
        getTFDTypeForInputs(tensorsOrDtype) :
        getTFDType(tensorsOrDtype)
  };
}

// TODO(yassogba) remove? who uses this?
export function createOpAttr(
    attrName: string, tensorsOrDtype: tf.Tensor|tf.Tensor[]|tf.DataType,
    value: ScalarLike): TFEOpAttr {
  if (isNullOrUndefined(tensorsOrDtype)) {
    throw new Error('Invalid input tensors value.');
  }
  return {name: attrName, type: nodeBackend().binding.TF_BOOL, value};
}

/** Returns the dtype number for a single or list of input Tensors. */
function getTFDTypeForInputs(tensors: tf.Tensor|tf.Tensor[]): number {
  if (isNullOrUndefined(tensors)) {
    throw new Error('Invalid input tensors value.');
  }
  if (isArray(tensors)) {
    for (let i = 0; i < tensors.length; i++) {
      return getTFDType(tensors[i].dtype);
    }
    return -1;
  } else {
    return getTFDType(tensors.dtype);
  }
}

export function ensureTensorflowBackend() {
  tf.util.assert(
      tf.getBackend() === 'tensorflow',
      () => `Expect the current backend to be "tensorflow", but got "${
          tf.getBackend()}"`);
}