/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/**
**************************************************************************
* \file dct8x8_kernel1.cu
* \brief Contains 1st CUDA implementations of DCT, IDCT and quantization
*routines,
*        used in JPEG internal data processing. Device code.
*
* This code implements first CUDA versions of forward and inverse Discrete
*Cosine
* Transform to blocks of image pixels (of 8x8 size), as in JPEG standard. The
*data
* processing is done using floating point representation.
* The routine that performs quantization of coefficients can be found in
* dct8x8_kernel_quantization.cu file.
*/
#pragma once
#include <cooperative_groups.h>

namespace cg = cooperative_groups;
#include "Common.h"

/**
*  This unitary matrix performs discrete cosine transform of rows of the matrix
* to the left
*/
__constant__ float DCTv8matrix[] = {
  0.3535533905932738f,  0.4903926402016152f,  0.4619397662556434f,  0.4157348061512726f,  0.3535533905932738f,  0.2777851165098011f,  0.1913417161825449f,  0.0975451610080642f,
  0.3535533905932738f,  0.4157348061512726f,  0.1913417161825449f, -0.0975451610080641f, -0.3535533905932737f, -0.4903926402016152f, -0.4619397662556434f, -0.2777851165098011f,
  0.3535533905932738f,  0.2777851165098011f, -0.1913417161825449f, -0.4903926402016152f, -0.3535533905932738f,  0.0975451610080642f,  0.4619397662556433f,  0.4157348061512727f,
  0.3535533905932738f,  0.0975451610080642f, -0.4619397662556434f, -0.2777851165098011f,  0.3535533905932737f,  0.4157348061512727f, -0.1913417161825450f, -0.4903926402016153f,
  0.3535533905932738f, -0.0975451610080641f, -0.4619397662556434f,  0.2777851165098009f,  0.3535533905932738f, -0.4157348061512726f, -0.1913417161825453f,  0.4903926402016152f,
  0.3535533905932738f, -0.2777851165098010f, -0.1913417161825452f,  0.4903926402016153f, -0.3535533905932733f, -0.0975451610080649f,  0.4619397662556437f, -0.4157348061512720f,
  0.3535533905932738f, -0.4157348061512727f,  0.1913417161825450f,  0.0975451610080640f, -0.3535533905932736f,  0.4903926402016152f, -0.4619397662556435f,  0.2777851165098022f,
  0.3535533905932738f, -0.4903926402016152f,  0.4619397662556433f, -0.4157348061512721f,  0.3535533905932733f, -0.2777851165098008f,  0.1913417161825431f, -0.0975451610080625f
};

// Temporary blocks
__shared__ float CurBlockLocal1[BLOCK_SIZE2];
__shared__ float CurBlockLocal2[BLOCK_SIZE2];

/**
**************************************************************************
*  Performs 1st implementation of 8x8 block-wise Forward Discrete Cosine
*Transform of the given
*  image plane and outputs result to the array of coefficients.
*
* \param Dst            [OUT] - Coefficients plane
* \param ImgWidth       [IN] - Stride of Dst
* \param OffsetXBlocks  [IN] - Offset along X in blocks from which to perform
*processing
* \param OffsetYBlocks  [IN] - Offset along Y in blocks from which to perform
*processing
*
* \return None
*/
__global__ void CUDAkernel1DCT(float *Dst, int ImgWidth, int OffsetXBlocks,
                               int OffsetYBlocks, cudaTextureObject_t TexSrc) {
  // Handle to thread block group
  cg::thread_block cta = cg::this_thread_block();
  // Block index
  const int bx = blockIdx.x + OffsetXBlocks;
  const int by = blockIdx.y + OffsetYBlocks;

  // Thread index (current coefficient)
  const int tx = threadIdx.x;
  const int ty = threadIdx.y;

  // Texture coordinates
  const float tex_x = (float)((bx << BLOCK_SIZE_LOG2) + tx) + 0.5f;
  const float tex_y = (float)((by << BLOCK_SIZE_LOG2) + ty) + 0.5f;

  // copy current image pixel to the first block
  CurBlockLocal1[(ty << BLOCK_SIZE_LOG2) + tx] =
      tex2D<float>(TexSrc, tex_x, tex_y);

  // synchronize threads to make sure the block is copied
  cg::sync(cta);

  // calculate the multiplication of DCTv8matrixT * A and place it in the second
  // block
  float curelem = 0;
  int DCTv8matrixIndex = 0 * BLOCK_SIZE + ty;
  int CurBlockLocal1Index = 0 * BLOCK_SIZE + tx;
#pragma unroll

  for (int i = 0; i < BLOCK_SIZE; i++) {
    curelem +=
        DCTv8matrix[DCTv8matrixIndex] * CurBlockLocal1[CurBlockLocal1Index];
    DCTv8matrixIndex += BLOCK_SIZE;
    CurBlockLocal1Index += BLOCK_SIZE;
  }

  CurBlockLocal2[(ty << BLOCK_SIZE_LOG2) + tx] = curelem;

  // synchronize threads to make sure the first 2 matrices are multiplied and
  // the result is stored in the second block
  cg::sync(cta);

  // calculate the multiplication of (DCTv8matrixT * A) * DCTv8matrix and place
  // it in the first block
  curelem = 0;
  int CurBlockLocal2Index = (ty << BLOCK_SIZE_LOG2) + 0;
  DCTv8matrixIndex = 0 * BLOCK_SIZE + tx;
#pragma unroll

  for (int i = 0; i < BLOCK_SIZE; i++) {
    curelem +=
        CurBlockLocal2[CurBlockLocal2Index] * DCTv8matrix[DCTv8matrixIndex];
    CurBlockLocal2Index += 1;
    DCTv8matrixIndex += BLOCK_SIZE;
  }

  CurBlockLocal1[(ty << BLOCK_SIZE_LOG2) + tx] = curelem;

  // synchronize threads to make sure the matrices are multiplied and the result
  // is stored back in the first block
  cg::sync(cta);

  // copy current coefficient to its place in the result array
  Dst[FMUL(((by << BLOCK_SIZE_LOG2) + ty), ImgWidth) +
      ((bx << BLOCK_SIZE_LOG2) + tx)] =
      CurBlockLocal1[(ty << BLOCK_SIZE_LOG2) + tx];
}

