#include "Kokkos_Core.hpp"
#include "KokkosKernels_default_types.hpp"
#include "KokkosKernels_Handle.hpp"
#include "KokkosKernels_IOUtils.hpp"
#include "KokkosSparse_spmv.hpp"
#include "KokkosSparse_CrsMatrix.hpp"
#include "KokkosSparse_gauss_seidel.hpp"
#include "KokkosBlas1_nrm2.hpp"

//Parallel Gauss-Seidel Preconditioner/Smoother
//  -Uses graph coloring to find independent row sets,
//   and applies GS to each set in parallel
//  -Here, use to solve a diagonally dominant linear system directly.

//Helper to print out colors in the shape of the grid
int main()
{
  using Scalar  = default_scalar;
  using Mag     = Kokkos::ArithTraits<Scalar>::mag_type;
  using Ordinal = default_lno_t;
  using Offset  = default_size_type;
  using ExecSpace = Kokkos::DefaultExecutionSpace;
  using MemSpace = typename ExecSpace::memory_space;
  using Device  = Kokkos::Device<ExecSpace, MemSpace>;
  using Handle  = KokkosKernels::Experimental::
    KokkosKernelsHandle<Offset, Ordinal, default_scalar, ExecSpace, MemSpace, MemSpace>;
  using Matrix  = KokkosSparse::CrsMatrix<Scalar, Ordinal, Device, void, Offset>;
  using Vector  = typename Matrix::values_type;
  constexpr Ordinal numRows = 10000;
  const Scalar one = Kokkos::ArithTraits<Scalar>::one();
  const Mag magOne = Kokkos::ArithTraits<Mag>::one();
  //Solve tolerance
  const Mag tolerance = 1e-6 * magOne;
  Kokkos::initialize();
  {
    //Generate a square, strictly diagonally dominant, but nonsymmetric matrix on which Gauss-Seidel should converge.
    //Get approx. 20 entries per row
    //Diagonals are 2x the absolute sum of all other entries.
    Offset nnz = numRows * 20;
    Matrix A = KokkosKernels::Impl::kk_generate_diagonally_dominant_sparse_matrix<Matrix>(numRows, numRows, nnz, 2, 100, 1.05 * one);
    std::cout << "Generated a matrix with " << numRows << " rows/cols, and " << nnz << " entries.\n";
    //Create a kernel handle, then a Gauss-Seidel handle with the default algorithm
    Handle handle;
    handle.create_gs_handle(KokkosSparse::GS_DEFAULT);
    //Set up Gauss-Seidel for the graph (matrix sparsity pattern)
    KokkosSparse::Experimental::gauss_seidel_symbolic(&handle, numRows, numRows, A.graph.row_map, A.graph.entries, false);
    //Set up Gauss-Seidel for the matrix values (numeric)
    //Another matrix with the same sparsity pattern could re-use the handle and symbolic phase, and only call numeric.
    KokkosSparse::Experimental::gauss_seidel_numeric(&handle, numRows, numRows, A.graph.row_map, A.graph.entries, A.values, false);
    //Now, preconditioner is ready to use. Set up an unknown vector (uninitialized) and randomized right-hand-side vector.
    Vector x(Kokkos::view_alloc(Kokkos::WithoutInitializing, "x"), numRows);
    Vector b(Kokkos::view_alloc(Kokkos::WithoutInitializing, "b"), numRows);
    Vector res(Kokkos::view_alloc(Kokkos::WithoutInitializing, "res"), numRows);
    auto bHost = Kokkos::create_mirror_view(b);
    for(Ordinal i = 0; i < numRows; i++)
      bHost(i) = 3 * ((one * rand()) / RAND_MAX);
    Kokkos::deep_copy(b, bHost);
    //Measure initial residual norm ||Ax - b||, where x is 0
    Mag initialRes = KokkosBlas::nrm2(b);
    Mag scaledResNorm = magOne;
    bool firstIter = true;
    //Iterate until reaching the tolerance
    int numIters = 0;
    while(scaledResNorm > tolerance)
    {
      //Run one sweep of forward Gauss-Seidel (SOR with omega = 1.0)
      //If this is the first iteration, tell apply:
      //  * to zero out x (it was uninitialized)
      //  * that b has changed since the previous apply (since there was no previous apply)
      KokkosSparse::Experimental::forward_sweep_gauss_seidel_apply(
          &handle, numRows, numRows,
          A.graph.row_map, A.graph.entries, A.values,
          x, b, firstIter, firstIter, one, 1);
      firstIter = false;
      //Now, compute the new residual norm using SPMV
      Kokkos::deep_copy(res, b);
      //Compute res := Ax - res (since res is now equal to b, this is Ax - b)
      KokkosSparse::spmv("N", one, A, x, -one, res);
      //Recompute the scaled norm
      scaledResNorm = KokkosBlas::nrm2(res) / initialRes;
      numIters++;
      std::cout << "Iteration " << numIters << " scaled residual norm: " << scaledResNorm << '\n';
    }
    std::cout << "SUCCESS: converged in " << numIters << " iterations.\n";
  }
  Kokkos::finalize();
  return 0;
}

