/*
 * Copyright (c) 2021, NVIDIA CORPORATION.
 *
 * 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.
 */
#pragma once
#include <raft/cudart_utils.h>
#include <raft/cuda_utils.cuh>
#include <raft/linalg/contractions.cuh>
#include <raft/linalg/norm.cuh>
#include <raft/vectorized.cuh>

#include <cstddef>

namespace raft {
namespace distance {

/**
 * @brief Device class for L1, L2 and cosine distance metrics.
 * @tparam useNorms       whether norms are needed
 * @tparam DataT          input data-type (for A and B matrices)
 * @tparam AccT           accumulation data-type
 * @tparam OutT           output data-type (for C and D matrices)
 * @tparam IdxT           index data-type
 * @tparam Policy         struct which tunes the Contraction kernel
 * @tparam CoreLambda     tells how to accumulate an x and y into 
                          acc. its signature:
    template <typename AccT, typename DataT> void core_lambda(AccT& acc,
      const DataT& x, const DataT& y)
 * @tparam EpilogueLambda applies an elementwise function to compute final 
    values. Its signature is:
    template <typename AccT, typename DataT> void epilogue_lambda
    (AccT acc[][], DataT* regxn, DataT* regyn);
 * @tparam FinalLambda the final lambda called on final distance value
 * @param[in] x input matrix
 * @param[in] y input matrix
 * @param[in] m number of rows of A and C/D
 * @param[in] n number of columns of B and C/D
 * @param[in] k number of cols of A and rows of B
 * @param[in] lda leading dimension of A
 * @param[in] ldb leading dimension of B
 * @param[in] ldd leading dimension of C/D
 * @param[in] xn row norms of input matrix A. Required for expanded L2, cosine
 * @param[in] yn row norms of input matrix B. Required for expanded L2, cosine
 * @param[output] pD output matrix
 * @param[in] smem shared mem buffer for intermediate storage of A, B, xn & yn.
 * @param core_op the core accumulation operation lambda
 * @param epilog_op the epilog operation lambda
 * @param fin_op the final gemm epilogue lambda
 */

template <bool useNorms, typename DataT, typename AccT, typename OutT,
          typename IdxT, typename Policy, typename CoreLambda,
          typename EpilogueLambda, typename FinalLambda,
          typename rowEpilogueLambda, bool isRowMajor = true,
          bool writeOut = true,
          typename BaseClass =
            raft::linalg::Contractions_NT<DataT, IdxT, Policy, isRowMajor>>
struct PairwiseDistances : public BaseClass {
 private:
  typedef Policy P;
  const DataT* xn;
  const DataT* yn;
  const DataT* const yBase;
  OutT* dOutput;
  char* smem;
  CoreLambda core_op;
  EpilogueLambda epilog_op;
  FinalLambda fin_op;
  rowEpilogueLambda rowEpilog_op;

  AccT acc[P::AccRowsPerTh][P::AccColsPerTh];

 public:
  // Constructor
  DI PairwiseDistances(const DataT* _x, const DataT* _y, IdxT _m, IdxT _n,
                       IdxT _k, IdxT _lda, IdxT _ldb, IdxT _ldd,
                       const DataT* _xn, const DataT* _yn, OutT* _dOutput,
                       char* _smem, CoreLambda _core_op,
                       EpilogueLambda _epilog_op, FinalLambda _fin_op,
                       rowEpilogueLambda _rowEpilog_op)
    : BaseClass(_x, _y, _m, _n, _k, _lda, _ldb, _ldd, _smem),
      xn(_xn),
      yn(_yn),
      yBase(_y),
      dOutput(_dOutput),
      smem(_smem),
      core_op(_core_op),
      epilog_op(_epilog_op),
      fin_op(_fin_op),
      rowEpilog_op(_rowEpilog_op) {}

  DI void run() {
    for (auto gridStrideY = blockIdx.y * P::Mblk; gridStrideY < this->m;
         gridStrideY += P::Mblk * gridDim.y) {
      for (auto gridStrideX = blockIdx.x * P::Nblk; gridStrideX < this->n;
           gridStrideX += P::Nblk * gridDim.x) {
        prolog(gridStrideX, gridStrideY);
        loop();
        epilog(gridStrideX, gridStrideY);
      }
      rowEpilog_op(gridStrideY);
    }
  }

