/*
 * ----------------------------------------------------------------------------
 * Programmer(s): Cody J. Balos @ LLNL
 * ----------------------------------------------------------------------------
 * SUNDIALS Copyright Start
 * Copyright (c) 2002-2024, Lawrence Livermore National Security
 * and Southern Methodist University.
 * All rights reserved.
 *
 * See the top-level LICENSE and NOTICE files for details.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 * SUNDIALS Copyright End
 * ----------------------------------------------------------------------------
 * This is the header file for the SUNMatrix that wraps the SuperLU-DIST
 * SLU_NR_loc type SuperMatrix.
 * ----------------------------------------------------------------------------
 */

#include <cmath>
#include <mpi.h>
#include <nvector/nvector_parallel.h>
#include <nvector/nvector_serial.h>
#include <stdio.h>
#include <stdlib.h>
#include <sundials/sundials_math.h>
#include <sundials/sundials_types.h>
#include <sunmatrix/sunmatrix_dense.h>
#include <sunmatrix/sunmatrix_slunrloc.h>
#include <superlu_ddefs.h>

#include "test_sunmatrix.h"

/* Control whethere to print on rank 0 only or all ranks
 * when using TEST_STATUS* macros */
static int print_all_ranks;
static int print_timing;

/* helper functions */
int csr_from_dense(SUNMatrix Ad, sunrealtype droptol, sunrealtype** matdata,
                   sunindextype** colind, sunindextype** rowptrs);

/* ----------------------------------------------------------------------------
 * Main SUNMatrix Testing Routine
 * --------------------------------------------------------------------------*/
