/* ******************************************************************************
 *
 *
 * 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
//
#include <execution/LaunchContext.h>
#include <exceptions/cuda_exception.h>
#include <system/op_boilerplate.h>
#include <loops/reduce_float.h>
#include <loops/scalar.h>
#include <loops/legacy_ops.h>
#include <helpers/DebugHelper.h>
#include <types/types.h>
#include <ops/specials_cuda.h>
#include <cuda.h>
#include <cuda_runtime.h>

    using namespace simdOps;

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z, typename OpType>
SD_KERNEL void simpleReduce(
    const void* x,
    const sd::LongType* outerXTadShapeInfo,
    const sd::LongType* innerXTadShapeInfo,
    void* extraParams,
    void* vreductionBuffer,
    void* z,
    const sd::LongType* zShapeInfo) {

  functions::reduce::ReduceFloatFunction<X,Z>::template transformCuda<OpType>(
      x,
      outerXTadShapeInfo,
      innerXTadShapeInfo,
      extraParams,
      vreductionBuffer,
      z,
      zShapeInfo);
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z, typename OpType>
SD_KERNEL void simpleScalar(
    const void* x,
    const sd::LongType* xShapeInfo,
    void* extraParams,
    void* z,
    const sd::LongType* zShapeInfo,
    sd::LongType* dimension,
    long long int dimensionLength,
    void* reductionBuffer,
    const sd::LongType* tadOnlyShapeInfo) {

  functions::reduce::ReduceFloatFunction<X, Z>::template execScalarCuda<OpType>(
      x,
      xShapeInfo,
      extraParams,
      z,
      zShapeInfo,
      reductionBuffer,
      tadOnlyShapeInfo);
}

namespace functions {
namespace reduce {

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
template <typename OpType>
SD_DEVICE void ReduceFloatFunction<X,Z>::aggregatePartials(
    void* vsPartials,
    sd::LongType tid,
    sd::LongType numItems,
    void* vextraParams) {

  using Y             = typename OpType::InterType;
  auto sPartials      = reinterpret_cast<Y*>(vsPartials);
  auto extraParams    = reinterpret_cast<Z*>(vextraParams);

  sd::LongType floorPow2 = numItems;

  if (floorPow2 & (floorPow2 - 1)) {
    while (floorPow2 & (floorPow2 - 1)) {
      floorPow2 &= (floorPow2 - 1);
    }
    if (tid >= floorPow2) {
      sPartials[tid - floorPow2] =
          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
    }
    __syncthreads();
  }

  for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
    if (tid < activeThreads && (tid + activeThreads) < numItems) {
      sPartials[tid] =
          OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
    }
    __syncthreads();
  }
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
template <typename OpType>
SD_DEVICE void ReduceFloatFunction<X,Z>::transformCuda(
    const void* vx,
    const sd::LongType* outerXTadShapeInfo,
    const sd::LongType* innerXTadShapeInfo,
    void* vextraParams,
    void* vreductionBuffer,
    void* vz,
    const sd::LongType* zShapeInfo) {

  auto x          = reinterpret_cast<const X*>(vx);
  auto z          = reinterpret_cast<Z*>(vz);
  auto extraParams= reinterpret_cast<Z*>(vextraParams);

  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
  __shared__ sd::LongType tadLen;
  __shared__ sd::LongType numTads;
  __shared__ bool sameOffsets;

  // Cache ranks/shape/stride for outer and z
  __shared__ sd::LongType outerRank;
  __shared__ const sd::LongType* outerStridePtr;
  __shared__ const sd::LongType* outerShapePtr;

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

  // Cache ranks/shape/stride for inner as well if needed
  __shared__ sd::LongType innerRank;
  __shared__ const sd::LongType* innerStridePtr;
  __shared__ const sd::LongType* innerShapePtr;

  if (threadIdx.x == 0) {
    outerRank      = shape::rank(outerXTadShapeInfo);
    outerShapePtr  = shape::shapeOf(outerXTadShapeInfo);
    outerStridePtr = shape::stride(outerXTadShapeInfo);

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

    innerRank      = shape::rank(innerXTadShapeInfo);
    innerShapePtr  = shape::shapeOf(innerXTadShapeInfo);
    innerStridePtr = shape::stride(innerXTadShapeInfo);

    sameOffsets    = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
    tadLen         = shape::length(innerXTadShapeInfo);
    numTads        = shape::length(outerXTadShapeInfo);
  }
  __syncthreads();

  sd::LongType coords[SD_MAX_RANK];
  sd::LongType zCoords[SD_MAX_RANK];
  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
    INDEX2COORDS(r, outerRank, outerShapePtr, coords);

    sd::LongType outerOffset;
    COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);

    sd::LongType zOffset;
    if (sameOffsets) {
      zOffset = outerOffset;
    } else {
      INDEX2COORDS(r, zRank, zShapePtr, zCoords);
      COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
    }

    const X* xTad = x + outerOffset;
    sPartials[threadIdx.x] = OpType::startingValue(xTad);

    // For the inner dimension
    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
      sd::LongType iCoords[SD_MAX_RANK];
      sd::LongType innerOffset;

      INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
      COORDS2INDEX(innerRank, innerStridePtr, iCoords, innerOffset);

      sPartials[threadIdx.x] = OpType::update(
          sPartials[threadIdx.x],
          OpType::op(xTad[innerOffset], extraParams),
          extraParams);
    }
    __syncthreads();

    aggregatePartials<OpType>(
        sPartials,
        threadIdx.x,
        sd::math::sd_min<int>(blockDim.x, tadLen),
        extraParams);
    __syncthreads();

    if (threadIdx.x == 0) {
      z[zOffset] =
          OpType::postProcess(sPartials[0], tadLen, extraParams);
    }
    __syncthreads();
  }
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
template <typename OpType>
SD_DEVICE void ReduceFloatFunction<X, Z>::execScalarCuda(
    const void* vx,
    const sd::LongType* xShapeInfo,
    void* vextraParams,
    void* vz,
    const sd::LongType* zShapeInfo,
    void* vreductionBuffer,
    const sd::LongType* tadOnlyShapeInfo) {

  auto x           = reinterpret_cast<const X*>(vx);
  auto z           = reinterpret_cast<Z*>(vz);
  auto extraParams = reinterpret_cast<Z*>(vextraParams);
  auto reductionBuffer = reinterpret_cast<Z*>(vreductionBuffer);

  using Y = typename OpType::InterType;

  __shared__ Y sPartials[SD_CUDA_BLOCK_SIZE];
  __shared__ sd::LongType length;

  // Cache rank/shape/stride
  __shared__ sd::LongType xRank;
  __shared__ const sd::LongType* xShapePtr;
  __shared__ const sd::LongType* xStridePtr;

  int tid = blockDim.x * blockIdx.x + threadIdx.x;
  if (threadIdx.x == 0) {
    length     = shape::length(xShapeInfo);
    xRank      = shape::rank(xShapeInfo);
    xShapePtr  = shape::shapeOf(xShapeInfo);
    xStridePtr = shape::stride(xShapeInfo);
  }
  __syncthreads();

  sPartials[threadIdx.x] = OpType::startingValue(x);

  sd::LongType gridSize = gridDim.x * blockDim.x;
  for (sd::LongType i = tid; i < length; i += gridSize) {
    sd::LongType xCoords[SD_MAX_RANK];
    sd::LongType xOffset;

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

    if(xOffset < length) {
      sPartials[threadIdx.x] = OpType::update(
          sPartials[threadIdx.x],
          OpType::op(x[xOffset], extraParams),
          extraParams);
    }

  }
  __syncthreads();

  aggregatePartials<OpType>(
      sPartials,
      threadIdx.x,
      sd::math::sd_min<int>(blockDim.x, length),
      extraParams);
  __syncthreads();

  if (gridDim.x > 1) {
    auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
    __shared__ bool amLast;

    if (threadIdx.x == 0) {
      reductionBuffer[blockIdx.x] = sPartials[0];
    }
    __threadfence();
    __syncthreads();

    if (threadIdx.x == 0) {
      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
      amLast              = (ticket == (gridDim.x - 1));
    }
    __syncthreads();

    if (amLast) {
      tc[16384]           = 0;
      sPartials[threadIdx.x] = OpType::startingValue(x);

      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
        sPartials[threadIdx.x] =
            OpType::update(sPartials[threadIdx.x],
                           reductionBuffer[i],
                           extraParams);
      }
      __syncthreads();

      aggregatePartials<OpType>(
          sPartials,
          threadIdx.x,
          sd::math::sd_min<int>(gridDim.x, blockDim.x),
          extraParams);
      __syncthreads();

      if (threadIdx.x == 0) {
        z[0] = OpType::postProcess(sPartials[0], length, extraParams);
      }
    }
  } else {
    if (threadIdx.x == 0) {
      auto tc   = reinterpret_cast<unsigned int*>(reductionBuffer);
      tc[16384] = 0;
      z[0]      = OpType::postProcess(sPartials[0], length, extraParams);
    }
  }
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
template<typename OpType>
SD_HOST void ReduceFloatFunction<X,Z>::intermediate(
    dim3 launchDims,
    cudaStream_t* stream,
    const void* x,
    const sd::LongType* dXShapeInfo,
    const sd::LongType* hXShapeInfo,
    void* extraParams,
    void* vreductionBuffer,
    void* z,
    const sd::LongType* dZShapeInfo,
    const sd::LongType* hZShapeInfo,
    const sd::LongType* dims) {

  if (shape::isEmptyConst(hXShapeInfo)) {
    const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
                                 ? sd::DataTypeUtils::nanOrZero<Z>()
                                 : static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));

    auto res = cudaMemcpyAsync(
        sd::LaunchContext::defaultContext()->getScalarPointer(),
        &startingVal,
        sizeof(Z),
        cudaMemcpyHostToDevice,
        *stream);
    if (res != 0) {
      throw sd::cuda_exception::build(
          "ReduceFloatFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
    }

    auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
    // scalar assign
    functions::scalar::ScalarTransform<Z, Z, Z>::executeCudaShaped(
        launchDims,
        stream,
        14,
        z,
        dZShapeInfo,
        hZShapeInfo,
        z,
        dZShapeInfo,
        hZShapeInfo,
        ptr,
        nullptr);
  } else {
    const int zRank   = shape::rank(hZShapeInfo);
    const int tadRank = shape::rank(hXShapeInfo) - zRank;

    auto outerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
        const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims), zRank);
    auto innerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
        const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims + zRank), tadRank);

    simpleReduce<X, Z, OpType>
        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
            x,
            reinterpret_cast<const sd::LongType*>(outerPack->special()),
            reinterpret_cast<const sd::LongType*>(innerPack->special()),
            extraParams,
            vreductionBuffer,
            z,
            dZShapeInfo);
  }
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
template<typename OpType>
SD_HOST void ReduceFloatFunction<X,Z>::intermediateScalar(
    dim3 launchDims,
    cudaStream_t* stream,
    const void* x,
    const sd::LongType* xShapeInfo,
    const sd::LongType* hXShapeInfo,
    void* extraParams,
    void* z,
    const sd::LongType* dZShapeInfo,
    const sd::LongType* hZShapeInfo,
    sd::LongType* dimension,
    sd::LongType dimensionLength,
    void* reductionBuffer,
    const sd::LongType* tadOnlyShapeInfo) {

  if (shape::isEmptyConst(hXShapeInfo)) {
    const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
                                 ? sd::DataTypeUtils::nanOrZero<Z>()
                                 : static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));

    auto res = cudaMemcpyAsync(z, &startingVal, sizeof(Z), cudaMemcpyHostToDevice, *stream);
    if (res != 0) {
      throw sd::cuda_exception::build(
          "ReduceFloatFunction<X,Z>::intermediateScalar: failed to copy resulting scalar", res);
    }
  } else {
    simpleScalar<X, Z, OpType><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
        x,
        xShapeInfo,
        extraParams,
        z,
        dZShapeInfo,
        dimension,
        dimensionLength,
        reductionBuffer,
        tadOnlyShapeInfo);
  }
  sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction intermediateScalar(...) failed");
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Y>
SD_HOST void ReduceFloatFunction<X,Y>::execReduceScalar(
    dim3 launchDims,
    cudaStream_t* stream,
    const int opNum,
    const void* x,
    const sd::LongType* xShapeInfo,
    const sd::LongType* hXShapeInfo,
    void* extraParams,
    void* z,
    const sd::LongType* dZShapeInfo,
    const sd::LongType* hZShapeInfo,
    sd::LongType* dimension,
    long long int dimensionLength,
    void* reductionBuffer,
    const sd::LongType* tadOnlyShapeInfo) {

  DISPATCH_BY_OPNUM_TT(
      intermediateScalar,
      PARAMS(
          launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z, dZShapeInfo, hZShapeInfo, dimension,
          dimensionLength, reductionBuffer, tadOnlyShapeInfo),
      OPS_A(REDUCE_FLOAT_OPS));
  sd::DebugHelper::checkErrorCode(stream, "execReduceScalarFloat(...) failed");
}

////////////////////////////////////////////////////////////////////////
template <typename X, typename Y>
SD_HOST void ReduceFloatFunction<X,Y>::execReduce(
    dim3 launchDims,
    cudaStream_t* stream,
    const int opNum,
    const void* x,
    const sd::LongType* dXShapeInfo,
    const sd::LongType* hXShapeInfo,
    void* extraParams,
    void* vreductionBuffer,
    void* z,
    const sd::LongType* dZShapeInfo,
    const sd::LongType* hZShapeInfo,
    const sd::LongType* dims) {

  if (shape::length(hZShapeInfo) == 1) {
    ReduceFloatFunction<X,Y>::execReduceScalar(
        launchDims, stream, opNum, x, dXShapeInfo, hXShapeInfo,
        extraParams, z, dZShapeInfo, hZShapeInfo,
        nullptr, 0, vreductionBuffer, nullptr);
  } else {
    DISPATCH_BY_OPNUM_TT(
        intermediate,
        PARAMS(
            launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams, vreductionBuffer, z, dZShapeInfo,
            hZShapeInfo, dims),
        OPS_A(REDUCE_FLOAT_OPS));
  }
  sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction execReduce(...) failed");
}

}  // namespace reduce
}  // namespace functions