/* ---------------------------------------------------------------------------
 * Programmer(s): David J. Gardner @ 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
 * ---------------------------------------------------------------------------
 * Example problem:
 *
 * The following is a simple example problem with a banded Jacobian, with the
 * program for its solution by CVODE. The problem is the semi-discrete form of
 * the advection-diffusion equation in 2-D:
 *
 *   u_t = u_xx + u_yy + 0.5 u_x
 *
 * on the rectangle 0 <= x <= 2, 0 <= y <= 1, and the time interval 0 <= t <= 1.
 * Homogeneous Dirichlet boundary conditions are posed, and the initial
 * condition is
 *
 *   u(x,y,0) = x (2-x) y (1-y) exp(5xy).
 *
 * The PDE is discretized on a uniform MX+2 by MY+2 grid with central
 * differencing, and with boundary values eliminated, leaving an ODE system of
 * size NEQ = MX*MY.
 *
 * This program solves the problem with the BDF method, Newton iteration with
 * the GMRES linear solver, and a user-supplied Jacobian-vector product routine.
 * It uses scalar relative and absolute tolerances.
 *
 * Output is printed at t = .1, .2, ..., 1. Run statistics (optional outputs)
 * are printed at the end.
 * ---------------------------------------------------------------------------*/

#include <cmath>
#include <cstdlib>
#include <cvode/cvode.h> // access CVODE fcts., consts.
#include <iostream>
#include <nvector/nvector_sycl.h> // access the SYCL NVector
#include <sundials/sundials_core.hpp>
#include <sunlinsol/sunlinsol_spgmr.h> // access the SPGMR SUNLinearSolver

// Real Constants

#define ATOL  SUN_RCONST(1.0e-5) // scalar absolute tolerance
#define T0    SUN_RCONST(0.0)    // initial time
#define T1    SUN_RCONST(0.1)    // first output time
#define DTOUT SUN_RCONST(0.1)    // output time increment
#define NOUT  10                 // number of output times

#define ZERO SUN_RCONST(0.0)
#define HALF SUN_RCONST(0.5)
#define ONE  SUN_RCONST(1.0)
#define TWO  SUN_RCONST(2.0)
#define FIVE SUN_RCONST(5.0)

#if defined(SUNDIALS_EXTENDED_PRECISION)
#define GSYM "Lg"
#define ESYM "Le"
#define FSYM "Lf"
#else
#define GSYM "g"
#define ESYM "e"
#define FSYM "f"
#endif

#if defined(SUNDIALS_INT64_T)
#define DSYM "ld"
#else
#define DSYM "d"
#endif

// User data sstructure contains model and discretization parameters
struct UserData
{
  sycl::queue* myQueue = NULL;              // SYCL queue
  sunindextype MX      = 10;                // interior nodes in the x-direction
  sunindextype MY      = 5;                 // interior nodes in the y-direction
  sunindextype NEQ     = MX * MY;           // number of equations
  sunrealtype xmax     = SUN_RCONST(2.0);   // x-domain boundary
  sunrealtype ymax     = SUN_RCONST(1.0);   // y-domain boundary
  sunrealtype dx       = xmax / (MX + 1);   // x-direction mesh spacing
  sunrealtype dy       = ymax / (MY + 1);   // y-directino mesh spacing
  sunrealtype hdcoef   = ONE / (dx * dx);   // x-diffusion
  sunrealtype vdcoef   = ONE / (dy * dy);   // y-diffusion
  sunrealtype hacoef   = HALF / (TWO * dx); // x-advection
};

// Functions Called by the Solver
static int f(sunrealtype t, N_Vector u, N_Vector udot, void* user_data);
static int jtv(N_Vector v, N_Vector Jv, sunrealtype t, N_Vector u, N_Vector fu,
               void* user_data, N_Vector tmp);

// Private Helper Functions
static void PrintHeader(sunrealtype reltol, sunrealtype abstol,
                        sunrealtype umax, UserData* data);
static void PrintOutput(sunrealtype t, sunrealtype umax, long int nst);
static void PrintFinalStats(void* cvode_mem);

// Private function to check function return values
static int check_retval(void* returnvalue, const char* funcname, int opt);

/* ---------------------------------------------------------------------------
 * Main Program
 * ---------------------------------------------------------------------------*/

