#include <cassert>
#include <vector>
#include <algorithm>

#include <stdexcept>

#include "Kokkos_Core.hpp"

#include "KokkosBatched_Util.hpp"

#define TEST_ASSERT(m, success)                                \
  if (!(m)) {                                                  \
    success = false;                                           \
    printf("FAILED: %s, at %d, %s\n", #m, __LINE__, __FILE__); \
  }

namespace KokkosBatched {

namespace Test {

typedef int ordinal_type;
typedef int size_type;
typedef double scalar_type;
#define BLOCKCRS_MAX_BLOCKSIZE 32
#define FLOP_MUL 1.0
#define FLOP_ADD 1.0

double LU_FlopCount(int mm, int nn) {
  double m = (double)mm;
  double n = (double)nn;
  if (m > n)
    return (FLOP_MUL * (0.5 * m * n * n - (1.0 / 6.0) * n * n * n +
                        0.5 * m * n - 0.5 * n * n + (2.0 / 3.0) * n) +
            FLOP_ADD * (0.5 * m * n * n - (1.0 / 6.0) * n * n * n -
                        0.5 * m * n + (1.0 / 6.0) * n));
  else
    return (FLOP_MUL * (0.5 * n * m * m - (1.0 / 6.0) * m * m * m +
                        0.5 * n * m - 0.5 * m * m + (2.0 / 3.0) * m) +
            FLOP_ADD * (0.5 * n * m * m - (1.0 / 6.0) * m * m * m -
                        0.5 * n * m + (1.0 / 6.0) * m));
}

double Trsm_Lower_FlopCountLower(int mm, int nn) {
  double m = (double)mm;
  double n = (double)nn;
  return (FLOP_MUL * (0.5 * m * n * (n + 1.0)) +
          FLOP_ADD * (0.5 * m * n * (n - 1.0)));
}

double Trsm_Upper_FlopCountUpper(int mm, int nn) {
  double m = (double)mm;
  double n = (double)nn;
  return (FLOP_MUL * (0.5 * m * n * (n + 1.0)) +
          FLOP_ADD * (0.5 * m * n * (n - 1.0)));
}

double Gemm_FlopCount(int mm, int nn, int kk) {
  double m = (double)mm;
  double n = (double)nn;
  double k = (double)kk;
  return (FLOP_MUL * (m * n * k) + FLOP_ADD * (m * n * k));
}

template <typename aViewType, typename bViewType>
double compute_relative_diff(const aViewType a, const bViewType b) {
  // Bring the vectors to the host. This is just a correctness checker.
  auto aa = Kokkos::create_mirror_view(a);
  Kokkos::deep_copy(aa, a);
  auto bb = Kokkos::create_mirror_view(b);
  Kokkos::deep_copy(bb, b);

  double diff2 = 0, norm2 = 0;
  for (ordinal_type i = 0, iend = aa.extent(0); i < iend; ++i)
    for (ordinal_type j = 0, jend = aa.extent(1); j < jend; ++j)
      for (ordinal_type k = 0, kend = aa.extent(2); k < kend; ++k)
        for (ordinal_type l = 0, lend = aa.extent(3); l < lend; ++l) {
          const double val  = aa.access(i, j, k, l),
                       diff = aa.access(i, j, k, l) - bb.access(i, j, k, l);
          diff2 += diff * diff;
          norm2 += val * val;
        }

  return std::sqrt(diff2 / norm2);
}

// Representation of a structured block mesh. The fastest index is k.
struct StencilShape {
  enum Enum { cross };
};

struct StructuredBlock {
  const ordinal_type ni, nj, nk;

  StructuredBlock(const ordinal_type ni_, const ordinal_type nj_,
                  const ordinal_type nk_)
      : ni(ni_), nj(nj_), nk(nk_), _njnk(nj_ * nk_) {}

  KOKKOS_INLINE_FUNCTION
  size_type size() const { return ni * nj * nk; }

  KOKKOS_INLINE_FUNCTION
  size_type ij2id(const ordinal_type i, const ordinal_type j) const {
    return i * nj + j;
  }

  KOKKOS_INLINE_FUNCTION
  void id2ij(const size_type id, ordinal_type &i, ordinal_type &j) const {
    i = id / nj;
    j = id % nj;
  }

  KOKKOS_INLINE_FUNCTION
  size_type ijk2id(const ordinal_type i, const ordinal_type j,
                   const ordinal_type k) const {
    return (i * nj + j) * nk + k;
  }

  KOKKOS_INLINE_FUNCTION
  void id2ijk(const size_type id, ordinal_type &i, ordinal_type &j,
              ordinal_type &k) const {
    i = id / _njnk;
    k = id % _njnk;
    j = k / nk;
    k = k % nk;
  }

 private:
  const ordinal_type _njnk;
};

template <typename ExecSpace, typename ArrayLayout>
struct CrsGraph {
  typedef ExecSpace exec_space;
  typedef ArrayLayout array_layout;

  typedef Kokkos::View<size_type *, array_layout, exec_space> row_ptr_type;
  typedef Kokkos::View<ordinal_type *, array_layout, exec_space> row_idx_type;
  typedef Kokkos::View<ordinal_type *, array_layout, exec_space> col_idx_type;

  row_ptr_type rowptr;
  row_idx_type rowidx;
  col_idx_type colidx;

  CrsGraph() : rowptr("rowptr", 1), rowidx("rowidx", 0), colidx("colidx", 0) {}

  KOKKOS_INLINE_FUNCTION
  bool isEmpty() const {
    return (rowptr.extent(0) <= 1 || colidx.extent(0) == 0 ||
            rowidx.extent(0) == 0);
  }

  KOKKOS_INLINE_FUNCTION
  ordinal_type NumRows() const {
    return (isEmpty() ? 0 : static_cast<ordinal_type>(rowptr.extent(0)) - 1);
  }

  KOKKOS_INLINE_FUNCTION
  size_type NumNonZeros() const {
    return (isEmpty() ? 0 : static_cast<size_type>(colidx.extent(0)));
  }
};

template <typename DstSpace, typename SrcSpace, typename ArrayLayout>
inline CrsGraph<DstSpace, ArrayLayout> create_mirror(
    const CrsGraph<SrcSpace, ArrayLayout> src) {
  CrsGraph<DstSpace, ArrayLayout> dst;

  dst.rowptr =
      Kokkos::create_mirror_view(typename DstSpace::memory_space(), src.rowptr);
  dst.rowidx =
      Kokkos::create_mirror_view(typename DstSpace::memory_space(), src.rowidx);
  dst.colidx =
      Kokkos::create_mirror_view(typename DstSpace::memory_space(), src.colidx);

  return dst;
}

template <typename DstSpace, typename SrcSpace, typename ArrayLayout>
inline void deep_copy(const CrsGraph<DstSpace, ArrayLayout> dst,
                      const CrsGraph<SrcSpace, ArrayLayout> src) {
  Kokkos::deep_copy(dst.rowptr, src.rowptr);
  Kokkos::deep_copy(dst.rowidx, src.rowidx);
  Kokkos::deep_copy(dst.colidx, src.colidx);
}

// Given a structured block and a stencil (at present, just a 3D 1-hop cross),
// construct a corresponding CRS graph.
template <typename ArrayLayout>
CrsGraph<Kokkos::DefaultHostExecutionSpace, ArrayLayout>
create_graph_host_for_structured_block(const StructuredBlock mesh,
                                       const StencilShape::Enum shape) {
  CrsGraph<Kokkos::DefaultHostExecutionSpace, ArrayLayout> graph;

  Kokkos::resize(graph.rowptr, mesh.size() + 1);
  graph.rowptr[0] = 0;

  std::vector<ordinal_type> colidx, rowidx;
  switch (shape) {
    case StencilShape::cross:
      for (ordinal_type c = 0; c < mesh.size(); ++c) {
        ordinal_type i, j, k, n = 0;

        mesh.id2ijk(c, i, j, k);

        rowidx.push_back(c);
        colidx.push_back(c);
        ++n;
        if (i > 0) {
          rowidx.push_back(c);
          colidx.push_back(mesh.ijk2id(i - 1, j, k));
          ++n;
        }
        if (i + 1 < mesh.ni) {
          rowidx.push_back(c);
          colidx.push_back(mesh.ijk2id(i + 1, j, k));
          ++n;
        }
        if (j > 0) {
          rowidx.push_back(c);
          colidx.push_back(mesh.ijk2id(i, j - 1, k));
          ++n;
        }
        if (j + 1 < mesh.nj) {
          rowidx.push_back(c);
          colidx.push_back(mesh.ijk2id(i, j + 1, k));
          ++n;
        }
        if (k > 0) {
          rowidx.push_back(c);
          colidx.push_back(mesh.ijk2id(i, j, k - 1));
          ++n;
        }
        if (k + 1 < mesh.nk) {
          rowidx.push_back(c);
          colidx.push_back(mesh.ijk2id(i, j, k + 1));
          ++n;
        }
        graph.rowptr[c + 1] = graph.rowptr[c] + n;
      }
      break;
  }
  assert(graph.rowptr[mesh.size()] == static_cast<size_type>(colidx.size()));
  assert(graph.rowptr[mesh.size()] == static_cast<size_type>(rowidx.size()));

  for (ordinal_type c = 0; c < mesh.size(); ++c)
    std::sort(colidx.begin() + graph.rowptr[c],
              colidx.begin() + graph.rowptr[c + 1]);

  const ordinal_type nnz = graph.rowptr[mesh.size()];
  Kokkos::resize(graph.colidx, nnz);
  Kokkos::resize(graph.rowidx, nnz);
  for (ordinal_type c = 0; c < nnz; ++c) {
    graph.colidx[c] = colidx[c];
    graph.rowidx[c] = rowidx[c];
  }
  return graph;
}

template <typename ExeSpace, typename ArrayLayout>
class BlockCrsMatrix {
 public:
  typedef ExeSpace exec_space;
  typedef ArrayLayout array_layout;
  typedef Test::CrsGraph<exec_space, array_layout> crs_graph_type;

  typedef scalar_type value_type;
  typedef Kokkos::View<scalar_type ***, array_layout, exec_space>
      value_array_type;

 private:
  crs_graph_type _graph;
  ordinal_type _blocksize;
  value_array_type _values;

 public:
  BlockCrsMatrix() : _graph(), _blocksize(), _values() {}

  BlockCrsMatrix(const BlockCrsMatrix &b)
      : _graph(b._graph), _blocksize(b._blocksize), _values(b._values) {}

  BlockCrsMatrix(const crs_graph_type graph, const ordinal_type blocksize)
      : _graph(graph),
        _blocksize(blocksize),
        _values("BlockCrsMatrix::_values", _graph.NumNonZeros(), _blocksize,
                _blocksize) {}

  BlockCrsMatrix(const crs_graph_type graph, const ordinal_type blocksize,
                 const value_array_type values)
      : _graph(graph), _blocksize(blocksize), _values(values) {}

  ordinal_type BlockSize() const { return _blocksize; }
  crs_graph_type CrsGraph() const { return _graph; }
  value_array_type Values() const { return _values; }
};

template <typename DstSpace, typename SrcSpace, typename ArrayLayout>
inline BlockCrsMatrix<DstSpace, ArrayLayout> create_mirror(
    const BlockCrsMatrix<SrcSpace, ArrayLayout> src) {
  const auto graph     = create_mirror<DstSpace>(src.CrsGraph());
  const auto blocksize = src.BlockSize();
  const auto values    = Kokkos::create_mirror_view(
      typename DstSpace::memory_space(), src.Values());
  return BlockCrsMatrix<DstSpace, ArrayLayout>(graph, blocksize, values);
}

template <typename DstSpace, typename SrcSpace, typename ArrayLayout>
inline void deep_copy(const BlockCrsMatrix<DstSpace, ArrayLayout> dst,
                      const BlockCrsMatrix<SrcSpace, ArrayLayout> src) {
  deep_copy(dst.CrsGraph(), src.CrsGraph());
  Kokkos::deep_copy(dst.Values(), src.Values());
}

template <typename ArrayLayout>
void fill_block_crs_matrix_host(
    BlockCrsMatrix<Kokkos::DefaultHostExecutionSpace, ArrayLayout> A) {
  // extract graph and blocksizes
  const auto graph             = A.CrsGraph();
  const auto values            = A.Values();
  const ordinal_type blocksize = A.BlockSize();

  scalar_type tmp[BLOCKCRS_MAX_BLOCKSIZE *
                  BLOCKCRS_MAX_BLOCKSIZE],  //[blocksize*blocksize],
      diag_block[BLOCKCRS_MAX_BLOCKSIZE]
                [BLOCKCRS_MAX_BLOCKSIZE],  //[blocksize][blocksize],
      offdiag_block[BLOCKCRS_MAX_BLOCKSIZE]
                   [BLOCKCRS_MAX_BLOCKSIZE];  //[blocksize][blocksize];

  Random<scalar_type> random;

  // for diagonal block, make spd
  {
    const ordinal_type iend = blocksize * blocksize;
    for (ordinal_type i = 0; i < iend; ++i) tmp[i] = 2 * (random.value() - 0.5);

    for (ordinal_type i = 0; i < blocksize; ++i)
      for (ordinal_type j = i; j < blocksize; ++j) {
        diag_block[i][j] = 0;
        for (ordinal_type k = 0; k < blocksize; ++k)
          diag_block[i][j] += tmp[i * blocksize + k] * tmp[j * blocksize + k];
        if (i != j)
          diag_block[j][i] = diag_block[i][j];  // symmetrize
        else
          diag_block[i][j] *= 0.5 * blocksize;  // improve condition
      }
  }

  {
    // for off diagonal; down-weight off-diag blocks to improve conditioning.
    for (ordinal_type i = 0; i < blocksize; ++i)
      for (ordinal_type j = 0; j < blocksize; ++j)
        offdiag_block[i][j] = 0.1 * 2 * (random.value() - 0.5);
  }

  for (ordinal_type r = 0; r < graph.NumRows(); ++r) {
    // random number generator (-1, 1)
    const ordinal_type cbegin = graph.rowptr(r), cend = graph.rowptr(r + 1);
    for (ordinal_type c = cbegin; c < cend; ++c) {
      auto block = Kokkos::subview(values, c, Kokkos::ALL(), Kokkos::ALL());

      if (graph.colidx(c) == r) {
        for (ordinal_type i = 0; i < blocksize; ++i)
          for (ordinal_type j = i; j < blocksize; ++j)
            block(i, j) = diag_block[i][j];
      } else {
        // for off diagonal; down-weight off-diag blocks to improve
        // conditioning.
        for (ordinal_type i = 0; i < blocksize; ++i)
          for (ordinal_type j = 0; j < blocksize; ++j)
            block(i, j) = offdiag_block[i][j];
      }
    }
  }
}

// nrhs should go after blocksize to match matrix dimensions consistently
template <typename ExeSpace, typename ArrayLayout>
class BlockMultiVector {
 public:
  typedef ExeSpace exec_space;
  typedef ArrayLayout array_layout;

  typedef scalar_type value_type;
  typedef Kokkos::View<scalar_type ***, array_layout, exec_space>
      value_array_type;

 private:
  value_array_type _values;

 public:
  BlockMultiVector(const ordinal_type nvecs, const ordinal_type nrows,
                   const ordinal_type blocksize)
      : _values("BlockMultiVector::_values", nvecs, nrows, blocksize) {}

  BlockMultiVector(const value_array_type values) : _values(values) {}

  ordinal_type NumVectors() const { return _values.extent(0); }
  ordinal_type NumRows() const { return _values.extent(1); }
  ordinal_type BlockSize() const { return _values.extent(2); }

  value_array_type Values() const { return _values; }
};

template <typename DstSpace, typename SrcSpace, typename ArrayLayout>
inline BlockMultiVector<DstSpace, ArrayLayout> create_mirror(
    const BlockMultiVector<SrcSpace, ArrayLayout> src) {
  return BlockMultiVector<DstSpace, ArrayLayout>(Kokkos::create_mirror_view(
      typename DstSpace::memory_space(), src.Values()));
}

template <typename DstSpace, typename SrcSpace, typename ArrayLayout>
inline void deep_copy(const BlockMultiVector<DstSpace, ArrayLayout> dst,
                      const BlockMultiVector<SrcSpace, ArrayLayout> src) {
  Kokkos::deep_copy(dst.Values(), src.Values());
}

template <typename ArrayLayout>
void fill_block_multi_vector_host(
    BlockMultiVector<Kokkos::DefaultHostExecutionSpace, ArrayLayout> B) {
  const ordinal_type jend = B.NumVectors(), iend = B.NumRows(),
                     kend = B.BlockSize();

  auto B_val = B.Values();

  for (ordinal_type j = 0; j < jend; ++j)
    for (ordinal_type i = 0; i < iend; ++i)
      for (ordinal_type k = 0; k < kend; ++k)
        B_val(j, i, k) = static_cast<double>((i + j + k) % 7) - 3;
}

template <typename ExecSpace, typename ValueType, typename ArrayLayout>
class BlockTridiagMatrices {
 public:
  typedef ExecSpace exec_space;
  typedef ValueType value_type;
  typedef ArrayLayout array_layout;

  typedef Kokkos::View<value_type ****, array_layout, exec_space>
      value_array_type;

 private:
  const ordinal_type _ntridiags, _nrows, _blocksize;
  // A B
  // C
  value_array_type _A, _B, _C;

 public:
  BlockTridiagMatrices(const ordinal_type ntridiags, const ordinal_type nrows,
                       const ordinal_type blocksize)
      : _ntridiags(ntridiags),
        _nrows(nrows),
        _blocksize(blocksize),
        _A("BlockTridiagMatrix::_A", _ntridiags, _nrows, _blocksize,
           _blocksize),
        _B("BlockTridiagMatrix::_B", _ntridiags, _nrows - 1, _blocksize,
           _blocksize),
        _C("BlockTridiagMatrix::_C", _ntridiags, _nrows - 1, _blocksize,
           _blocksize) {}

  BlockTridiagMatrices(const ordinal_type ntridiags, const ordinal_type nrows,
                       const ordinal_type blocksize, const value_array_type A_,
                       const value_array_type B_, const value_array_type C_)
      : _ntridiags(ntridiags),
        _nrows(nrows),
        _blocksize(blocksize),
        _A(A_),
        _B(B_),
        _C(C_) {}

  value_array_type A() const { return _A; }
  value_array_type B() const { return _B; }
  value_array_type C() const { return _C; }

  ordinal_type BlockSize() const { return _blocksize; }
  ordinal_type NumRows() const { return _nrows; }
  ordinal_type NumTridiagMatrices() const { return _ntridiags; }
};

template <typename ExecSpace, typename ValueType, typename ArrayLayout>
BlockTridiagMatrices<ExecSpace, ValueType, ArrayLayout>
create_block_tridiag_matrices(const ordinal_type ntridiags,
                              const ordinal_type nrows,
                              const ordinal_type blocksize) {
  return BlockTridiagMatrices<ExecSpace, ValueType, ArrayLayout>(
      adjustDimension<ValueType>(ntridiags), nrows, blocksize);
}

template <typename DstSpace, typename SrcSpace, typename ValueType,
          typename ArrayLayout>
inline BlockTridiagMatrices<DstSpace, ValueType, ArrayLayout> create_mirror(
    const BlockTridiagMatrices<SrcSpace, ValueType, ArrayLayout> src) {
  return BlockTridiagMatrices<DstSpace, ValueType, ArrayLayout>(
      src.NumTridiagMatrices(), src.NumRows(), src.BlockSize(),
      Kokkos::create_mirror_view(typename DstSpace::memory_space(), src.A()),
      Kokkos::create_mirror_view(typename DstSpace::memory_space(), src.B()),
      Kokkos::create_mirror_view(typename DstSpace::memory_space(), src.C()));
}

template <typename DstSpace, typename SrcSpace, typename ValueType,
          typename ArrayLayout>
inline void deep_copy(
    const BlockTridiagMatrices<DstSpace, ValueType, ArrayLayout> dst,
    const BlockTridiagMatrices<SrcSpace, ValueType, ArrayLayout> src) {
  Kokkos::deep_copy(dst.A(), src.A());
  Kokkos::deep_copy(dst.B(), src.B());
  Kokkos::deep_copy(dst.C(), src.C());
}

template <typename ViewType>
KOKKOS_INLINE_FUNCTION typename std::enable_if<
    std::is_same<typename ViewType::value_type, scalar_type>::value,
    scalar_type &>::type
tdiag_val(const ViewType &A, const ordinal_type &t, const ordinal_type &i,
          const ordinal_type &ii, const ordinal_type &jj) {
  return A(t, i, ii, jj);
}

template <typename ViewType>
KOKKOS_INLINE_FUNCTION typename std::enable_if<
    !std::is_same<typename ViewType::value_type, scalar_type>::value,
    scalar_type &>::type
tdiag_val(const ViewType &A, const ordinal_type &t, const ordinal_type &i,
          const ordinal_type &ii, const ordinal_type &jj) {
  typedef typename ViewType::value_type value_type;
  return A(t / value_type::vector_length, i, ii,
           jj)[t % value_type::vector_length];
}

template <typename ExeSpace, typename ValueType, typename ArrayLayout>
class PartitionedBlockMultiVector {
 public:
  typedef ExeSpace exec_space;
  typedef ValueType value_type;
  typedef ArrayLayout array_layout;

  typedef Kokkos::View<value_type ****, array_layout, exec_space>
      value_array_type;

 private:
  value_array_type _values;

 public:
  PartitionedBlockMultiVector(const ordinal_type nparts,
                              const ordinal_type nvectors,
                              const ordinal_type nrows,
                              const ordinal_type blocksize)
      : _values("BlockMultiVector::_values", nparts, nvectors, nrows,
                blocksize) {}

  PartitionedBlockMultiVector(const value_array_type values)
      : _values(values) {}

  ordinal_type NumPartitions() const { return _values.extent(0); }
  ordinal_type NumVectors() const { return _values.extent(1); }
  ordinal_type NumRows() const { return _values.extent(2); }
  ordinal_type BlockSize() const { return _values.extent(3); }

  value_array_type Values() const { return _values; }
};

template <typename ExecSpace, typename ValueType, typename ArrayLayout>
PartitionedBlockMultiVector<ExecSpace, ValueType, ArrayLayout>
create_partitioned_block_multi_vector(const ordinal_type nparts,
                                      const ordinal_type nvectors,
                                      const ordinal_type nrows,
                                      const ordinal_type blocksize) {
  return PartitionedBlockMultiVector<ExecSpace, ValueType, ArrayLayout>(
      adjustDimension<ValueType>(nparts), nvectors, nrows, blocksize);
}

template <typename DstSpace, typename SrcSpace, typename ValueType,
          typename ArrayLayout>
inline PartitionedBlockMultiVector<DstSpace, ValueType, ArrayLayout>
create_mirror(
    const PartitionedBlockMultiVector<SrcSpace, ValueType, ArrayLayout> src) {
  return PartitionedBlockMultiVector<DstSpace, ValueType, ArrayLayout>(
      Kokkos::create_mirror_view(typename DstSpace::memory_space(),
                                 src.Values()));
}

template <typename DstSpace, typename SrcSpace, typename ValueType,
          typename ArrayLayout>
inline void deep_copy(
    const PartitionedBlockMultiVector<DstSpace, ValueType, ArrayLayout> dst,
    const PartitionedBlockMultiVector<SrcSpace, ValueType, ArrayLayout> src) {
  Kokkos::deep_copy(dst.Values(), src.Values());
}

template <typename ValueType, typename ArrayLayout>
void fill_partitioned_block_multi_vector_host(
    PartitionedBlockMultiVector<Kokkos::DefaultHostExecutionSpace, ValueType,
                                ArrayLayout>
        B,
    const ordinal_type ninj) {
  const ordinal_type iend = ninj,  // B.NumPartitions(),
      jend = B.NumVectors(), kend = B.NumRows(), lend = B.BlockSize();

  auto B_val = B.Values();
  for (ordinal_type i = 0; i < iend; ++i)
    for (ordinal_type j = 0; j < jend; ++j)
      for (ordinal_type k = 0; k < kend; ++k)
        for (ordinal_type l = 0; l < lend; ++l)
          tdiag_val(B_val, i, j, k, l) =
              static_cast<double>((i + j + k + l) % 7) - 3;
}

template <typename ExecSpace, typename ArrayLayout>
class BlockCrsMatrixVectorProductByRow {
 public:
  typedef BlockCrsMatrix<ExecSpace, ArrayLayout> block_crs_matrix_type;
  typedef typename block_crs_matrix_type::crs_graph_type crs_graph_type;
  typedef BlockMultiVector<ExecSpace, ArrayLayout> block_multi_vector_type;

 private:
  ConstUnmanagedViewType<typename crs_graph_type::row_ptr_type> _rowptr;
  ConstUnmanagedViewType<typename crs_graph_type::col_idx_type> _colidx;

  ConstUnmanagedViewType<typename block_crs_matrix_type::value_array_type> _A;
  ConstUnmanagedViewType<typename block_multi_vector_type::value_array_type> _x;
  /**/ UnmanagedViewType<typename block_multi_vector_type::value_array_type> _y;

  ordinal_type _blocksize;

 public:
  // A thread maps to a point row of the matrix.
  // loop = blksize*m
  KOKKOS_INLINE_FUNCTION
  void operator()(const ordinal_type idx) const {
    // index of blockrow and row in a block
    const ordinal_type i  = idx / _blocksize;
    const ordinal_type ii = idx % _blocksize;

    // loop over multivectors
    const ordinal_type jend = _y.extent(0);
    for (ordinal_type j = 0; j < jend; ++j) {
      scalar_type tmp = 0;

      // block row
      const ordinal_type cbegin = _rowptr(i), cend = _rowptr(i + 1);

      for (ordinal_type c = cbegin; c < cend; ++c) {
        const ordinal_type col = _colidx(c);
        for (ordinal_type jj = 0; jj < _blocksize; ++jj)
          tmp += _A(col, ii, jj) * _x(j, col, jj);
      }
      _y(j, i, ii) = tmp;
    }
  }

  void run(const block_crs_matrix_type A, const block_multi_vector_type x,
           const block_multi_vector_type y) {
    _rowptr = A.CrsGraph().rowptr;
    _colidx = A.CrsGraph().colidx;

    _blocksize = A.BlockSize();

    _A = A.Values();
    _x = x.Values();
    _y = y.Values();

    Kokkos::RangePolicy<ExecSpace> policy(0, _x.extent(1) * _blocksize);
    Kokkos::parallel_for("BlockCrsMatrixVectorProductByRow::run", policy,
                         *this);
  }
};

template <typename ExecSpace, typename ArrayLayout>
class BlockCrsMatrixVectorProductByBlockRow {
 public:
  typedef BlockCrsMatrix<ExecSpace, ArrayLayout> block_crs_matrix_type;
  typedef typename block_crs_matrix_type::crs_graph_type crs_graph_type;
  typedef BlockMultiVector<ExecSpace, ArrayLayout> block_multi_vector_type;

 private:
  ConstUnmanagedViewType<typename crs_graph_type::row_ptr_type> _rowptr;
  ConstUnmanagedViewType<typename crs_graph_type::col_idx_type> _colidx;

  ConstUnmanagedViewType<typename block_crs_matrix_type::value_array_type> _A;
  ConstUnmanagedViewType<typename block_multi_vector_type::value_array_type> _x;
  /**/ UnmanagedViewType<typename block_multi_vector_type::value_array_type> _y;

  ordinal_type _blocksize;

 public:
  // A thread maps to a row block of the matrix.
  // loop = m
  KOKKOS_INLINE_FUNCTION
  void operator()(const ordinal_type i) const {
    // loop over multivector colums
    const ordinal_type jend = _y.extent(0);
    for (ordinal_type j = 0; j < jend; ++j) {
      // set zero
      for (ordinal_type ii = 0; ii < _blocksize; ++ii) _y(j, i, ii) = 0;

      // block row
      const ordinal_type cbegin = _rowptr(i), cend = _rowptr(i + 1);

      for (ordinal_type c = cbegin; c < cend; ++c) {
        const ordinal_type col = _colidx(c);
        for (ordinal_type ii = 0; ii < _blocksize; ++ii) {
          scalar_type tmp = 0;
          for (ordinal_type jj = 0; jj < _blocksize; ++jj)
            tmp += _A(col, ii, jj) * _x(j, col, jj);
          _y(j, i, ii) += tmp;
        }
      }
    }
  }

  void run(const block_crs_matrix_type A, const block_multi_vector_type x,
           const block_multi_vector_type y) {
    _rowptr = A.CrsGraph().rowptr;
    _colidx = A.CrsGraph().colidx;

    _blocksize = A.BlockSize();

    _A = A.Values();
    _x = x.Values();
    _y = y.Values();

    Kokkos::RangePolicy<ExecSpace> policy(0, _x.extent(1));
    Kokkos::parallel_for(policy, *this);
  }
};

template <typename ExecSpace, typename ValueType, typename ArrayLayout>
class ExtractBlockTridiagMatrices {
 public:
  typedef ExecSpace exec_space;
  typedef ValueType value_type;
  typedef ArrayLayout array_layout;

  typedef StructuredBlock structured_block_mesh_type;
  typedef BlockCrsMatrix<exec_space, array_layout> block_crs_matrix_type;
  typedef typename block_crs_matrix_type::crs_graph_type crs_graph_type;
  typedef BlockTridiagMatrices<exec_space, value_type, array_layout>
      block_tridiag_matrices_type;

 private:
  structured_block_mesh_type _mesh;
  ordinal_type _blocksize;

  ConstUnmanagedViewType<typename crs_graph_type::row_ptr_type> _rowptr;
  ConstUnmanagedViewType<typename crs_graph_type::row_idx_type> _rowidx;
  ConstUnmanagedViewType<typename crs_graph_type::col_idx_type> _colidx;

  ConstUnmanagedViewType<typename block_crs_matrix_type::value_array_type> _A;
  /**/ UnmanagedViewType<typename block_tridiag_matrices_type::value_array_type>
      _TA, _TB, _TC;

 public:
  ExtractBlockTridiagMatrices(const structured_block_mesh_type mesh)
      : _mesh(mesh) {}

  template <typename TViewType, typename AViewType>
  KOKKOS_INLINE_FUNCTION void elementwise_copy(
      const TViewType &T, const AViewType &A, const ordinal_type ij,
      const ordinal_type k, const ordinal_type c,
      const ordinal_type blocksize) const {
    for (ordinal_type ii = 0; ii < blocksize; ++ii)
      for (ordinal_type jj = 0; jj < blocksize; ++jj)
        tdiag_val(T, ij, k, ii, jj) = A(c, ii, jj);
  }

  // A thread maps nonzero blocks
  KOKKOS_INLINE_FUNCTION
  void operator()(const ordinal_type c) const {
    const ordinal_type row = _rowidx[c], col = _colidx[c];

    ordinal_type ri, rj, rk, ci, cj, ck;
    _mesh.id2ijk(row, ri, rj, rk);
    _mesh.id2ijk(col, ci, cj, ck);

    if (ri == ci && rj == cj) {
      const ordinal_type ij = _mesh.ij2id(ri, rj);
      // consider connectivity to k-direction
      switch (rk - ck) {
        case 1: elementwise_copy(_TC, _A, ij, ck, c, _blocksize); break;
        case 0: elementwise_copy(_TA, _A, ij, rk, c, _blocksize); break;
        case -1: elementwise_copy(_TB, _A, ij, rk, c, _blocksize); break;
      }
    }
  }

  void run(const block_crs_matrix_type A, const block_tridiag_matrices_type T) {
    _rowptr = A.CrsGraph().rowptr;
    _rowidx = A.CrsGraph().rowidx;
    _colidx = A.CrsGraph().colidx;

    _A = A.Values();

    _TA = T.A();
    _TB = T.B();
    _TC = T.C();

    _blocksize = A.BlockSize();
    Kokkos::RangePolicy<ExecSpace> policy(0, _A.extent(0));
    Kokkos::parallel_for(policy, *this);
  }

  template <typename TViewType, typename AViewType>
  bool elementwise_check(const TViewType &T, const AViewType &A,
                         const ordinal_type ij, const ordinal_type k,
                         const ordinal_type c,
                         const ordinal_type blocksize) const {
    const auto eps = 1e2 * std::numeric_limits<scalar_type>::epsilon();
    for (ordinal_type ii = 0; ii < blocksize; ++ii)
      for (ordinal_type jj = 0; jj < blocksize; ++jj)
        if (Kokkos::ArithTraits<scalar_type>::abs(tdiag_val(T, ij, k, ii, jj) -
                                                  A(c, ii, jj)) >= eps)
          return false;
    return true;
  }

  bool check() const {
    auto rowptr = Kokkos::create_mirror_view(_rowptr);
    Kokkos::deep_copy(rowptr, _rowptr);
    auto colidx = Kokkos::create_mirror_view(_colidx);
    Kokkos::deep_copy(colidx, _colidx);
    auto TA = Kokkos::create_mirror_view(_TA);
    Kokkos::deep_copy(TA, _TA);
    auto TB = Kokkos::create_mirror_view(_TB);
    Kokkos::deep_copy(TB, _TB);
    auto TC = Kokkos::create_mirror_view(_TC);
    Kokkos::deep_copy(TC, _TC);
    auto A = Kokkos::create_mirror_view(_A);
    Kokkos::deep_copy(A, _A);

    const ordinal_type ijend = adjustDimension<value_type>(_mesh.ni * _mesh.nj),
                       kend  = _mesh.nk;

    assert(ijend == ordinal_type(TA.extent(0)));
    assert((kend - 0) == ordinal_type(TA.extent(1)));
    assert(ijend == ordinal_type(TB.extent(0)));
    assert((kend - 1) == ordinal_type(TB.extent(1)));
    assert(ijend == ordinal_type(TC.extent(0)));
    assert((kend - 1) == ordinal_type(TC.extent(1)));

    for (ordinal_type ij = 0; ij < ijend; ++ij) {
      ordinal_type i, j;
      _mesh.id2ij(ij, i, j);

      for (ordinal_type k = 0; k < kend; ++k) {
        const ordinal_type row = _mesh.ijk2id(i, j, k), idx_begin = rowptr[row],
                           idx_end = rowptr[row + 1];

        // check
        bool found[3] = {}, same[3] = {};
        for (ordinal_type idx = idx_begin; idx < idx_end; ++idx) {
          switch (row - colidx[idx]) {
            case 1:
              same[2]  = elementwise_check(TC, A, ij, k - 1, idx, _blocksize);
              found[2] = true;
              break;
            case 0:
              same[0]  = elementwise_check(TA, A, ij, k, idx, _blocksize);
              found[0] = true;
              break;
            case -1:
              same[1]  = elementwise_check(TB, A, ij, k, idx, _blocksize);
              found[1] = true;
              break;
          }
        }
        if (k == 0)
          assert(found[0] & same[0] && found[1] & same[1]);
        else if (k == (kend - 1))
          assert(found[0] & same[0] && found[2] & same[2]);
        else
          assert(found[0] & same[0] && found[1] & same[1] &&
                 found[2] & same[2]);
      }
    }
    return true;
  }
};

inline bool eq(const std::string &a, const char *const b1,
               const char *const b2 = 0) {
  return (a == std::string(b1) || (b2 && a == std::string(b2)) ||
          a == std::string("-") + std::string(b1));
}

// Command-line argument parser and holder.
template <typename ExecSpace>
struct Input {
  bool quiet, check;
  ordinal_type ni, nj, nk;
  ordinal_type bs;    // block size
  ordinal_type nrhs;  // #vectors in multivector
  ordinal_type opf, ops;
  StencilShape::Enum stencil_shape;

  Input(int argc, char **argv) {
    quiet = false;
    check = false;
    ni = nj = nk = 10;
    bs           = 5;
    nrhs         = 1;
    if (std::is_same<typename ExecSpace::memory_space,
                     Kokkos::HostSpace>::value) {
      opf = 0;
      ops = 0;  // range policy default
    } else {
      opf = 1;
      ops = 1;  // team is default
    }
    stencil_shape = StencilShape::cross;

    for (ordinal_type i = 1; i < argc; ++i) {
      const std::string &token = argv[i];
      if (eq(token, "-nijk"))
        ni = nj = nk = std::atoi(argv[++i]);
      else if (eq(token, "-ni"))
        ni = std::atoi(argv[++i]);
      else if (eq(token, "-nj"))
        nj = std::atoi(argv[++i]);
      else if (eq(token, "-nk"))
        nk = std::atoi(argv[++i]);
      else if (eq(token, "-bs"))
        bs = std::atoi(argv[++i]);
      else if (eq(token, "-nrhs"))
        nrhs = std::atoi(argv[++i]);
      else if (eq(token, "-opf"))
        opf = std::atoi(argv[++i]);
      else if (eq(token, "-ops"))
        ops = std::atoi(argv[++i]);
      else if (eq(token, "-c", "-check"))
        check = true;
    }
    if (nk <= 1) throw std::runtime_error("k dimension is <= 1; must be >= 2.");
    if (!quiet) print(std::cout);
  }

  void print(std::ostream &os) const {
    os << "<I> ni " << ni << " nj " << nj << " nk " << nk << " bs " << bs
       << " nrhs " << nrhs << " opf " << opf << " ops " << ops << " sc "
       << stencil_shape << "\n";
  }
};

}  // namespace Test
}  // namespace KokkosBatched
