/**
 * @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 {Scalar, Tensor, Tensor1D, tidy, util} from '@tensorflow/tfjs-core';
// tslint:disable-next-line: no-imports-from-dist
import * as tfOps from '@tensorflow/tfjs-core/dist/ops/ops_for_converter';

import {NamedTensorsMap} from '../../data/types';
import {ExecutionContext} from '../../executor/execution_context';
import {InternalOpExecutor, Node} from '../types';

import {getParamValue} from './utils';

export const executeOp: InternalOpExecutor =
    (node: Node, tensorMap: NamedTensorsMap, context: ExecutionContext,
     ops = tfOps): Tensor[] => {
      switch (node.op) {
        case 'ConcatV2':
        case 'Concat': {
          const n = getParamValue('n', node, tensorMap, context) as number;
          const axis =
              getParamValue('axis', node, tensorMap, context) as number;
          let inputs =
              getParamValue('tensors', node, tensorMap, context) as Tensor[];
          inputs = inputs.slice(0, n);
          return [ops.concat(inputs, axis)];
        }
        case 'Gather': {
          const input = getParamValue('x', node, tensorMap, context) as Tensor;
          const indices =
              getParamValue('indices', node, tensorMap, context) as Tensor1D;
          return [ops.gather(input, ops.cast(indices, 'int32'), 0)];
        }
        case 'GatherV2': {
          const axis =
              getParamValue('axis', node, tensorMap, context) as number;
          const batchDims =
              getParamValue('batchDims', node, tensorMap, context) as number;
          const input = getParamValue('x', node, tensorMap, context) as Tensor;
          const indices =
              getParamValue('indices', node, tensorMap, context) as Tensor1D;
          return [ops.gather(
              input, ops.cast(indices, 'int32'), axis, batchDims)];
        }
        case 'Reverse': {
          const dims =
              getParamValue('dims', node, tensorMap, context) as boolean[];
          const axis = [];
          for (let i = 0; i < dims.length; i++) {
            if (dims[i]) {
              axis.push(i);
            }
          }
          const input = getParamValue('x', node, tensorMap, context) as Tensor;
          return [ops.reverse(input, axis)];
        }
        case 'ReverseV2': {
          const axis =
              getParamValue('axis', node, tensorMap, context) as number[];
          const input = getParamValue('x', node, tensorMap, context) as Tensor;
          return [ops.reverse(input, axis)];
        }
        case 'Slice': {
          // tslint:disable-next-line:no-any
          const begin = getParamValue('begin', node, tensorMap, context) as any;
          // tslint:disable-next-line:no-any
          const size = getParamValue('size', node, tensorMap, context) as any;
          return [ops.slice(
              getParamValue('x', node, tensorMap, context) as Tensor, begin,
              size)];
        }
        case 'StridedSlice': {
          const begin =
              getParamValue('begin', node, tensorMap, context) as number[];
          const end =
              getParamValue('end', node, tensorMap, context) as number[];
          const strides =
              getParamValue('strides', node, tensorMap, context) as number[];
          const beginMask =
              getParamValue('beginMask', node, tensorMap, context) as number;
          const endMask =
              getParamValue('endMask', node, tensorMap, context) as number;
          const ellipsisMask =
              getParamValue('ellipsisMask', node, tensorMap, context) as number;
          const newAxisMask =
              getParamValue('newAxisMask', node, tensorMap, context) as number;
          const shrinkAxisMask =
              getParamValue('shrinkAxisMask', node, tensorMap, context) as
              number;
          const tensor = getParamValue('x', node, tensorMap, context) as Tensor;

          return [ops.stridedSlice(
              tensor, begin, end, strides, beginMask, endMask, ellipsisMask,
              newAxisMask, shrinkAxisMask)];
        }
        case 'Pack': {
          return tidy(() => {
            const axis =
                getParamValue('axis', node, tensorMap, context) as number;
            const tensors =
                getParamValue('tensors', node, tensorMap, context) as Tensor[];
            // Reshape the tensors to the first tensor's shape if they don't
            // match.
            const shape = tensors[0].shape;
            const squeezedShape = ops.squeeze(tensors[0]).shape;
            const mapped = tensors.map(tensor => {
              const sameShape = util.arraysEqual(tensor.shape, shape);
              if (!sameShape &&
                  !util.arraysEqual(ops.squeeze(tensor).shape, squeezedShape)) {
                throw new Error('the input tensors shape does not match');
              }
              return sameShape ? tensor : ops.reshape(tensor, shape);
            });
            return [ops.stack(mapped, axis)];
          });
        }
        case 'Unpack': {
          const axis =
              getParamValue('axis', node, tensorMap, context) as number;
          const tensor =
              getParamValue('tensor', node, tensorMap, context) as Tensor;
          return ops.unstack(tensor, axis);
        }
        case 'Tile': {
          const reps =
              getParamValue('reps', node, tensorMap, context) as number[];
          return [ops.tile(
              getParamValue('x', node, tensorMap, context) as Tensor, reps)];
        }
        case 'Split':
        case 'SplitV': {
          const axis =
              getParamValue('axis', node, tensorMap, context) as number;
          const numOrSizeSplits =
              getParamValue('numOrSizeSplits', node, tensorMap, context) as
                  number |
              number[];
          const tensor = getParamValue('x', node, tensorMap, context) as Tensor;

          return ops.split(tensor, numOrSizeSplits, axis);
        }
        case 'ScatterNd': {
          const indices =
              getParamValue('indices', node, tensorMap, context) as Tensor;
          const values =
              getParamValue('values', node, tensorMap, context) as Tensor;
          const shape =
              getParamValue('shape', node, tensorMap, context) as number[];
          return [ops.scatterND(indices, values, shape)];
        }
        case 'GatherNd': {
          const x = getParamValue('x', node, tensorMap, context) as Tensor;
          const indices =
              getParamValue('indices', node, tensorMap, context) as Tensor;
          return [ops.gatherND(x, indices)];
        }
        case 'SparseToDense': {
          const indices =
              getParamValue('sparseIndices', node, tensorMap, context) as
              Tensor;
          const shape =
              getParamValue('outputShape', node, tensorMap, context) as
              number[];
          const sparseValues =
              getParamValue('sparseValues', node, tensorMap, context) as Tensor;
          const defaultValue =
              getParamValue('defaultValue', node, tensorMap, context) as Scalar;
          return [ops.sparseToDense(
              indices, sparseValues, shape,
              sparseValues.dtype === defaultValue.dtype ?
                  defaultValue :
                  ops.cast(defaultValue, sparseValues.dtype))];
        }
        case 'TensorScatterUpdate': {
          const indices =
              getParamValue('indices', node, tensorMap, context) as Tensor;
          const values =
              getParamValue('values', node, tensorMap, context) as Tensor;
          const tensor =
              getParamValue('tensor', node, tensorMap, context) as Tensor;
          return [ops.tensorScatterUpdate(tensor, indices, values)];
        }
        default:
          throw TypeError(`Node type ${node.op} is not implemented`);
      }
    };

export const CATEGORY = 'slice_join';