int main(int argc, char** argv)
{
  // SUNDIALS simulation context
  sundials::Context sunctx;

  // return flag value
  int retval;

  // Create an in-order GPU queue
#if SYCL_LANGUAGE_VERSION >= 2020 && !defined(SUNDIALS_SYCL_2020_UNSUPPORTED)
  sycl::queue myQueue(sycl::gpu_selector_v,
                      sycl::property_list{sycl::property::queue::in_order{}});
#else
  sycl::gpu_selector selector;
  sycl::queue myQueue(selector,
                      sycl::property_list{sycl::property::queue::in_order{}});
#endif

  sycl::device dev = myQueue.get_device();
  std::cout << "Running on " << (dev.get_info<sycl::info::device::name>())
            << std::endl;

  // Create user data and set queue
  UserData data;
  data.myQueue = &myQueue;

  // Create a SYCL vector
  N_Vector u = N_VNew_Sycl(data.NEQ, &myQueue, sunctx);
  if (check_retval((void*)u, "N_VNew_Sycl", 0)) { return 1; }

  // Extract host pointer to solution vector data on the host
  sunrealtype* udata = N_VGetArrayPointer(u);

  // Load initial profile into u vector
  for (sunindextype tid = 0; tid < data.NEQ; tid++)
  {
    sunindextype i = tid / data.MY; // x-node index
    sunindextype j = tid % data.MY; // y-node index

    sunrealtype x = (i + 1) * data.dx; // x location
    sunrealtype y = (j + 1) * data.dy; // y location

    udata[tid] = x * (data.xmax - x) * y * (data.ymax - y) *
                 std::exp(FIVE * x * y);
  }

  // Copy data to device
  N_VCopyToDevice_Sycl(u);

  // Create CVODE and specify the Backward Differentiation Formula
  void* cvode_mem = CVodeCreate(CV_BDF, sunctx);
  if (check_retval((void*)cvode_mem, "CVodeCreate", 0)) { return 1; }

  // Specify the right hand side function in f(t,u), initial condition (t0, u0)
  retval = CVodeInit(cvode_mem, f, T0, u);
  if (check_retval(&retval, "CVodeInit", 1)) { return 1; }

  // Specify the scalar relative tolerance and scalar absolute tolerance
  sunrealtype reltol = ZERO;
  sunrealtype abstol = ATOL;
  retval             = CVodeSStolerances(cvode_mem, reltol, abstol);
  if (check_retval(&retval, "CVodeSStolerances", 1)) { return 1; }

  // Set the pointer to user-defined data
  retval = CVodeSetUserData(cvode_mem, &data);
  if (check_retval(&retval, "CVodeSetUserData", 1)) { return 1; }

  // Create SPGMR solver without preconditioning and default Krylov dimension
  SUNLinearSolver LS = SUNLinSol_SPGMR(u, SUN_PREC_NONE, 0, sunctx);
  if (check_retval(&retval, "SUNLinSol_SPGMR", 1)) { return 1; }

  // Attach the linear sovler to CVODE
  retval = CVodeSetLinearSolver(cvode_mem, LS, NULL);
  if (check_retval(&retval, "CVodeSetLinearSolver", 1)) { return 1; }

  // Set the Jacobian-times-vector function
  retval = CVodeSetJacTimes(cvode_mem, NULL, jtv);
  if (check_retval(&retval, "CVodeSetJacTimesVecFn", 1)) { return 1; }

  // In loop over output points: call CVODE, print results, test for errors
  sunrealtype umax = N_VMaxNorm(u);
  PrintHeader(reltol, abstol, umax, &data);

  sunrealtype tout = T1; // output time
  sunrealtype t;         // CVODE return time
  long int nst;          // number of time steps

  for (int iout = 0; iout < NOUT; iout++)
  {
    // Advance in time
    retval = CVode(cvode_mem, tout, u, &t, CV_NORMAL);
    if (check_retval(&retval, "CVode", 1)) { break; }

    // Output status
    retval = CVodeGetNumSteps(cvode_mem, &nst);
    if (check_retval(&retval, "CVodeGetNumSteps", 1)) { break; }

    umax = N_VMaxNorm(u);
    PrintOutput(t, umax, nst);

    // Update output time
    tout += DTOUT;
  }

  PrintFinalStats(cvode_mem); // Print some final statistics

  N_VDestroy(u);         // Free the u vector
  CVodeFree(&cvode_mem); // Free the integrator memory
  SUNLinSolFree(LS);     // Free linear solver memory

  return 0;
}