/**
**************************************************************************
*  Performs 1st implementation of 8x8 block-wise Inverse Discrete Cosine
*Transform of the given
*  DCT coefficients plane and outputs result to the image array
*
* \param Dst            [OUT] - Image plane
* \param ImgWidth       [IN] - Stride of Dst
* \param OffsetXBlocks  [IN] - Offset along X in blocks from which to perform
*processing
* \param OffsetYBlocks  [IN] - Offset along Y in blocks from which to perform
*processing
*
* \return None
*/
__global__ void CUDAkernel1IDCT(float *Dst, int ImgWidth, int OffsetXBlocks,
                                int OffsetYBlocks, cudaTextureObject_t TexSrc) {
  // Handle to thread block group
  cg::thread_block cta = cg::this_thread_block();
  // Block index
  int bx = blockIdx.x + OffsetXBlocks;
  int by = blockIdx.y + OffsetYBlocks;

  // Thread index (current image pixel)
  int tx = threadIdx.x;
  int ty = threadIdx.y;

  // Texture coordinates
  const float tex_x = (float)((bx << BLOCK_SIZE_LOG2) + tx) + 0.5f;
  const float tex_y = (float)((by << BLOCK_SIZE_LOG2) + ty) + 0.5f;

  // copy current image pixel to the first block
  CurBlockLocal1[(ty << BLOCK_SIZE_LOG2) + tx] =
      tex2D<float>(TexSrc, tex_x, tex_y);

  // synchronize threads to make sure the block is copied
  cg::sync(cta);

  // calculate the multiplication of DCTv8matrix * A and place it in the second
  // block
  float curelem = 0;
  int DCTv8matrixIndex = (ty << BLOCK_SIZE_LOG2) + 0;
  int CurBlockLocal1Index = 0 * BLOCK_SIZE + tx;
#pragma unroll

  for (int i = 0; i < BLOCK_SIZE; i++) {
    curelem +=
        DCTv8matrix[DCTv8matrixIndex] * CurBlockLocal1[CurBlockLocal1Index];
    DCTv8matrixIndex += 1;
    CurBlockLocal1Index += BLOCK_SIZE;
  }

  CurBlockLocal2[(ty << BLOCK_SIZE_LOG2) + tx] = curelem;

  // synchronize threads to make sure the first 2 matrices are multiplied and
  // the result is stored in the second block
  cg::sync(cta);

  // calculate the multiplication of (DCTv8matrix * A) * DCTv8matrixT and place
  // it in the first block
  curelem = 0;
  int CurBlockLocal2Index = (ty << BLOCK_SIZE_LOG2) + 0;
  DCTv8matrixIndex = (tx << BLOCK_SIZE_LOG2) + 0;
#pragma unroll

  for (int i = 0; i < BLOCK_SIZE; i++) {
    curelem +=
        CurBlockLocal2[CurBlockLocal2Index] * DCTv8matrix[DCTv8matrixIndex];
    CurBlockLocal2Index += 1;
    DCTv8matrixIndex += 1;
  }

  CurBlockLocal1[(ty << BLOCK_SIZE_LOG2) + tx] = curelem;

  // synchronize threads to make sure the matrices are multiplied and the result
  // is stored back in the first block
  cg::sync(cta);

  // copy current coefficient to its place in the result array
  Dst[FMUL(((by << BLOCK_SIZE_LOG2) + ty), ImgWidth) +
      ((bx << BLOCK_SIZE_LOG2) + tx)] =
      CurBlockLocal1[(ty << BLOCK_SIZE_LOG2) + tx];
}