static char help[] = "Tests MatCreateDenseCUDA(), MatDenseCUDAPlaceArray(), MatDenseCUDAReplaceArray(), and MatDenseCUDAResetArray()\n\
  or MatCreateDenseHIP(), MatDenseHIPPlaceArray(), MatDenseHIPReplaceArray(), and MatDenseHIPResetArray(),\n\
  as well as MatDiagonalScale() on device.\n";

#include <petscmat.h>
#include <petscpkg_version.h>

#if !defined(PETSC_PKG_CUDA_VERSION_GE)
  #define PETSC_PKG_CUDA_VERSION_GE(MAJ, MIN, PAT) 0
#endif
#define PetscConcat3(a, b, c) PetscConcat(PetscConcat(a, b), c)
#if PetscDefined(HAVE_CUDA)
  #define DEVICE_TAG CUDA
#elif PetscDefined(HAVE_HIP)
  #define DEVICE_TAG HIP
#else
  #define DEVICE_TAG
#endif
#define VecDeviceGetArray          PetscConcat3(Vec, DEVICE_TAG, GetArray)
#define VecDeviceRestoreArray      PetscConcat3(Vec, DEVICE_TAG, RestoreArray)
#define MatCreateDenseDevice       PetscConcat(MatCreateDense, DEVICE_TAG)
#define MatDenseDeviceReplaceArray PetscConcat3(MatDense, DEVICE_TAG, ReplaceArray)
#define MatDenseDevicePlaceArray   PetscConcat3(MatDense, DEVICE_TAG, PlaceArray)
#define MatDenseDeviceResetArray   PetscConcat3(MatDense, DEVICE_TAG, ResetArray)

static PetscErrorCode MatMult_S(Mat S, Vec x, Vec y)
{
  Mat A;

  PetscFunctionBeginUser;
  PetscCall(MatShellGetContext(S, &A));
  PetscCall(MatMult(A, x, y));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscBool test_cusparse_transgen = PETSC_FALSE;

static PetscErrorCode MatMultTranspose_S(Mat S, Vec x, Vec y)
{
  Mat A;

  PetscFunctionBeginUser;
  PetscCall(MatShellGetContext(S, &A));
  PetscCall(MatMultTranspose(A, x, y));

  /* alternate transgen true and false to test code logic */
  PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, test_cusparse_transgen));
  test_cusparse_transgen = (PetscBool)!test_cusparse_transgen;
  PetscFunctionReturn(PETSC_SUCCESS);
}