/* ---------------------------------------------------------------------------
 * Functions called by the solver
 * ---------------------------------------------------------------------------*/

// Compute the ODE right-hand side function f(t,u).
static int f(sunrealtype t, N_Vector u, N_Vector udot, void* user_data)
{
  UserData* data = static_cast<UserData*>(user_data);

  // Extract needed constants from data
  const size_t MX         = static_cast<size_t>(data->MX);
  const size_t MY         = static_cast<size_t>(data->MY);
  const sunrealtype hordc = data->hdcoef;
  const sunrealtype horac = data->hacoef;
  const sunrealtype verdc = data->vdcoef;

  // Extract pointers to vector data
  const sunrealtype* udata = N_VGetDeviceArrayPointer(u);
  sunrealtype* dudata      = N_VGetDeviceArrayPointer(udot);

  data->myQueue->submit(
    [&](sycl::handler& h)
    {
      h.parallel_for(sycl::range{MX, MY},
                     [=](sycl::id<2> idx)
                     {
                       sunindextype i   = idx[0];
                       sunindextype j   = idx[1];
                       sunindextype tid = i * MY + j;

                       sunrealtype uij = udata[tid];
                       sunrealtype udn = (j == 0) ? ZERO : udata[tid - 1];
                       sunrealtype uup = (j == MY - 1) ? ZERO : udata[tid + 1];
                       sunrealtype ult = (i == 0) ? ZERO : udata[tid - MY];
                       sunrealtype urt = (i == MX - 1) ? ZERO : udata[tid + MY];

                       // Set diffusion and advection terms and load into udot
                       sunrealtype hdiff = hordc * (ult - TWO * uij + urt);
                       sunrealtype vdiff = verdc * (uup - TWO * uij + udn);
                       sunrealtype hadv  = horac * (urt - ult);

                       dudata[tid] = hdiff + vdiff + hadv;
                     });
    });

  return 0;
}

// Jacobian-times-vector routine.
static int jtv(N_Vector v, N_Vector Jv, sunrealtype t, N_Vector u, N_Vector fu,
               void* user_data, N_Vector tmp)
{
  UserData* data = static_cast<UserData*>(user_data);

  // Extract needed constants from data
  const size_t MX         = static_cast<size_t>(data->MX);
  const size_t MY         = static_cast<size_t>(data->MY);
  const sunrealtype hordc = data->hdcoef;
  const sunrealtype horac = data->hacoef;
  const sunrealtype verdc = data->vdcoef;

  // Extract pointers to vector data
  const sunrealtype* vdata = N_VGetDeviceArrayPointer(v);
  sunrealtype* Jvdata      = N_VGetDeviceArrayPointer(Jv);

  data->myQueue->submit(
    [&](sycl::handler& h)
    {
      h.parallel_for(sycl::range{MX, MY},
                     [=](sycl::id<2> idx)
                     {
                       sunindextype i   = idx[0];
                       sunindextype j   = idx[1];
                       sunindextype tid = i * MY + j;

                       // set the tid-th element of Jv
                       Jvdata[tid] = -TWO * (verdc + hordc) * vdata[tid];

                       if (i != 0)
                         Jvdata[tid] += (hordc - horac) * vdata[tid - MY];
                       if (i != MX - 1)
                         Jvdata[tid] += (hordc + horac) * vdata[tid + MY];
                       if (j != 0) Jvdata[tid] += verdc * vdata[tid - 1];
                       if (j != MY - 1) Jvdata[tid] += verdc * vdata[tid + 1];
                     });
    });

  return 0;
}

/* ---------------------------------------------------------------------------
 * Private helper functions
 * ---------------------------------------------------------------------------*/

// Print first lines of output (problem description)
static void PrintHeader(sunrealtype reltol, sunrealtype abstol,
                        sunrealtype umax, UserData* data)
{
  std::cout << "\n2-D Advection-Diffusion Equation" << std::endl;
  std::cout << "Mesh dimensions = " << data->MX << " X " << data->MY << std::endl;
  std::cout << "Total system size = " << data->NEQ << std::endl;
  std::cout << "Tolerance parameters: reltol = " << reltol
            << "   abstol = " << abstol << std::endl
            << std::endl;
  std::cout << "At t = " << T0 << "      max.norm(u) = " << umax << std::endl;
  return;
}

// Print current value
static void PrintOutput(sunrealtype t, sunrealtype umax, long int nst)
{
  std::cout << "At t = " << t << "   max.norm(u) = " << umax
            << "   nst = " << nst << std::endl;
  return;
}

