/* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * 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.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

//
// @author raver119@gmail.com
//

#ifndef SCALAR_CU
#define SCALAR_CU

#include <cuda.h>
#include <cuda_runtime.h>
#include <system/op_boilerplate.h>
#include <types/types.h>
#include <loops/pairwise_instantiations.h>

#include "../../types/types.h"
#include "loops/scalar.h"

using namespace simdOps;

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z, typename OpType>
SD_KERNEL static void scalarSimpleShaped(
    void const* x,
    void const* y,
    sd::LongType const* xShapeInfo,
    void* params,
    void* z,
    sd::LongType const* zShapeInfo,
    sd::LongType* allocationBuffer) {

  auto xVals    = reinterpret_cast<X const*>(x);
  auto scalar   = reinterpret_cast<Y const*>(y)[0];
  auto pVals    = reinterpret_cast<Z*>(params);
  auto zVals    = reinterpret_cast<Z*>(z);

  const int totalThreads = gridDim.x * blockDim.x;
  const int tid          = blockIdx.x * blockDim.x + threadIdx.x;

  __shared__ sd::LongType length;
  __shared__ sd::LongType xRank;
  __shared__ const sd::LongType* xShapePtr;
  __shared__ const sd::LongType* xStridePtr;

  __shared__ sd::LongType zRank;
  __shared__ const sd::LongType* zShapePtr;
  __shared__ const sd::LongType* zStridePtr;

  if (threadIdx.x == 0) {
    length     = shape::length(xShapeInfo);
    xRank      = shape::rank(xShapeInfo);
    xShapePtr  = shape::shapeOf(xShapeInfo);
    xStridePtr = shape::stride(xShapeInfo);

    zRank      = shape::rank(zShapeInfo);
    zShapePtr  = shape::shapeOf(zShapeInfo);
    zStridePtr = shape::stride(zShapeInfo);

    //need special handling for scalars as strides may not be set for scalars
    if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
      z[0] = OpType::op(x[0], scalar[0], extraParams);
      return;

    }
  }
  __syncthreads();

  //need special handling for scalars as strides may not be set for scalars
  if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
    return;

  }
  for (sd::LongType i = tid; i < length; i += totalThreads) {
    sd::LongType xCoords[SD_MAX_RANK];
    sd::LongType zCoords[SD_MAX_RANK];
    sd::LongType xOffset;
    sd::LongType zOffset;

    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);

    INDEX2COORDS(i, zRank, zShapePtr, zCoords);
    COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);

    zVals[zOffset] = OpType::op(xVals[xOffset], scalar, pVals);
  }
}

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z, typename OpType>
SD_KERNEL static void scalarAlongDimension(
    void const* x,
    sd::LongType const* xShapeInfo,
    void* extraParams,
    void* z,
    sd::LongType const* zShapeInfo,
    void const* scalars,
    sd::LongType* dimension,
    sd::LongType dimensionLength,
    sd::LongType const* tadShapeInfo,
    sd::LongType const* tadOffsets,
    sd::LongType const* tadShapeInfoZ,
    sd::LongType const* tadOffsetsZ) {

  auto xVals     = reinterpret_cast<X const*>(x);
  auto zVals     = reinterpret_cast<Z*>(z);
  auto pVals     = reinterpret_cast<Z*>(extraParams);
  auto sVals     = reinterpret_cast<Y const*>(scalars);

  if (tadShapeInfoZ == nullptr) {
    tadShapeInfoZ = tadShapeInfo;
    tadOffsetsZ   = tadOffsets;
  }

  __shared__ sd::LongType tadLength;
  __shared__ sd::LongType numTads;

  __shared__ sd::LongType xRank;
  __shared__ const sd::LongType* xShapePtr;
  __shared__ const sd::LongType* xStridePtr;

  __shared__ sd::LongType tadRank;
  __shared__ const sd::LongType* tadShapePtr;
  __shared__ const sd::LongType* tadStridePtr;

  __shared__ sd::LongType tadRankZ;
  __shared__ const sd::LongType* tadShapePtrZ;
  __shared__ const sd::LongType* tadStridePtrZ;

  if (threadIdx.x == 0) {
    tadLength      = shape::length(tadShapeInfo);
    numTads        = shape::length(xShapeInfo) / tadLength;

    xRank          = shape::rank(xShapeInfo);
    xShapePtr      = shape::shapeOf(xShapeInfo);
    xStridePtr     = shape::stride(xShapeInfo);

    tadRank        = shape::rank(tadShapeInfo);
    tadShapePtr    = shape::shapeOf(tadShapeInfo);
    tadStridePtr   = shape::stride(tadShapeInfo);

    tadRankZ       = shape::rank(tadShapeInfoZ);
    tadShapePtrZ   = shape::shapeOf(tadShapeInfoZ);
    tadStridePtrZ  = shape::stride(tadShapeInfoZ);
  }
  __syncthreads();

  for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
    auto zPart   = zVals + tadOffsetsZ[r];
    auto xPart   = xVals + tadOffsets[r];
    auto scalarV = sVals[r];

    for (sd::LongType f = threadIdx.x; f < tadLength; f += blockDim.x) {
      sd::LongType xCoords[SD_MAX_RANK];
      sd::LongType zCoords[SD_MAX_RANK];
      sd::LongType xOffset;
      sd::LongType zOffset;

      INDEX2COORDS(f, tadRank, tadShapePtr, xCoords);
      COORDS2INDEX(tadRank, tadStridePtr, xCoords, xOffset);

      INDEX2COORDS(f, tadRankZ, tadShapePtrZ, zCoords);
      COORDS2INDEX(tadRankZ, tadStridePtrZ, zCoords, zOffset);

      zPart[zOffset] = OpType::op(xPart[xOffset], scalarV, pVals);
    }
  }
}

namespace functions {
namespace scalar {

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z>
template <typename OpType>
void SD_HOST ScalarTransform<X, Y, Z>::intermediateShaped(
    dim3& launchDims,
    cudaStream_t* stream,
    void const* vx,
    sd::LongType const* xShapeInfo,
    sd::LongType const* hxShapeInfo,  // Unused in the kernel
    void* vz,
    sd::LongType const* zShapeInfo,
    sd::LongType const* hzShapeInfo,  // Unused in the kernel
    void const* vscalar,
    void* vextraParams,
    sd::LongType* allocPointer) {

  scalarSimpleShaped<X, Y, Z, OpType>
      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
          vx,
          vscalar,
          xShapeInfo,
          vextraParams,
          vz,
          zShapeInfo,
          allocPointer);
  sd::DebugHelper::checkErrorCode(stream, "scalarSimpleShaped(...) failed");
}

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z>
template <typename OpType>
void SD_HOST ScalarTransform<X, Y, Z>::intermediateAlongDimension(
    dim3& launchDims,
    cudaStream_t* stream,
    void const* x,
    sd::LongType const* xShapeInfo,
    void* z,
    sd::LongType const* zShapeInfo,
    void const* scalars,
    void* extraParams,
    sd::LongType* dimension,
    sd::LongType dimensionLength,
    sd::LongType const* tadShapeInfo,
    sd::LongType const* tadOffsets,
    sd::LongType const* tadShapeInfoZ,
    sd::LongType const* tadOffsetsZ) {

  scalarAlongDimension<X, Y, Z, OpType>
      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
          x,
          xShapeInfo,
          extraParams,
          z,
          zShapeInfo,
          scalars,
          dimension,
          dimensionLength,
          tadShapeInfo,
          tadOffsets,
          tadShapeInfoZ,
          tadOffsetsZ);
  sd::DebugHelper::checkErrorCode(stream, "scalarAlongDimension(...) failed");
}

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z>
void ScalarTransform<X, Y, Z>::executeCudaShaped(
    dim3& launchDims,
    cudaStream_t* stream,
    int opNum,
    void const* vx,
    sd::LongType const* xShapeInfo,
    sd::LongType const* hxShapeInfo,
    void* vz,
    sd::LongType const* zShapeInfo,
    sd::LongType const* hzShapeInfo,
    void const* vscalar,
    void* vextraParams) {

  DISPATCH_BY_OPNUM_TTT(
      intermediateShaped,
      PARAMS(launchDims, stream, vx, xShapeInfo, hxShapeInfo, vz, zShapeInfo, hzShapeInfo, vscalar, vextraParams, nullptr),
      SCALAR_OPS);
}


        ////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z>
void ScalarTransform<X, Y, Z>::executeCudaAlongDimension(
    dim3& launchDims,
    cudaStream_t* stream,
    int opNum,
    void const* vx,
    sd::LongType const* xShapeInfo,
    void* vz,
    sd::LongType const* zShapeInfo,
    void const* vscalars,
    void* vextraParams,
    sd::LongType* dimension,
    sd::LongType dimensionLength,
    sd::LongType const* tadShapeInfo,
    sd::LongType const* tadOffsets,
    sd::LongType const* tadShapeInfoZ,
    sd::LongType const* tadOffsetsZ) {

  DISPATCH_BY_OPNUM_TTT(
      intermediateAlongDimension,
      PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo, vscalars, vextraParams, dimension, dimensionLength,
             tadShapeInfo, tadOffsets, tadShapeInfoZ, tadOffsetsZ),
      SCALAR_OPS);
}

}  // namespace scalar
}  // namespace functions

#endif // SCALAR_CU