/*

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

*/

#ifndef GOOSEFEM_MATRIXDIAGONALPARTITIONED_H
#define GOOSEFEM_MATRIXDIAGONALPARTITIONED_H

#include "config.h"

namespace GooseFEM {

class MatrixDiagonalPartitioned {
public:
    // Constructors
    MatrixDiagonalPartitioned() = default;

    MatrixDiagonalPartitioned(
        const xt::xtensor<size_t, 2>& conn,
        const xt::xtensor<size_t, 2>& dofs,
        const xt::xtensor<size_t, 1>& iip);

    // 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
    size_t nnu() const;   // number of unknown DOFs
    size_t nnp() const;   // number of prescribed DOFs

    // DOF lists
    xt::xtensor<size_t, 2> dofs() const; // DOFs
    xt::xtensor<size_t, 1> iiu() const;  // unknown DOFs
    xt::xtensor<size_t, 1> iip() const;  // prescribed DOFs

    // 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;

    void dot_u(
        const xt::xtensor<double, 1>& x_u,
        const xt::xtensor<double, 1>& x_p,
        xt::xtensor<double, 1>& b_u) const;

    void dot_p(
        const xt::xtensor<double, 1>& x_u,
        const xt::xtensor<double, 1>& x_p,
        xt::xtensor<double, 1>& b_p) const;

    // Solve:
    // x_u = A_uu \ ( b_u - A_up * x_p ) = A_uu \ b_u
    void solve(const xt::xtensor<double, 2>& b, xt::xtensor<double, 2>& x); // modified with "x_u"

    void solve(const xt::xtensor<double, 1>& b, xt::xtensor<double, 1>& x); // modified with "x_u"

    void solve_u(
        const xt::xtensor<double, 1>& b_u,
        const xt::xtensor<double, 1>& x_p,
        xt::xtensor<double, 1>& x_u);

    // Get right-hand-size for corresponding to the prescribed DOFs:
    // b_p = A_pu * x_u + A_pp * x_p = A_pp * x_p
    void reaction(
        const xt::xtensor<double, 2>& x, xt::xtensor<double, 2>& b) const; // modified with "b_p"

    void reaction(
        const xt::xtensor<double, 1>& x, xt::xtensor<double, 1>& b) const; // modified with "b_p"

    void reaction_p(
        const xt::xtensor<double, 1>& x_u,
        const xt::xtensor<double, 1>& x_p,
        xt::xtensor<double, 1>& b_p) const;

    // Return matrix as diagonal matrix (column)
    xt::xtensor<double, 1> Todiagonal() 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, 1> Dot_u(
        const xt::xtensor<double, 1>& x_u, const xt::xtensor<double, 1>& x_p) const;

    xt::xtensor<double, 1> Dot_p(
        const xt::xtensor<double, 1>& x_u, const xt::xtensor<double, 1>& x_p) const;

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

    xt::xtensor<double, 1> Solve_u(
        const xt::xtensor<double, 1>& b_u, const xt::xtensor<double, 1>& x_p);

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

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

    xt::xtensor<double, 1> Reaction_p(
        const xt::xtensor<double, 1>& x_u, const xt::xtensor<double, 1>& x_p) const;

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

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

    // Bookkeeping
    xt::xtensor<size_t, 2> m_conn; // connectivity                      [nelem, nne ]
    xt::xtensor<size_t, 2> m_dofs; // DOF-numbers per node              [nnode, ndim]
    xt::xtensor<size_t, 2> m_part; // DOF-numbers per node, renumbered  [nnode, ndim]
    xt::xtensor<size_t, 1> m_iiu;  // DOF-numbers that are unknown      [nnu]
    xt::xtensor<size_t, 1> m_iip;  // DOF-numbers that are prescribed   [nnp]

    // 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
    size_t m_nnu;   // number of unknown DOFs
    size_t m_nnp;   // number of prescribed DOFs

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

} // namespace GooseFEM

#include "MatrixDiagonalPartitioned.hpp"

#endif
