/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

/**
 * @addtogroup surf2D surf2D
 * @{
 * @ingroup SurfaceTest
 */

#include <hip_array_common.hh>
#include <hip_test_common.hh>
#include <hip_texture_helper.hh>
#include "surf_common.h"

#pragma clang diagnostic ignored "-Wunused-variable"
#pragma clang diagnostic ignored "-Wunused-parameter"

#define LOG_DATA 0

template <typename T> __global__ void surf2DKernelR(hipSurfaceObject_t surfaceObject, T* outputData,
                                                    int width, int height) {
#if !__HIP_NO_IMAGE_SUPPORT
  int x = blockIdx.x * blockDim.x + threadIdx.x;
  int y = blockIdx.y * blockDim.y + threadIdx.y;
  if (x < width && y < height) {
    surf2Dread<T>(outputData + y * width + x, surfaceObject, x * sizeof(T), y);
  }
#endif
}

template <typename T> __global__ void surf2DKernelW(hipSurfaceObject_t surfaceObject, T* inputData,
                                                    int width, int height) {
#if !__HIP_NO_IMAGE_SUPPORT
  int x = blockIdx.x * blockDim.x + threadIdx.x;
  int y = blockIdx.y * blockDim.y + threadIdx.y;
  if (x < width && y < height) {
    surf2Dwrite<T>(inputData[y * width + x], surfaceObject, x * sizeof(T), y);
  }
#endif
}

template <typename T> __global__ void surf2DKernelRW(hipSurfaceObject_t surfaceObject,
                                                     hipSurfaceObject_t outputSurfObj, int width,
                                                     int height) {
#if !__HIP_NO_IMAGE_SUPPORT
  int x = blockIdx.x * blockDim.x + threadIdx.x;
  int y = blockIdx.y * blockDim.y + threadIdx.y;
  if (x < width && y < height) {
    T data;
    surf2Dread<T>(&data, surfaceObject, x * sizeof(T), y);
    surf2Dwrite<T>(data, outputSurfObj, x * sizeof(T), y);
  }
#endif
}

template <typename T> static void runTestR(const int width, const int height) {
  unsigned int size = width * height * sizeof(T);
  T* hData = (T*)malloc(size);
  memset(hData, 0, size);
  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      initVal(hData[i * width + j]);
    }
  }

  hipChannelFormatDesc channelDesc = hipCreateChannelDesc<T>();
  hipArray_t hipArray = nullptr;
  HIP_CHECK(hipMallocArray(&hipArray, &channelDesc, width, height, hipArraySurfaceLoadStore));

  // Need set source pitch, but we don't have any padding here
  const size_t spitch = width * sizeof(T);
  HIP_CHECK(
      hipMemcpy2DToArray(hipArray, 0, 0, hData, spitch, spitch, height, hipMemcpyHostToDevice));

  hipResourceDesc resDesc;
  memset(&resDesc, 0, sizeof(resDesc));
  resDesc.resType = hipResourceTypeArray;
  resDesc.res.array.array = hipArray;

  // Create surface object
  hipSurfaceObject_t surfaceObject = 0;
  HIP_CHECK(hipCreateSurfaceObject(&surfaceObject, &resDesc));

  T* hOutputData = nullptr;
  HIP_CHECK(hipHostMalloc((void**)&hOutputData, size));
  memset(hOutputData, 0, size);

  dim3 dimBlock(16, 16, 1);
  dim3 dimGrid((width + dimBlock.x - 1) / dimBlock.x, (height + dimBlock.y - 1) / dimBlock.y, 1);
  surf2DKernelR<T><<<dimGrid, dimBlock>>>(surfaceObject, hOutputData, width, height);
  HIP_CHECK(hipGetLastError());
  HIP_CHECK(hipDeviceSynchronize());

  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      int index = i * width + j;
      if (!isEqual(hData[index], hOutputData[index])) {
        printf("Difference [ %d %d ]:%s ----%s\n", i, j, getString(hData[index]).c_str(),
               getString(hOutputData[index]).c_str());
        REQUIRE(false);
      }
    }
  }

  HIP_CHECK(hipDestroySurfaceObject(surfaceObject));
  HIP_CHECK(hipFreeArray(hipArray));
  free(hData);
  HIP_CHECK(hipHostFree(hOutputData));
}