int main(int argc, char* argv[])
{
  int fails     = 0;              /* counter for test failures        */
  int globfails = 0;              /* counter for test failures        */
  int nprocs, nprow, npcol, rank; /* process grid size and rank       */
  MPI_Status mpistatus;           /* MPI status                       */
  MPI_Comm comm;                  /* MPI communicator                 */
  gridinfo_t grid;                /* SuperLU-DIST process grid        */

  sunindextype M, N;              /* matrix size                      */
  sunrealtype* matdata;           /* pointer to matrix data           */
  sunindextype *rowptrs, *colind; /* indexptrs and indexvals          */
  SUNMatrix D, A, I;              /* global and local matrices        */
  SuperMatrix *Asuper, *Isuper;   /* SLU_NR_loc SuperMatrices         */
  NRformat_loc* Istore;           /* SuperMatrix->Store               */
  N_Vector gx, gy, x, y;          /* test vectors                     */
  sunrealtype *xdata, *ydata;     /* vector data                      */
  sunindextype M_local;           /* num rows in local matrix         */
  sunindextype Mnprocs;           /* M/nprocs                         */
  sunindextype Mrem;              /* M%nprocs                         */
  sunindextype NNZ_local;         /* num of nnz in local matrix       */
  sunindextype fst_row;           /* global. index of 1st local row   */
  int square;                     /* is A a square matrix             */
  sunindextype i, j, k;           /* just some iteration variables    */
  SUNContext sunctx;

  MPI_Init(&argc, &argv);

  comm = MPI_COMM_WORLD;
  MPI_Comm_size(comm, &nprocs);
  MPI_Comm_rank(comm, &rank);

  if (SUNContext_Create(comm, &sunctx))
  {
    printf("ERROR: SUNContext_Create failed\n");
    return (-1);
  }

  if (argc < 5)
  {
    if (rank == 0)
    {
      printf("ERROR: FIVE (5) arguments required: matrix size, nprow, npcol, "
             "print_all_ranks, print_timing\n");
    }
    return (1);
  }

  /* should TEST_STATUS print on all ranks */
  print_all_ranks = atoi(argv[4]);
  SetPrintAllRanks(print_all_ranks);

  print_timing = atoi(argv[5]);
  SetTiming(print_timing);

  /* Validate the process grid size */
  nprow = atoi(argv[2]);
  npcol = atoi(argv[3]);
  if (nprow <= 0 || npcol <= 0)
  {
    if (rank == 0)
    {
      printf("ERROR: nprow and npcol must be positive integers \n");
    }
    return (1);
  }
  else if (nprocs < (nprow * npcol))
  {
    if (rank == 0) { printf("ERROR: nprocs must be >= nprow*npcol \n"); }
    return (1);
  }

  /* Extract matrix size arguments */
  M = N = (sunindextype)atol(argv[1]);

  /* Setup the process grid */
  superlu_gridinit(MPI_COMM_WORLD, nprow, npcol, &grid);
  /* excess processes just exit */
  if (grid.iam >= nprow * npcol)
  {
    superlu_gridexit(&grid);
    MPI_Finalize();
    return (0);
  }
  /* now update the number of procs available */
  nprocs = nprow * npcol;

  /* Determine local matrix information */
  M_local = Mnprocs = M / nprocs;
  Mrem              = M - M_local * nprocs;
  fst_row           = M_local * grid.iam;
  /* if there is remaining rows, the last process gets them */
  if ((M_local * nprocs) != M)
  {
    if (grid.iam == (nprocs - 1))
    { /* last proc. gets all*/
      M_local = M - M_local * (nprocs - 1);
    }
  }

  /* Initialize matrices and vectors to null */
  D       = NULL;
  A       = NULL;
  I       = NULL;
  x       = NULL;
  y       = NULL;
  gx      = NULL;
  gy      = NULL;
  matdata = NULL;
  rowptrs = NULL;
  colind  = NULL;

  /* Setup global matrix */
  if (grid.iam == 0)
  {
    /* Create the matrix as dense first */
    D = SUNDenseMatrix(M, N, sunctx);

    /* Fill matrix with uniform random data in [0,1/N] */
    for (k = 0; k < N; k++)
    {
      i          = rand() % N;
      j          = rand() % N;
      matdata    = SUNDenseMatrix_Column(D, j);
      matdata[i] = (sunrealtype)rand() / (sunrealtype)RAND_MAX / N;
    }

    /* Add identity to matrix */
    fails = SUNMatScaleAddI(ONE, D);
    if (fails)
    {
      printf(">>> FAIL: SUNDenseMatrix SUNMatScaleAddI failure\n");
      return (1);
    }

    /* create the global vectors */
    gx    = N_VNew_Serial(N, sunctx);
    gy    = N_VNew_Serial(N, sunctx);
    xdata = N_VGetArrayPointer(gx);
    ydata = N_VGetArrayPointer(gy);

    /* fill x with random data, then find y = Dx using Matvec */
    for (i = 0; i < N; i++)
    {
      xdata[i] = (sunrealtype)rand() / (sunrealtype)RAND_MAX;
    }
    if (SUNMatMatvec(D, gx, gy))
    {
      TEST_STATUS(">>> FAIL: SUNMatMatvec failed for dense matrix\n", grid.iam);
      SUNMatDestroy(D);
      N_VDestroy(gx);
      N_VDestroy(gy);
      return (1);
    }

    fails = csr_from_dense(D, ZERO, &matdata, &colind, &rowptrs);
    if (fails != 0)
    {
      printf(">>> FAIL: csr_from_dense failure \n");
      return (1);
    }

    /* distribute matrix and vectors */
    for (i = 1; i < nprocs; i++)
    {
      sunindextype *coltemp, *rowtemp;
      sunrealtype* datatemp;
      sunindextype fstrow_temp, M_localtemp;

      M_localtemp = (i == (nprocs - 1)) ? M_local + Mrem : Mnprocs;
      fstrow_temp = Mnprocs * i;

      /* send number of local NNZ */
      NNZ_local = rowptrs[fstrow_temp + M_localtemp] - rowptrs[fstrow_temp];
      MPI_Send(&NNZ_local, 1, MPI_SUNINDEXTYPE, i, i, grid.comm);

      /* send out rowptrs */
      rowtemp = &rowptrs[fstrow_temp];
      MPI_Send(rowtemp, M_localtemp + 1, MPI_SUNINDEXTYPE, i, i, grid.comm);

      /* send corresponding column indices */
      coltemp = &colind[rowptrs[fstrow_temp]];
      MPI_Send(coltemp, NNZ_local, MPI_SUNINDEXTYPE, i, i, grid.comm);

      /* send corresponding data */
      datatemp = &matdata[rowptrs[fstrow_temp]];
      MPI_Send(datatemp, NNZ_local, MPI_SUNREALTYPE, i, i, grid.comm);

      /* send vector data */
      datatemp = &xdata[fstrow_temp];
      MPI_Send(datatemp, M_localtemp, MPI_SUNREALTYPE, i, i, grid.comm);

      datatemp = &ydata[fstrow_temp];
      MPI_Send(datatemp, M_localtemp, MPI_SUNREALTYPE, i, i, grid.comm);
    }

    NNZ_local = rowptrs[M_local] - rowptrs[0];

    /* Create local SuperLU-DIST SLU_NR_loc SuperMatrix */
    Asuper = NULL;
    Asuper = (SuperMatrix*)malloc(sizeof(SuperMatrix));

    /* Allocates Asuper->Store and sets structure members */
    dCreate_CompRowLoc_Matrix_dist(Asuper, M, N, NNZ_local, M_local, fst_row,
                                   matdata, colind, rowptrs, SLU_NR_loc, SLU_D,
                                   SLU_GE);

    A = SUNMatrix_SLUNRloc(Asuper, &grid, sunctx);
    if (A == NULL)
    {
      fails++;
      TEST_STATUS(">>> FAIL: Failed to create SUNMatrix_SLUNRloc\n", grid.iam);
      SUNMatDestroy(D);
      N_VDestroy(gx);
      N_VDestroy(gy);
      return (fails);
    }

    /* last entry in rowptrs should be the num of local nz */
    rowptrs[M_local] = NNZ_local;

    /* make the local NVectors */
    x = N_VMake_Parallel(grid.comm, M_local, N, xdata, sunctx);
    y = N_VMake_Parallel(grid.comm, M_local, N, ydata, sunctx);
  }
  else
  {
    sunindextype shift;

    /* recieve number of local nnz */
    MPI_Recv(&NNZ_local, 1, MPI_SUNINDEXTYPE, 0, grid.iam, grid.comm, &mpistatus);

    /* Allocate memory for matrix members */
    matdata = (sunrealtype*)malloc(NNZ_local * sizeof(sunrealtype));
    colind  = (sunindextype*)malloc(NNZ_local * sizeof(sunindextype));
    rowptrs = (sunindextype*)malloc((M_local + 1) * sizeof(sunindextype));

    /* receive distributed matrix */
    MPI_Recv(rowptrs, M_local + 1, MPI_SUNINDEXTYPE, 0, grid.iam, grid.comm,
             &mpistatus);
    MPI_Recv(colind, NNZ_local, MPI_SUNINDEXTYPE, 0, grid.iam, grid.comm,
             &mpistatus);
    MPI_Recv(matdata, NNZ_local, MPI_SUNREALTYPE, 0, grid.iam, grid.comm,
             &mpistatus);

    /* localize rowptrs */
    shift = rowptrs[0];
    for (i = 0; i < M_local; i++) { rowptrs[i] = rowptrs[i] - shift; }
    rowptrs[M_local] = NNZ_local;

    /* Create local SuperLU-DIST SuperMatrix */
    Asuper = NULL;
    Asuper = (SuperMatrix*)malloc(sizeof(SuperMatrix));
    dCreate_CompRowLoc_Matrix_dist(Asuper, M, N, NNZ_local, M_local, fst_row,
                                   matdata, colind, rowptrs, SLU_NR_loc, SLU_D,
                                   SLU_GE);

    /* Create local SuperLU-DIST SUNMatrix */
    A = SUNMatrix_SLUNRloc(Asuper, &grid, sunctx);
    if (A == NULL)
    {
      fails++;
      TEST_STATUS(">>> FAIL: Failed to create SUNMatrix_SLUNRloc\n", grid.iam);
      Destroy_CompRowLoc_Matrix_dist(Asuper);
      return (fails);
    }

    /* make the local NVectors */
    x     = N_VNew_Parallel(grid.comm, M_local, N, sunctx);
    y     = N_VNew_Parallel(grid.comm, M_local, N, sunctx);
    xdata = N_VGetArrayPointer(x);
    ydata = N_VGetArrayPointer(y);

    /* recieve vectors */
    MPI_Recv(xdata, M_local, MPI_SUNREALTYPE, 0, grid.iam, grid.comm, &mpistatus);
    MPI_Recv(ydata, M_local, MPI_SUNREALTYPE, 0, grid.iam, grid.comm, &mpistatus);
  }

  /* Create local I matrix that is the same shape as A (which was
   * made with the diagonal) since SuperLU requires the the matrices
   * to be the same shape and have the same sparsity pattern. */
  square = is_square(A);
  if (square)
  {
    /* clone/copy A so we keep colind, rowptrs */
    I = SUNMatClone(A);
    if (I == NULL)
    {
      fails++;
      TEST_STATUS(">>> FAIL: Failed to alloc I\n", grid.iam);
    }
    fails += SUNMatCopy(A, I);

    Isuper  = SUNMatrix_SLUNRloc_SuperMatrix(I);
    Istore  = (NRformat_loc*)Isuper->Store;
    matdata = (sunrealtype*)Istore->nzval;
    rowptrs = (sunindextype*)Istore->rowptr;
    colind  = (sunindextype*)Istore->colind;

    for (i = 0; i < M_local; i++)
    {
      for (j = 0; j < rowptrs[i + 1] - rowptrs[i]; j++)
      {
        if (colind[rowptrs[i] + j] == (fst_row + i))
        {
          matdata[rowptrs[i] + j] = ONE;
        }
        else { matdata[rowptrs[i] + j] = ZERO; }
      }
    }
  }

  /* Commence SUNMatrix module tests */
  fails += Test_SUNMatGetID(A, SUNMATRIX_SLUNRLOC, grid.iam);
  fails += Test_SUNMatClone(A, grid.iam);
  fails += Test_SUNMatCopy(A, grid.iam);
  fails += Test_SUNMatZero(A, grid.iam);
  fails += Test_SUNMatScaleAdd(A, I, grid.iam);
  if (square) { fails += Test_SUNMatScaleAddI(A, I, grid.iam); }
  fails += Test_SUNMatMatvecSetup(A, grid.iam);
  fails += Test_SUNMatMatvec(A, x, y, grid.iam);
  fails += Test_SUNMatSpace(A, grid.iam);

  if (fails)
  {
    TEST_STATUS2("FAIL: SUNMatrix module SUNMatrix_SLUNRloc failed %i "
                 "tests.\n\n",
                 fails, grid.iam);
  }
  else
  {
    TEST_STATUS("SUCCESS: SUNMatrix_SLUNRloc passed all tests.\n\n", grid.iam);
  }

  /* Destroy the SuperLU-DIST SuperMatrices */
  Destroy_CompRowLoc_Matrix_dist(Asuper);
  free(Asuper);
  Asuper = NULL;

  /* Destroy matricies and vectors */
  if (grid.iam == 0)
  {
    SUNMatDestroy(D);
    N_VDestroy(gx);
    N_VDestroy(gy);
  }
  SUNMatDestroy(A);
  SUNMatDestroy(I);
  N_VDestroy(x);
  N_VDestroy(y);

  /* Check all procs for errors */
  MPI_Allreduce(&fails, &globfails, 1, MPI_INT, MPI_MAX, grid.comm);

  superlu_gridexit(&grid);
  SUNContext_Free(&sunctx);
  MPI_Finalize();
  return (globfails);
}

