#ifndef __JMATH__JMATRIXNS__
#define __JMATH__JMATRIXNS__

#include <vector>
#include <utility>

#include "JMath/JMatrixND.hh"
#include "JMath/JVectorND.hh"
#include "JLang/JException.hh"


/**
 * \author mdejong
 */

namespace JMATH {}
namespace JPP { using namespace JMATH; }

namespace JMATH {

  using JLANG::JDivisionByZero;
  using JLANG::JIndexOutOfRange;


  /**
   * N x N symmetric matrix
   */
  struct JMatrixNS :
    public JMatrixND
  {
    /**
     * Default constructor.
     */
    JMatrixNS() :
      JMatrixND()
    {}


    /**
     * Constructor (identity matrix).
     *
     * \param  size     dimension
     */
    JMatrixNS(const size_t size) :
      JMatrixND(size)
    {}


    /**
     * Contructor.
     *
     * \param   A      matrix
     */
    JMatrixNS(const JMatrixND& A) :
      JMatrixND(A)
    {}


    /**
     * Assignment operator.
     *
     * \param  A        matrix
     */
    JMatrixNS& operator=(const JMatrixNS& A)
    {
      this->set(A);

      return *this;
    }


    /**
     * Invert matrix according LDU decomposition.
     */
    void invert()
    {
      using namespace std;

      size_t  i, row, col, root;

      double* p0;
      double* p1;
      double* p2;
      double* p3;

      const double* c0;

      double v0;
      double v1;
      double v2;
      double v3;

      double w0;

      P.resize(this->size());

      // LDU decomposition

      for (root = 0; root != this->size(); ++root) {

	double value = (*this)(root, root);
	size_t pivot = root;

	for (row = root; ++row != this->size(); ) {
	  if (fabs((*this)(row, root)) > fabs(value)) {
	    value = (*this)(row, root);
	    pivot = row;
	  }
	}

	if (value == 0.0) {
	  THROW(JDivisionByZero, "LDU decomposition zero pivot at " << root);
	}

	if (pivot != root) {
	  this->rswap(root, pivot);
	}

	P[root] = pivot;

        for (row = root + 1; row + 4 <= this->size(); row += 4) {            // process rows by four

          p0 = (*this)[row + 0]  +  root;
          p1 = (*this)[row + 1]  +  root;
          p2 = (*this)[row + 2]  +  root;
          p3 = (*this)[row + 3]  +  root;

          v0 = (*p0++ /= value);
          v1 = (*p1++ /= value);
          v2 = (*p2++ /= value);
          v3 = (*p3++ /= value);

	  c0 = (*this)[root] + root + 1;

          for (i = root + 1; i + 4 <= this->size(); i += 4, p0 += 4, p1 += 4, p2 += 4, p3 += 4, c0 += 4) {
            p0[0] -= v0 * c0[0];  p0[1] -= v0 * c0[1];  p0[2] -= v0 * c0[2];  p0[3] -= v0 * c0[3];
            p1[0] -= v1 * c0[0];  p1[1] -= v1 * c0[1];  p1[2] -= v1 * c0[2];  p1[3] -= v1 * c0[3];
            p2[0] -= v2 * c0[0];  p2[1] -= v2 * c0[1];  p2[2] -= v2 * c0[2];  p2[3] -= v2 * c0[3];
            p3[0] -= v3 * c0[0];  p3[1] -= v3 * c0[1];  p3[2] -= v3 * c0[2];  p3[3] -= v3 * c0[3];
          }

	  for ( ; i != this->size(); ++i, ++p0, ++p1, ++p2, ++p3, ++c0) {
	    *p0 -= v0 * (*c0);
	    *p1 -= v1 * (*c0);
	    *p2 -= v2 * (*c0);
	    *p3 -= v3 * (*c0);
	  }
	}

        for ( ; row != this->size(); ++row) {	                             // remaining rows

          p0 = (*this)[row + 0]  +  root;

          v0 = (*p0++ /= value);

	  c0 = (*this)[root] + root + 1;

          for (i = root + 1; i + 4 <= this->size(); i += 4, p0 += 4, c0 += 4) {
            *p0 -= v0 * c0[0];  p0[1] -= v0 * c0[1];  p0[2] -= v0 * c0[2];  p0[3] -= v0 * c0[3];
	  }

	  for ( ; i != this->size(); ++i, ++p0, ++c0) {
	    *p0 -= v0 * (*c0);
	  }
        }
      }


      JMatrixND_t& A = getInstance();

      A.reset();

      for (i = 0; i != this->size(); ++i) {
	A(i,i) = 1.0;
      }


      // forward substitution

      col = 0;
      
      for ( ; col + 4 <= this->size(); col += 4) {
	  
	for (row = 0; row != this->size(); ++row) {

	  p0 = A[col + 0];
	  p1 = A[col + 1];
	  p2 = A[col + 2];
	  p3 = A[col + 3];

	  i  = P[row];

	  v0 = p0[i];
	  v1 = p1[i];
	  v2 = p2[i];
	  v3 = p3[i];

	  p0[i] = p0[row];
	  p1[i] = p1[row];
	  p2[i] = p2[row];
	  p3[i] = p3[row];

	  for (i = 0, c0 = (*this)[row]; i != row; ++c0, ++p0, ++p1, ++p2, ++p3, ++i) {
	    v0 -= (*c0) * (*p0);
	    v1 -= (*c0) * (*p1);
	    v2 -= (*c0) * (*p2);
	    v3 -= (*c0) * (*p3);
	  }

	  A[col + 0][row] = v0;
	  A[col + 1][row] = v1;
	  A[col + 2][row] = v2;
	  A[col + 3][row] = v3;
	}
      }

      for ( ; col != this->size(); ++col) {

	for (row = 0; row != this->size(); ++row) {

	  p0 = A[col];
	  
	  i  = P[row];

	  v0 = p0[i];

	  p0[i] = p0[row];

	  for (i = 0, c0 = (*this)[row] + i; i != row; ++c0, ++p0, ++i) {
	    v0 -= (*c0) * (*p0);
	  }

	  A[col][row] = v0;
	}
      }

      // backward substitution

      col = 0;
      
      for ( ; col + 4 <= this->size(); col += 4) {

	for (int row = this->size() - 1; row >= 0; --row) {

	  p0 = A[col + 0] + row;
	  p1 = A[col + 1] + row;
	  p2 = A[col + 2] + row;
	  p3 = A[col + 3] + row;

	  v0 = *p0;
	  v1 = *p1;
	  v2 = *p2;
	  v3 = *p3;

	  ++p0;
	  ++p1;
	  ++p2;
	  ++p3;

	  for (i = row + 1, c0 = (*this)[row] + i; i != this->size(); ++c0, ++p0, ++p1, ++p2, ++p3, ++i) {
	    v0 -= (*c0) * (*p0);
	    v1 -= (*c0) * (*p1);
	    v2 -= (*c0) * (*p2);
	    v3 -= (*c0) * (*p3);
	  }

	  w0 = 1.0 / (*this)[row][row];

	  A[col + 0][row] = v0 * w0;
	  A[col + 1][row] = v1 * w0;
	  A[col + 2][row] = v2 * w0;
	  A[col + 3][row] = v3 * w0;
	}
      }
     
      for ( ; col != this->size(); ++col) {

	for (int row = this->size() - 1; row >= 0; --row) {

	  p0 = A[col] + row;

	  v0 = *p0;

	  ++p0;

	  for (i = row + 1, c0 = (*this)[row] + i; i != this->size(); ++c0, ++p0, ++i) {
	    v0 -= (*c0) * (*p0);
	  }

	  w0 = 1.0 / (*this)[row][row];

	  A[col][row] = v0 * w0;
	}
      }

      A.transpose();

      this->swap(A);
    }


    /**
     * Get solution of equation <tt>A x = b</tt>.
     *
     * \param  u        column vector; b on input, x on output(I/O)
     */
    template<class JVectorND_t>
    void solve(JVectorND_t& u)
    {
      using namespace std;

      size_t  i, row, root;

      double* p0;
      double* p1;
      double* p2;
      double* p3;

      const double* c0;

      double v0;
      double v1;
      double v2;
      double v3;

      P.resize(this->size());

      // LDU decomposition

      for (root = 0; root != this->size(); ++root) {

	double value = (*this)(root, root);
	size_t pivot = root;

	for (row = root; ++row != this->size(); ) {
	  if (fabs((*this)(row, root)) > fabs(value)) {
	    value = (*this)(row, root);
	    pivot = row;
	  }
	}

	if (value == 0.0) {
	  THROW(JDivisionByZero, "LDU decomposition zero pivot at " << root);
	}

	if (pivot != root) {
	  this->rswap(root, pivot);
	}

	P[root] = pivot;

        for (row = root + 1; row + 4 <= this->size(); row += 4) {            // process rows by four

          p0 = (*this)[row + 0]  +  root;
          p1 = (*this)[row + 1]  +  root;
          p2 = (*this)[row + 2]  +  root;
          p3 = (*this)[row + 3]  +  root;

          v0 = (*p0++ /= value);
          v1 = (*p1++ /= value);
          v2 = (*p2++ /= value);
          v3 = (*p3++ /= value);

	  c0 = (*this)[root] + root + 1;

          for (i = root + 1; i + 4 <= this->size(); i += 4, p0 += 4, p1 += 4, p2 += 4, p3 += 4, c0 += 4) {
            p0[0] -= v0 * c0[0];  p0[1] -= v0 * c0[1];  p0[2] -= v0 * c0[2];  p0[3] -= v0 * c0[3];
            p1[0] -= v1 * c0[0];  p1[1] -= v1 * c0[1];  p1[2] -= v1 * c0[2];  p1[3] -= v1 * c0[3];
            p2[0] -= v2 * c0[0];  p2[1] -= v2 * c0[1];  p2[2] -= v2 * c0[2];  p2[3] -= v2 * c0[3];
            p3[0] -= v3 * c0[0];  p3[1] -= v3 * c0[1];  p3[2] -= v3 * c0[2];  p3[3] -= v3 * c0[3];
          }

	  for ( ; i != this->size(); ++i, ++p0, ++p1, ++p2, ++p3, ++c0) {
	    *p0 -= v0 * (*c0);
	    *p1 -= v1 * (*c0);
	    *p2 -= v2 * (*c0);
	    *p3 -= v3 * (*c0);
	  }
	}

        for ( ; row != this->size(); ++row) {	                             // remaining rows

          p0 = (*this)[row + 0]  +  root;

          v0 = (*p0++ /= value);

	  c0 = (*this)[root] + root + 1;

          for (i = root + 1; i + 4 <= this->size(); i += 4, p0 += 4, c0 += 4) {
            *p0 -= v0 * c0[0];  p0[1] -= v0 * c0[1];  p0[2] -= v0 * c0[2];  p0[3] -= v0 * c0[3];
	  }

	  for ( ; i != this->size(); ++i, ++p0, ++c0) {
	    *p0 -= v0 * (*c0);
	  }
        }
      }

      // forward substitution

      for (row = 0; row != this->size(); ++row) {

	i      = P[row];
	v0     = u[i];
	u[i]   = u[row];

	for (i = 0, c0 = (*this)[row] + i; i != row; ++c0, ++i) {
	  v0 -= *c0 * u[i];
	}

	u[row] = v0;
      }

      // backward substitution

      for (int row = this->size() - 1; row >= 0; --row) {

	v0     = u[row];

	for (i = row + 1, c0 = (*this)[row] + i; i < this->size(); ++c0, ++i) {
	  v0 -= *c0 * u[i];
	}

	u[row] = v0 / (*this)[row][row];
      }
    }


    /**
     * Update inverted matrix at given diagonal element.
     *
     * If A is the original matrix and A' is such that for each <tt>(i,j) != (k,k)</tt>:
     * <pre>
     *       A'[i][j] = A[i][j]
     *       A'[k][k] = A[k][k] + value
     * </pre>
     * then <tt>JMatrixNS::update(k, value)</tt> will modify the already inverted matrix
     * of A such that it will be the inverse of A'.
     *
     * See also <a href="https://arxiv.org/abs/1708.07622"> arXiv</a>.\n
     * This algorithm can be considered a special case of the Sherman–Morrison formula.
     *
     * \param  k        index of diagonal element
     * \param  value    value to add
     */
    void update(const size_t k, const double value)
    {
      if (k >= this->size()) {
        THROW(JIndexOutOfRange, "JMatrixNS::update(): index out of range " << k);
      }

      size_t  i, row;
      double* col;

      const double del  =  (*this)(k,k);
      const double din  =  1.0 / (1.0 + value * del);

      double* root = (*this)[k];

      for (row = 0; row != this->size(); ++row) {

        if (row != k) {

          const double val = (*this)(row, k);

          for (i = 0, col = (*this)[row]; i != this->size(); ++i, ++col) {
            if (i != k) {
              (*col) -= val * root[i] * value * din;
            }
          }

          (*this)(row, k) *= din;
        }
      }

      for (i = 0, col = root; i != this->size(); ++i, ++col) {
        if (i != k) {
          (*col) *= din;
        }
      }

      root[k] = del * din;
    }

  private:
    std::vector<int> P;      //!< permutation matrix
  };
}

#endif