template <typename T> static void runTestW(const int width, const int height) {
  unsigned int size = width * height * sizeof(T);
  T* hData = nullptr;
  HIP_CHECK(hipHostMalloc((void**)&hData, size));
  memset(hData, 0, size);

  hipChannelFormatDesc channelDesc = hipCreateChannelDesc<T>();
  hipArray_t hipArray = nullptr;
  HIP_CHECK(hipMallocArray(&hipArray, &channelDesc, width, height, hipArraySurfaceLoadStore));

  // Need set source pitch, but we don't have any padding here
  const size_t spitch = width * sizeof(T);
  HIP_CHECK(
      hipMemcpy2DToArray(hipArray, 0, 0, hData, spitch, spitch, height, hipMemcpyHostToDevice));

  hipResourceDesc resDesc;
  memset(&resDesc, 0, sizeof(resDesc));
  resDesc.resType = hipResourceTypeArray;
  resDesc.res.array.array = hipArray;

  // Create surface object
  hipSurfaceObject_t surfaceObject = 0;
  HIP_CHECK(hipCreateSurfaceObject(&surfaceObject, &resDesc));

  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      initVal(hData[i * width + j]);
    }
  }

  dim3 dimBlock(16, 16, 1);
  dim3 dimGrid((width + dimBlock.x - 1) / dimBlock.x, (height + dimBlock.y - 1) / dimBlock.y, 1);
  surf2DKernelW<T><<<dimGrid, dimBlock>>>(surfaceObject, hData, width, height);
  HIP_CHECK(hipGetLastError());
  HIP_CHECK(hipDeviceSynchronize());

  T* hOutputData = (T*)malloc(size);

  memset(hOutputData, 0, size);
  HIP_CHECK(hipMemcpy2DFromArray(hOutputData, spitch, hipArray, 0, 0, spitch, height,
                                 hipMemcpyDeviceToHost));

  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      int index = i * width + j;
      if (!isEqual(hData[index], hOutputData[index])) {
        printf("Difference [ %d %d ]:%s ----%s\n", i, j, getString(hData[index]).c_str(),
               getString(hOutputData[index]).c_str());
        REQUIRE(false);
      }
    }
  }

  HIP_CHECK(hipDestroySurfaceObject(surfaceObject));
  HIP_CHECK(hipFreeArray(hipArray));
  HIP_CHECK(hipHostFree(hData));
  free(hOutputData);
}

template <typename T> static void runTestRW(const int width, const int height) {
  unsigned int size = width * height * sizeof(T);
  T* hData = (T*)malloc(size);
  memset(hData, 0, size);
  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      initVal(hData[i * width + j]);
    }
  }
#if LOG_DATA
  printf("hData: ");
  for (int i = 0; i < 32; i++) {
    printf("%s  ", getString(hData[i]).c_str());
  }
  printf("\n");
#endif

  hipChannelFormatDesc channelDesc = hipCreateChannelDesc<T>();
  hipArray_t hipArray = nullptr, hipOutArray = nullptr;
  HIP_CHECK(hipMallocArray(&hipArray, &channelDesc, width, height, hipArraySurfaceLoadStore));

  // Need set source pitch, but we don't have any padding here
  const size_t spitch = width * sizeof(T);
  HIP_CHECK(
      hipMemcpy2DToArray(hipArray, 0, 0, hData, spitch, spitch, height, hipMemcpyHostToDevice));

  hipResourceDesc resDesc;
  memset(&resDesc, 0, sizeof(resDesc));
  resDesc.resType = hipResourceTypeArray;
  resDesc.res.array.array = hipArray;

  // Create surface object
  hipSurfaceObject_t surfaceObject = 0;
  HIP_CHECK(hipCreateSurfaceObject(&surfaceObject, &resDesc));

  HIP_CHECK(hipMallocArray(&hipOutArray, &channelDesc, width, height, hipArraySurfaceLoadStore));

  hipResourceDesc resOutDesc;
  memset(&resOutDesc, 0, sizeof(resOutDesc));
  resOutDesc.resType = hipResourceTypeArray;
  resOutDesc.res.array.array = hipOutArray;

  hipSurfaceObject_t outSurfaceObject = 0;
  HIP_CHECK(hipCreateSurfaceObject(&outSurfaceObject, &resOutDesc));

  dim3 dimBlock(16, 16, 1);
  dim3 dimGrid((width + dimBlock.x - 1) / dimBlock.x, (height + dimBlock.y - 1) / dimBlock.y, 1);
  surf2DKernelRW<T><<<dimGrid, dimBlock>>>(surfaceObject, outSurfaceObject, width, height);
  HIP_CHECK(hipGetLastError());
  HIP_CHECK(hipDeviceSynchronize());

  T* hOutputData = (T*)malloc(size);

  memset(hOutputData, 0, size);
  HIP_CHECK(hipMemcpy2DFromArray(hOutputData, spitch, hipOutArray, 0, 0, spitch, height,
                                 hipMemcpyDeviceToHost));