/* ----------------------------------------------------------------------
 * Check matrix
 * --------------------------------------------------------------------*/
int check_matrix(SUNMatrix A, SUNMatrix B, sunrealtype tol)
{
  int failure = 0;
  gridinfo_t* grid;
  SuperMatrix *Asuper, *Bsuper;
  NRformat_loc *Astore, *Bstore;
  sunrealtype *Adata, *Bdata;
  sunindextype *Aindexptrs, *Bindexptrs;
  sunindextype *Aindexvals, *Bindexvals;
  sunindextype i, Arows, Brows, Annz, Bnnz;

  grid = SUNMatrix_SLUNRloc_ProcessGrid(A);

  /* get matrix pointers */
  Asuper     = SUNMatrix_SLUNRloc_SuperMatrix(A);
  Astore     = (NRformat_loc*)Asuper->Store;
  Adata      = (sunrealtype*)Astore->nzval;
  Aindexptrs = Astore->rowptr;
  Aindexvals = Astore->colind;
  Arows      = Astore->m_loc;
  Annz       = Aindexptrs[Arows];

  Bsuper     = SUNMatrix_SLUNRloc_SuperMatrix(B);
  Bstore     = (NRformat_loc*)Bsuper->Store;
  Bdata      = (sunrealtype*)Bstore->nzval;
  Bindexptrs = Bstore->rowptr;
  Bindexvals = Bstore->colind;
  Brows      = Bstore->m_loc;
  Bnnz       = Bindexptrs[Brows];

  /* matrices must have same sparsetype, shape and actual data lengths */
  if (SUNMatGetID(A) != SUNMatGetID(B))
  {
    TEST_STATUS3(">>> ERROR: check_matrix: Different storage types (%d vs "
                 "%d)\n",
                 SUNMatGetID(A), SUNMatGetID(B), grid->iam);
    return (1);
  }
  if (Asuper->nrow != Bsuper->nrow)
  {
    TEST_STATUS3(">>> ERROR: check_matrix: Different numbers of rows (%ld vs "
                 "%ld)\n",
                 (long int)Asuper->nrow, (long int)Bsuper->nrow, grid->iam);
    return (1);
  }
  if (Asuper->ncol != Bsuper->ncol)
  {
    TEST_STATUS3(">>> ERROR: check_matrix: Different numbers of cols (%ld vs "
                 "%ld)\n",
                 (long int)Asuper->nrow, (long int)Bsuper->nrow, grid->iam);
    return (1);
  }
  if (Arows != Brows)
  {
    TEST_STATUS3(">>> ERROR: check_matrix: Different numbers of local rows "
                 "(%ld vs %ld)\n",
                 (long int)Arows, (long int)Brows, grid->iam);
    return (1);
  }
  if (Annz != Bnnz)
  {
    TEST_STATUS3(">>> ERROR: check_matrix: Different numbers of nonzeros (%ld "
                 "vs %ld)\n",
                 (long int)Annz, (long int)Bnnz, grid->iam);
    return (1);
  }

  /* compare sparsity patterns */
  for (i = 0; i < Arows; i++) { failure += (Aindexptrs[i] != Bindexptrs[i]); }
  if (failure > ZERO)
  {
    TEST_STATUS(">>> ERROR: check_matrix: Different indexptrs \n", grid->iam);
    return (1);
  }
  for (i = 0; i < Annz; i++) { failure += (Aindexvals[i] != Bindexvals[i]); }
  if (failure > ZERO)
  {
    TEST_STATUS(">>> ERROR: check_matrix: Different indexvals \n", grid->iam);
    return (1);
  }

  /* compare global start index */
  if (Astore->fst_row != Bstore->fst_row)
  {
    TEST_STATUS(">>> ERRROR: check_matrix: Different globalidx \n", grid->iam);
    return (1);
  }

  /* compare matrix values */
  for (i = 0; i < Annz; i++)
  {
    failure += SUNRCompareTol(Adata[i], Bdata[i], tol);
    if (SUNRCompareTol(Adata[i], Bdata[i], tol))
    {
      TEST_STATUS3("Adata[%ld] != Bdata[%ld] \n", (long int)i, (long int)i,
                   grid->iam);
    }
  }
  if (failure > ZERO)
  {
    TEST_STATUS(">>> ERROR: check_matrix: Different entries \n", grid->iam);
    return (1);
  }

  return (0);
}