 private:
  DI void updateIndicesY() {
    const auto stride = P::Nblk * gridDim.x;
    if (isRowMajor) {
      this->y += stride * this->ldb;
    } else {
      this->y += stride;
    }
    this->yrowid += stride;
  }

  DI void updateIndicesXY() {
    const auto stride = P::Mblk * gridDim.y;
    if (isRowMajor) {
      this->x += stride * this->lda;
      this->yrowid = IdxT(blockIdx.x) * P::Nblk + this->srowid;
      this->y = yBase + this->yrowid * this->ldb;
    } else {
      this->x += stride;
      this->yrowid = IdxT(blockIdx.x) * P::Nblk;
      this->y = yBase + this->yrowid + this->srowid * this->ldb;
    }
    this->xrowid += stride;
  }

  DI void ldgNextGridStride(IdxT gridStrideX, IdxT gridStrideY) {
    // Fetch next grid stride ldg if within range
    if ((gridStrideX + gridDim.x * P::Nblk) < this->n) {
      updateIndicesY();
      this->ldgXY(0);
    } else if ((gridStrideY + gridDim.y * P::Mblk) < this->m) {
      updateIndicesXY();
      this->ldgXY(0);
    }
  }

  DI void prolog(IdxT gridStrideX, IdxT gridStrideY) {
    if (gridStrideX == blockIdx.x * P::Nblk) {
      this->ldgXY(0);
    }

#pragma unroll
    for (int i = 0; i < P::AccRowsPerTh; ++i) {
#pragma unroll
      for (int j = 0; j < P::AccColsPerTh; ++j) {
        acc[i][j] = BaseClass::Zero;
      }
    }

    this->stsXY();
    __syncthreads();
    this->pageWr ^= 1;
  }

  DI void loop() {
    for (int kidx = P::Kblk; kidx < this->k; kidx += P::Kblk) {
      this->ldgXY(kidx);
      accumulate();  // on the previous k-block
      this->stsXY();
      __syncthreads();
      this->pageWr ^= 1;
      this->pageRd ^= 1;
    }
    accumulate();  // last iteration
    // This is needed for making sure next grid stride of
    // non-norm based metrics uses previously accumulated buffer so
    // it doesn't make shmem dirty until previous iteration
    // is complete.
    this->pageRd ^= 1;
  }

  DI void accumulate() {
#pragma unroll
    for (int ki = 0; ki < P::Kblk; ki += P::Veclen) {
      this->ldsXY(ki);
#pragma unroll
      for (int i = 0; i < P::AccRowsPerTh; ++i) {
#pragma unroll
        for (int j = 0; j < P::AccColsPerTh; ++j) {
#pragma unroll
          for (int v = 0; v < P::Veclen; ++v) {
            core_op(acc[i][j], this->regx[i][v], this->regy[j][v]);
          }
        }
      }
    }
  }

