/* =================================================================================================

(c - GPLv3) T.W.J. de Geus (Tom) | tom@geus.me | www.geus.me | github.com/tdegeus/GooseFEM

================================================================================================= */

#ifndef GOOSEFEM_MATRIXDIAGONAL_H
#define GOOSEFEM_MATRIXDIAGONAL_H

// -------------------------------------------------------------------------------------------------

#include "config.h"

// =================================================================================================

namespace GooseFEM {

// -------------------------------------------------------------------------------------------------

class MatrixDiagonal
{
public:

  // Constructors

  MatrixDiagonal() = default;

  MatrixDiagonal(
    const xt::xtensor<size_t,2>& conn,
    const xt::xtensor<size_t,2>& dofs);

  // Dimensions

  size_t nelem() const; // number of elements
  size_t nne()   const; // number of nodes per element
  size_t nnode() const; // number of nodes
  size_t ndim()  const; // number of dimensions
  size_t ndof()  const; // number of DOFs

  // DOF lists

  xt::xtensor<size_t,2> dofs() const; // DOFs

  // Set matrix components

  void set(const xt::xtensor<double,1>& A);

  // assemble from matrices stored per element [nelem, nne*ndim, nne*ndim]
  // WARNING: ignores any off-diagonal terms

  void assemble(const xt::xtensor<double,3>& elemmat);

  // Dot-product:
  // b_i = A_ij * x_j

  void dot(
    const xt::xtensor<double,2>& x,
          xt::xtensor<double,2>& b) const;

  void dot(
    const xt::xtensor<double,1>& x,
          xt::xtensor<double,1>& b) const;

  // Solve:
  // x = A \ b

  void solve(
    const xt::xtensor<double,2>& b,
          xt::xtensor<double,2>& x);

  void solve(
    const xt::xtensor<double,1>& b,
          xt::xtensor<double,1>& x);

  // Return matrix as diagonal matrix (column)

  xt::xtensor<double,1> AsDiagonal() const;

  // Auto-allocation of the functions above

  xt::xtensor<double,2> Dot(const xt::xtensor<double,2>& x) const;

  xt::xtensor<double,1> Dot(const xt::xtensor<double,1>& x) const;

  xt::xtensor<double,2> Solve(const xt::xtensor<double,2>& b);

  xt::xtensor<double,1> Solve(const xt::xtensor<double,1>& b);

private:

  // The diagonal matrix, and its inverse (re-used to solve different RHS)
  xt::xtensor<double,1> m_A;
  xt::xtensor<double,1> m_inv;

  // Signal changes to data compare to the last inverse
  bool m_factor=false;

  // Bookkeeping
  xt::xtensor<size_t,2> m_conn; // connectivity         [nelem, nne ]
  xt::xtensor<size_t,2> m_dofs; // DOF-numbers per node [nnode, ndim]

  // Dimensions
  size_t m_nelem; // number of elements
  size_t m_nne;   // number of nodes per element
  size_t m_nnode; // number of nodes
  size_t m_ndim;  // number of dimensions
  size_t m_ndof;  // number of DOFs

  // Compute inverse (automatically evaluated by "solve")
  void factorize();
};

// -------------------------------------------------------------------------------------------------

} // namespace ...

// =================================================================================================

#include "MatrixDiagonal.hpp"

// =================================================================================================

#endif