int check_matrix_entry(SUNMatrix A, sunrealtype val, sunrealtype tol)
{
  int failure = 0;
  SuperMatrix* Asuper;
  NRformat_loc* Astore;
  sunrealtype* Adata;
  sunindextype i, nnz_loc;

  /* get matrix pointers */
  Asuper = SUNMatrix_SLUNRloc_SuperMatrix(A);
  Astore = (NRformat_loc*)Asuper->Store;
  Adata  = (sunrealtype*)Astore->nzval;

  /* compare data */
  nnz_loc = Astore->nnz_loc;
  for (i = 0; i < nnz_loc; i++)
  {
    if (SUNRCompareTol(Adata[i], val, tol) != 0)
    {
      printf("rhs=%g\n", std::abs(val) * tol);
      printf("  Adata[%ld] = %g != %g (err = %g)\n", (long int)i, Adata[i], val,
             std::abs(Adata[i] - val));
      failure++;
    }
  }

  if (failure > ZERO) { return (1); }
  else { return (0); }
}

int check_vector(N_Vector x, N_Vector y, sunrealtype tol)
{
  int failure = 0;
  sunrealtype *xdata, *ydata;
  sunindextype xldata, yldata;
  sunindextype i;

  /* get vector data */
  xdata = N_VGetArrayPointer(x);
  ydata = N_VGetArrayPointer(y);

  /* check data lengths */
  xldata = N_VGetLocalLength_Parallel(x);
  yldata = N_VGetLocalLength_Parallel(y);

  if (xldata != yldata)
  {
    TEST_STATUS(">>> ERROR: check_vector: Different data array lengths \n", 0);
    return (1);
  }

  /* check vector data */
  for (i = 0; i < xldata; i++)
  {
    failure += SUNRCompareTol(xdata[i], ydata[i], tol);
  }

  if (failure > ZERO) { return (1); }
  else { return (0); }
}