  DI void epilog(IdxT gridStrideX, IdxT gridStrideY) {
    if (useNorms) {
      DataT* sxNorm = (DataT*)(&smem[P::SmemSize]);
      DataT* syNorm = (&sxNorm[P::Mblk]);

      // Load x & y norms required by this threadblock in shmem buffer
      if (gridStrideX == blockIdx.x * P::Nblk) {
        for (int i = threadIdx.x; i < P::Mblk; i += P::Nthreads) {
          auto idx = gridStrideY + i;
          sxNorm[i] = idx < this->m ? xn[idx] : 0;
        }
      }

      for (int i = threadIdx.x; i < P::Nblk; i += P::Nthreads) {
        auto idx = gridStrideX + i;
        syNorm[i] = idx < this->n ? yn[idx] : 0;
      }

      __syncthreads();

      DataT regxn[P::AccRowsPerTh], regyn[P::AccColsPerTh];
#pragma unroll
      for (int i = 0; i < P::AccRowsPerTh; ++i) {
        regxn[i] = sxNorm[i * P::AccThRows + (threadIdx.x / P::AccThCols)];
      }
#pragma unroll
      for (int i = 0; i < P::AccColsPerTh; ++i) {
        regyn[i] = syNorm[i * P::AccThCols + (threadIdx.x % P::AccThCols)];
      }

      // Overlap ldg with epilog computation
      ldgNextGridStride(gridStrideX, gridStrideY);
      epilog_op(acc, regxn, regyn, gridStrideX, gridStrideY);
    } else {
      // Overlap ldg with epilog computation
      ldgNextGridStride(gridStrideX, gridStrideY);
      epilog_op(acc, nullptr, nullptr, gridStrideX, gridStrideY);
    }

    if (writeOut) {
      IdxT starty = gridStrideY + this->accrowid;
      IdxT startx = gridStrideX + this->acccolid;

#pragma unroll
      for (int i = 0; i < P::AccRowsPerTh; ++i) {
        auto rowId = starty + i * P::AccThRows;
#pragma unroll
        for (int j = 0; j < P::AccColsPerTh; ++j) {
          auto colId = startx + j * P::AccThCols;
          if (rowId < this->m && colId < this->n) {
            dOutput[rowId * this->n + colId] = fin_op(acc[i][j], 0);
          }
        }
      }
    }
  }
};  // struct PairwiseDistances

/**
 * @brief the distance matrix calculation kernel for L1, L2 and cosine
 * @tparam useNorms       whether norms are needed
 * @tparam DataT          input data-type (for A and B matrices)
 * @tparam AccT           accumulation data-type
 * @tparam OutT           output data-type (for C and D matrices)
 * @tparam IdxT           index data-type
 * @tparam Policy         struct which tunes the Contraction kernel
 * @tparam CoreLambda     lambda which implements accumulation operation
 * @tparam EpilogueLambda lambda which implements operation for calculating
                          final value.
 * @tparam FinalLambda    final lambda called on final distance value
 * @tparam isRowMajor     true if input/output is row major(default),
                          false for column major
 *
 * @param[in]       x input matrix
 * @param[in]       y input matrix
 * @param[in]       xn row norms of input matrix A.
 * @param[in]       yn row norms of input matrix B.
 * @param[in]       m number of rows of A and C/D
 * @param[in]       n number of columns of B and C/D
 * @param[in]       k number of cols of A and rows of B
 * @param[in]       lda leading dimension of A
 * @param[in]       ldb leading dimension of B
 * @param[in]       ldd leading dimension of C/D
 * @param[output]   pD output matrix
 * @param core_op   the core lambda
 * @param epilog_op the epilogue lambda
 * @param fin_op    the final gemm epilogue lambda
 */

template <bool useNorms, typename DataT, typename AccT, typename OutT,
          typename IdxT, typename Policy, typename CoreLambda,
          typename EpilogueLambda, typename FinalLambda, bool isRowMajor = true,
          bool writeOut = true>
__global__ __launch_bounds__(
  Policy::Nthreads,
  2) void pairwiseDistanceMatKernel(const DataT* x, const DataT* y,
                                    const DataT* _xn, const DataT* _yn, IdxT m,
                                    IdxT n, IdxT k, IdxT lda, IdxT ldb,
                                    IdxT ldd, OutT* dOutput, CoreLambda core_op,
                                    EpilogueLambda epilog_op,
                                    FinalLambda fin_op) {
  extern __shared__ char smem[];
  auto rowEpilog = [] __device__(IdxT starty) { return; };

  PairwiseDistances<useNorms, DataT, AccT, OutT, IdxT, Policy, CoreLambda,
                    EpilogueLambda, FinalLambda, decltype(rowEpilog),
                    isRowMajor, writeOut>
    obj(x, y, m, n, k, lda, ldb, ldd, _xn, _yn, dOutput, smem, core_op,
        epilog_op, fin_op, rowEpilog);
  obj.run();
}

template <typename P, typename IdxT, typename T>
dim3 launchConfigGenerator(IdxT m, IdxT n, std::size_t sMemSize, T func) {
  const auto numSMs = raft::getMultiProcessorCount();
  int numBlocksPerSm = 0;
  dim3 grid;

  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
    &numBlocksPerSm, func, P::Nthreads, sMemSize));
  std::size_t minGridSize = numSMs * numBlocksPerSm;
  std::size_t yChunks = raft::ceildiv<int>(m, P::Mblk);
  std::size_t xChunks = raft::ceildiv<int>(n, P::Nblk);
  grid.y = yChunks > minGridSize ? minGridSize : yChunks;
  grid.x = (minGridSize - grid.y) <= 0 ? 1 : xChunks;
  if (grid.x != 1) {
    std::size_t i = 1;
    while (grid.y * i < minGridSize) {
      i++;
    }
    grid.x = i >= xChunks ? xChunks : i;
  }

  return grid;
}

};  // namespace distance
};  // namespace raft