// Get and print some final statistics
static void PrintFinalStats(void* cvode_mem)
{
  long lenrw, leniw;
  long lenrwLS, leniwLS;
  long int nst, nfe, nsetups, nni, ncfn, netf;
  long int nli, npe, nps, ncfl, nfeLS;
  int retval;

  retval = CVodeGetWorkSpace(cvode_mem, &lenrw, &leniw);
  check_retval(&retval, "CVodeGetWorkSpace", 1);
  retval = CVodeGetNumSteps(cvode_mem, &nst);
  check_retval(&retval, "CVodeGetNumSteps", 1);
  retval = CVodeGetNumRhsEvals(cvode_mem, &nfe);
  check_retval(&retval, "CVodeGetNumRhsEvals", 1);
  retval = CVodeGetNumLinSolvSetups(cvode_mem, &nsetups);
  check_retval(&retval, "CVodeGetNumLinSolvSetups", 1);
  retval = CVodeGetNumErrTestFails(cvode_mem, &netf);
  check_retval(&retval, "CVodeGetNumErrTestFails", 1);
  retval = CVodeGetNumNonlinSolvIters(cvode_mem, &nni);
  check_retval(&retval, "CVodeGetNumNonlinSolvIters", 1);
  retval = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn);
  check_retval(&retval, "CVodeGetNumNonlinSolvConvFails", 1);

  retval = CVodeGetLinWorkSpace(cvode_mem, &lenrwLS, &leniwLS);
  check_retval(&retval, "CVodeGetLinWorkSpace", 1);
  retval = CVodeGetNumLinIters(cvode_mem, &nli);
  check_retval(&retval, "CVodeGetNumLinIters", 1);
  retval = CVodeGetNumPrecEvals(cvode_mem, &npe);
  check_retval(&retval, "CVodeGetNumPrecEvals", 1);
  retval = CVodeGetNumPrecSolves(cvode_mem, &nps);
  check_retval(&retval, "CVodeGetNumPrecSolves", 1);
  retval = CVodeGetNumLinConvFails(cvode_mem, &ncfl);
  check_retval(&retval, "CVodeGetNumLinConvFails", 1);
  retval = CVodeGetNumLinRhsEvals(cvode_mem, &nfeLS);
  check_retval(&retval, "CVodeGetNumLinRhsEvals", 1);

  std::cout << "\nFinal Statistics.. \n\n";
  std::cout << "lenrw   = " << lenrw << "     leniw   = " << leniw << "\n";
  std::cout << "lenrwLS = " << lenrwLS << "     leniwLS = " << leniwLS << "\n";
  std::cout << "nst     = " << nst << "\n";
  std::cout << "nfe     = " << nfe << "     nfeLS   = " << nfeLS << "\n";
  std::cout << "nni     = " << nni << "     nli     = " << nli << "\n";
  std::cout << "nsetups = " << nsetups << "     netf    = " << netf << "\n";
  std::cout << "npe     = " << npe << "     nps     = " << nps << "\n";
  std::cout << "ncfn    = " << ncfn << "     ncfl    = " << ncfl << "\n\n";

  return;
}

/* Check function return value...
     opt == 0 means SUNDIALS function allocates memory so check if
              returned NULL pointer
     opt == 1 means SUNDIALS function returns an integer value so check if
              retval >= 0
     opt == 2 means function allocates memory so check if returned
              NULL pointer */
static int check_retval(void* returnvalue, const char* funcname, int opt)
{
  int* retval;

  if (opt == 0 && returnvalue == NULL)
  {
    // Check if SUNDIALS function returned NULL pointer - no memory allocated
    std::cerr << "\nSUNDIALS_ERROR: " << funcname
              << " failed - returned NULL pointer\n\n";
    return 1;
  }
  else if (opt == 1)
  {
    // Check if retval < 0
    retval = static_cast<int*>(returnvalue);
    if (*retval < 0)
    {
      std::cerr << "\nSUNDIALS_ERROR: " << funcname
                << " failed with retval = " << *retval << "\n\n";
      return 1;
    }
  }
  else if (opt == 2 && returnvalue == NULL)
  {
    // Check if function returned NULL pointer - no memory allocated
    std::cerr << "\nMEMORY_ERROR: " << funcname
              << " failed - returned NULL pointer\n\n";
    return 1;
  }

  return 0;
}