sunbooleantype has_data(SUNMatrix A)
{
  SuperMatrix* Asuper;
  NRformat_loc* Astore;
  sunrealtype* Adata;

  /* get matrix pointers */
  Asuper = SUNMatrix_SLUNRloc_SuperMatrix(A);
  Astore = (NRformat_loc*)Asuper->Store;
  Adata  = (sunrealtype*)Astore->nzval;

  if (Adata == NULL) { return SUNFALSE; }
  else { return SUNTRUE; }
}

sunbooleantype is_square(SUNMatrix A)
{
  SuperMatrix* Asuper = SUNMatrix_SLUNRloc_SuperMatrix(A);
  if (Asuper->nrow == Asuper->ncol) { return SUNTRUE; }
  else { return SUNFALSE; }
}

int csr_from_dense(SUNMatrix Ad, sunrealtype droptol, sunrealtype** matdata,
                   sunindextype** colind, sunindextype** rowptrs)
{
  sunindextype i, j, nnz;
  sunindextype M, N;

  if (droptol < ZERO) { return -1; }
  if (SUNMatGetID(Ad) != SUNMATRIX_DENSE) { return -1; }

  /* set size of new matrix */
  M = SUNDenseMatrix_Rows(Ad);
  N = SUNDenseMatrix_Columns(Ad);

  /* determine total number of nonzeros */
  nnz = 0;
  for (j = 0; j < N; j++)
  {
    for (i = 0; i < M; i++)
    {
      nnz += (std::abs(SM_ELEMENT_D(Ad, i, j)) > droptol);
    }
  }

  /* allocate */
  (*matdata) = (sunrealtype*)malloc(nnz * sizeof(sunrealtype));
  (*colind)  = (sunindextype*)malloc(nnz * sizeof(sunindextype));
  (*rowptrs) = (sunindextype*)malloc((M + 1) * sizeof(sunindextype));

  /* copy nonzeros from Ad into As, based on CSR/CSC type */
  nnz = 0;
  for (i = 0; i < M; i++)
  {
    (*rowptrs)[i] = nnz;
    for (j = 0; j < N; j++)
    {
      if (std::abs(SM_ELEMENT_D(Ad, i, j)) > droptol)
      {
        (*colind)[nnz]    = j;
        (*matdata)[nnz++] = SM_ELEMENT_D(Ad, i, j);
      }
    }
  }
  (*rowptrs)[M] = nnz;

  return 0;
}

void sync_device(SUNMatrix A)
{
  /* not running on GPU, just return */
  return;
}