#if LOG_DATA
  printf("dData: ");
  for (int i = 0; i < 32; i++) {
    printf("%s  ", getString(hOutputData[i]).c_str());
  }
  printf("\n");
#endif

  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      int index = i * width + j;
      if (!isEqual(hData[index], hOutputData[index])) {
        printf("Difference [ %d %d ]:%s ----%s\n", i, j, getString(hData[index]).c_str(),
               getString(hOutputData[index]).c_str());
        REQUIRE(false);
      }
    }
  }

  HIP_CHECK(hipDestroySurfaceObject(surfaceObject));
  HIP_CHECK(hipDestroySurfaceObject(outSurfaceObject));
  HIP_CHECK(hipFreeArray(hipArray));
  HIP_CHECK(hipFreeArray(hipOutArray));
  free(hData);
  free(hOutputData);
}

/**
 * Test Description
 * ------------------------
 *    - Basic test for `surf2Dread` with different types and dimensions.
 * Test source
 * ------------------------
 *    - unit/surface/surf2D.cc
 * Test requirements
 * ------------------------
 *    - HIP_VERSION >= 5.7
 */
HIP_TEMPLATE_TEST_CASE(Unit_surf2Dread_Positive_Basic, char, unsigned char, short, unsigned short, int,
                   unsigned int, float) {
  CHECK_IMAGE_SUPPORT;

  const int width = GENERATE(31, 67);
  const int height = GENERATE(131, 263);
  runTestR<vector_type_helper_t<TestType, 0>>(width, height);
  runTestR<vector_type_helper_t<TestType, 1>>(width, height);
  runTestR<vector_type_helper_t<TestType, 2>>(width, height);
  runTestR<vector_type_helper_t<TestType, 4>>(width, height);
}

/**
 * Test Description
 * ------------------------
 *    - Basic test for `surf2Dwrite` with different types and dimensions.
 * Test source
 * ------------------------
 *    - unit/surface/surf2D.cc
 * Test requirements
 * ------------------------
 *    - HIP_VERSION >= 5.7
 */
HIP_TEMPLATE_TEST_CASE(Unit_surf2Dwrite_Positive_Basic, char, unsigned char, short, unsigned short, int,
                   unsigned int, float) {
  CHECK_IMAGE_SUPPORT;

  const int width = GENERATE(31, 67);
  const int height = GENERATE(131, 263);
  runTestW<vector_type_helper_t<TestType, 0>>(width, height);
  runTestW<vector_type_helper_t<TestType, 1>>(width, height);
  runTestW<vector_type_helper_t<TestType, 2>>(width, height);
  runTestW<vector_type_helper_t<TestType, 4>>(width, height);
}

/**
 * Test Description
 * ------------------------
 *    - Basic test for `surf2Dread` and `surf2Dwrite` together, with different types and dimensions.
 * Test source
 * ------------------------
 *    - unit/surface/surf2D.cc
 * Test requirements
 * ------------------------
 *    - HIP_VERSION >= 5.7
 */
HIP_TEMPLATE_TEST_CASE(Unit_surf2D_Positive_ReadWrite, char, unsigned char, short, unsigned short, int,
                   unsigned int, float) {
  CHECK_IMAGE_SUPPORT;

  const int width = GENERATE(31, 67);
  const int height = GENERATE(131, 263);
  runTestRW<vector_type_helper_t<TestType, 0>>(width, height);
  runTestRW<vector_type_helper_t<TestType, 1>>(width, height);
  runTestRW<vector_type_helper_t<TestType, 2>>(width, height);
  runTestRW<vector_type_helper_t<TestType, 4>>(width, height);
}

/**
 * End doxygen group SurfaceTest.
 * @}
 */