int main(int argc, char **argv)
{
  Mat          A, B, C, S;
  Vec          t, v, ll, rr;
  PetscScalar *vv, *aa;
  // We met a mysterious cudaErrorMisalignedAddress error on some systems with cuda-12.0,1 but not
  // with prior cuda-11.2,3,7,8 versions. Making nloc an even number somehow 'fixes' the problem.
  // See more at https://gitlab.com/petsc/petsc/-/merge_requests/6225
#if PETSC_PKG_CUDA_VERSION_GE(12, 0, 0)
  PetscInt n = 32;
#else
  PetscInt n = 30;
#endif
  PetscInt  k = 6, l = 0, i, Istart, Iend, nloc, bs, test = 1;
  PetscBool flg, reset, use_shell = PETSC_FALSE;
  VecType   vtype;

  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc, &argv, NULL, help));
  PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
  PetscCall(PetscOptionsGetInt(NULL, NULL, "-k", &k, NULL));
  PetscCall(PetscOptionsGetInt(NULL, NULL, "-l", &l, NULL));
  PetscCall(PetscOptionsGetInt(NULL, NULL, "-test", &test, NULL));
  PetscCall(PetscOptionsGetBool(NULL, NULL, "-use_shell", &use_shell, NULL));
  PetscCheck(k >= 0, PETSC_COMM_WORLD, PETSC_ERR_USER, "k %" PetscInt_FMT " must be positive", k);
  PetscCheck(l >= 0, PETSC_COMM_WORLD, PETSC_ERR_USER, "l %" PetscInt_FMT " must be positive", l);
  PetscCheck(l <= k, PETSC_COMM_WORLD, PETSC_ERR_USER, "l %" PetscInt_FMT " must be smaller or equal than k %" PetscInt_FMT, l, k);

  /* sparse matrix */
  PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
  PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, n, n));
  PetscCall(MatSetType(A, PetscDefined(HAVE_CUDA) ? MATAIJCUSPARSE : MATAIJHIPSPARSE));
  PetscCall(MatSetOptionsPrefix(A, "A_"));
  PetscCall(MatSetFromOptions(A));
  PetscCall(MatSetUp(A));

  /* test special case for MATSEQAIJCUSPARSE and MATSEQAIJHIPSPARSE to generate explicit transpose (not default) */
  PetscCall(MatSetOption(A, MAT_FORM_EXPLICIT_TRANSPOSE, PETSC_TRUE));

  PetscCall(MatGetOwnershipRange(A, &Istart, &Iend));
  for (i = Istart; i < Iend; i++) {
    if (i > 0) PetscCall(MatSetValue(A, i, i - 1, -1.0, INSERT_VALUES));
    if (i < n - 1) PetscCall(MatSetValue(A, i, i + 1, -1.0, INSERT_VALUES));
    PetscCall(MatSetValue(A, i, i, 2.0, INSERT_VALUES));
  }
  PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));

  /* template vector */
  PetscCall(MatCreateVecs(A, NULL, &t));
  PetscCall(VecGetType(t, &vtype));

  /* long vector, contains the stacked columns of an nxk dense matrix */
  PetscCall(VecGetLocalSize(t, &nloc));
  PetscCall(VecGetBlockSize(t, &bs));
  PetscCall(VecCreate(PetscObjectComm((PetscObject)t), &v));
  PetscCall(VecSetType(v, vtype));
  PetscCall(VecSetSizes(v, k * nloc, k * n));
  PetscCall(VecSetBlockSize(v, bs));
  PetscCall(VecSetRandom(v, NULL));

  /* dense matrix that contains the columns of v */
  PetscCall(VecDeviceGetArray(v, &vv));

  /* test few cases for MATDENSECUDA and MATDENSEHIP handling pointers */
  switch (test) {
  case 1:
    PetscCall(MatCreateDenseDevice(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, vv, &B)); /* pass a pointer to avoid allocation of storage */
    PetscCall(MatDenseDeviceReplaceArray(B, NULL));                                                         /* replace with a null pointer, the value after BVRestoreMat */
    PetscCall(MatDenseDevicePlaceArray(B, vv + l * nloc));                                                  /* set the actual pointer */
    reset = PETSC_TRUE;
    break;
  case 2:
    PetscCall(MatCreateDenseDevice(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &B));
    PetscCall(MatDenseDevicePlaceArray(B, vv + l * nloc)); /* set the actual pointer */
    reset = PETSC_TRUE;
    break;
  default:
    PetscCall(MatCreateDenseDevice(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, vv + l * nloc, &B));
    reset = PETSC_FALSE;
    break;
  }

  /* test MatMatMult() */
  if (use_shell) {
    /* we could have called the general converter below, but we explicitly set the operations
       ourselves to test MatProductSymbolic_X_Dense, MatProductNumeric_X_Dense code */
    /* PetscCall(MatConvert(A, MATSHELL, MAT_INITIAL_MATRIX, &S)); */
    PetscCall(MatCreateShell(PetscObjectComm((PetscObject)v), nloc, nloc, n, n, A, &S));
    PetscCall(MatShellSetOperation(S, MATOP_MULT, (PetscErrorCodeFn *)MatMult_S));
    PetscCall(MatShellSetOperation(S, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_S));
    PetscCall(MatShellSetVecType(S, vtype));
  } else {
    PetscCall(PetscObjectReference((PetscObject)A));
    S = A;
  }

  PetscCall(MatCreateDenseDevice(PetscObjectComm((PetscObject)v), nloc, PETSC_DECIDE, n, k - l, NULL, &C));

  /* test MatMatMult() */
  PetscCall(MatProductCreateWithMat(S, B, NULL, C));
  PetscCall(MatProductSetType(C, MATPRODUCT_AB));
  PetscCall(MatProductSetFromOptions(C));
  PetscCall(MatProductSymbolic(C));
  PetscCall(MatProductNumeric(C));
  PetscCall(MatMatMultEqual(S, B, C, 10, &flg));
  if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatMatMult()\n"));

  /* test MatTransposeMatMult() */
  PetscCall(MatProductCreateWithMat(S, B, NULL, C));
  PetscCall(MatProductSetType(C, MATPRODUCT_AtB));
  PetscCall(MatProductSetFromOptions(C));
  PetscCall(MatProductSymbolic(C));
  PetscCall(MatProductNumeric(C));
  PetscCall(MatTransposeMatMultEqual(S, B, C, 10, &flg));
  if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatTransposeMatMult()\n"));

  PetscCall(MatDestroy(&C));
  PetscCall(MatDestroy(&S));

  /* test MatDiagonalScale() */
  PetscCall(MatDuplicate(B, MAT_COPY_VALUES, &C));
  PetscCall(MatBindToCPU(C, PETSC_TRUE));
  PetscCall(MatCreateVecs(B, &rr, &ll));
  PetscCall(VecSetRandom(ll, NULL));
  PetscCall(VecSetRandom(rr, NULL));
  PetscCall(MatDiagonalScale(C, ll, rr));
  PetscCall(MatDiagonalScale(B, ll, rr));
  PetscCall(MatMultEqual(B, C, 10, &flg));
  if (!flg) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error MatDiagonalScale()\n"));

  PetscCall(VecDestroy(&ll));
  PetscCall(VecDestroy(&rr));
  PetscCall(MatDestroy(&C));

  /* finished using B */
  PetscCall(MatDenseGetArrayAndMemType(B, &aa, NULL));
  PetscCheck(vv == aa - l * nloc, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong array");
  PetscCall(MatDenseRestoreArrayAndMemType(B, &aa));
  if (reset) PetscCall(MatDenseDeviceResetArray(B));
  PetscCall(VecDeviceRestoreArray(v, &vv));

  if (test == 1) {
    PetscCall(MatDenseGetArrayAndMemType(B, &aa, NULL));
    PetscCheck(!aa, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Expected a null pointer");
    PetscCall(MatDenseRestoreArrayAndMemType(B, &aa));
  }

  /* free work space */
  PetscCall(MatDestroy(&B));
  PetscCall(MatDestroy(&A));
  PetscCall(VecDestroy(&t));
  PetscCall(VecDestroy(&v));
  PetscCall(PetscFinalize());
  return 0;
}

/*TEST

  build:
    requires: defined(PETSC_DEVICELANGUAGE_CXX)

  testset:
    nsize: {{1 2}}
    args: -test {{0 1 2}} -k 6 -l {{0 5}} -use_shell {{0 1}}
    output_file: output/empty.out
    test:
      requires: cuda
      suffix: 1_cuda
      args: -A_mat_type {{aij aijcusparse}}
    test:
      requires: hip
      suffix: 1_hip
      args: -A_mat_type aijhipsparse

TEST*/
