#ifndef __JMATH__JMATRIXNS__
#define __JMATH__JMATRIXNS__

#include <vector>
#include <utility>

#include "./JMatrixND.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
  {
    /**
     * Type definition of permutation matrix.
     */
    typedef std::vector< std::pair<size_t,size_t> > JPermutationMatrix;


    /**
     * 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;
      const double* c1;
      const double* c2;
      const double* c3;

      const double* a0;
      const double* a1;
      const double* a2;
      const double* a3;

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

      double w0;
      double w1;
      double w2;
      double w3;	  

      JPermutationMatrix P;


      // 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.push_back(make_pair(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);
	  }
        }
      }


      // invert D

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

	double& pivot = (*this)(i,i);

	if (pivot == 0.0) {
	  THROW(JDivisionByZero, "Invert D zero pivot at " << i);
	}

	pivot = 1.0 / pivot;
      }


      // invert U

      JMatrixND_t& A = getInstance();

      A.set(*this);
      A.transpose();

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

	p0 = (*this)[row + 0]  +  row + 0;
	p1 = (*this)[row + 1]  +  row + 1;
	p2 = (*this)[row + 2]  +  row + 2;
	p3 = (*this)[row + 3]  +  row + 3;
	
	v0 = *(p0++);
	v1 = *(p1++);
	v2 = *(p2++);
	v3 = *(p3++);

	for (col = row + 1; col + 4 <= this->size(); ++col, ++p0, ++p1, ++p2, ++p3) {

	  w0 = (*p0) * -v0;
	  w1 = (*p1) * -v1;
	  w2 = (*p2) * -v2;
	  w3 = (*p3) * -v3;

	  c0 = (*this)[row + 0]  +  row + 0 + 1;
	  c1 = (*this)[row + 1]  +  row + 1 + 1;
	  c2 = (*this)[row + 2]  +  row + 2 + 1;
	  c3 = (*this)[row + 3]  +  row + 3 + 1;

	  a0 = A[col + 0]  +  row + 0 + 1;
	  a1 = A[col + 1]  +  row + 1 + 1;
	  a2 = A[col + 2]  +  row + 2 + 1;
	  a3 = A[col + 3]  +  row + 3 + 1;

	  for ( ; c0 + 4 < p0; c0 += 4, c1 += 4, c2 += 4, c3 += 4, a0 += 4, a1 += 4, a2 += 4, a3 += 4) {
	    w0 -= c0[0]*a0[0] + c0[1]*a0[1] + c0[2]*a0[2] + c0[3]*a0[3]; 
	    w1 -= c1[0]*a1[0] + c1[1]*a1[1] + c1[2]*a1[2] + c1[3]*a1[3]; 
	    w2 -= c2[0]*a2[0] + c2[1]*a2[1] + c2[2]*a2[2] + c2[3]*a2[3]; 
	    w3 -= c3[0]*a3[0] + c3[1]*a3[1] + c3[2]*a3[2] + c3[3]*a3[3]; 
	  }

	  for ( ; c0 != p0; ++c0, ++c1, ++c2, ++c3, ++a0, ++a1, ++a2, ++a3) {
	    w0 -= (*c0) * (*a0);
	    w1 -= (*c1) * (*a1);
	    w2 -= (*c2) * (*a2);
	    w3 -= (*c3) * (*a3);
	  }

	  *p0 = w0 * (*this)(col + 0, col + 0);
	  *p1 = w1 * (*this)(col + 1, col + 1);
	  *p2 = w2 * (*this)(col + 2, col + 2);
	  *p3 = w3 * (*this)(col + 3, col + 3);
	}

	// last 3 columns

	w0 = (*p0) * -v0;
	w1 = (*p1) * -v1;
	w2 = (*p2) * -v2;

	c0 = (*this)[row + 0]  +  row + 0 + 1;
	c1 = (*this)[row + 1]  +  row + 1 + 1;
	c2 = (*this)[row + 2]  +  row + 2 + 1;

	a0 = A[col + 0]  +  row + 0 + 1;
	a1 = A[col + 1]  +  row + 1 + 1;
	a2 = A[col + 2]  +  row + 2 + 1;

	for ( ; c0 + 4 < p0; c0 += 4, c1 += 4, c2 += 4, c3 += 4, a0 += 4, a1 += 4, a2 += 4, a3 += 4) {
	  w0 -= c0[0]*a0[0] + c0[1]*a0[1] + c0[2]*a0[2] + c0[3]*a0[3]; 
	  w1 -= c1[0]*a1[0] + c1[1]*a1[1] + c1[2]*a1[2] + c1[3]*a1[3]; 
	  w2 -= c2[0]*a2[0] + c2[1]*a2[1] + c2[2]*a2[2] + c2[3]*a2[3]; 
	}

	for ( ; c0 != p0; ++c0, ++c1, ++c2, ++a0, ++a1, ++a2) {
	  w0 -= (*c0) * (*a0);
	  w1 -= (*c1) * (*a1);
	  w2 -= (*c2) * (*a2);
	}
	  
	*p0 = w0 * (*this)(col + 0, col + 0);
	*p1 = w1 * (*this)(col + 1, col + 1);
	*p2 = w2 * (*this)(col + 2, col + 2);

	++col; ++p0; ++p1;

	// last 2 columns

	w0 = (*p0) * -v0;
	w1 = (*p1) * -v1;

	c0 = (*this)[row + 0]  +  row + 0 + 1;
	c1 = (*this)[row + 1]  +  row + 1 + 1;

	a0 = A[col + 0]  +  row + 0 + 1;
	a1 = A[col + 1]  +  row + 1 + 1;

	for ( ; c0 + 4 < p0; c0 += 4, c1 += 4, c2 += 4, c3 += 4, a0 += 4, a1 += 4, a2 += 4, a3 += 4) {
	  w0 -= c0[0]*a0[0] + c0[1]*a0[1] + c0[2]*a0[2] + c0[3]*a0[3]; 
	  w1 -= c1[0]*a1[0] + c1[1]*a1[1] + c1[2]*a1[2] + c1[3]*a1[3]; 
	}

	for ( ; c0 != p0; ++c0, ++c1, ++a0, ++a1) {
	  w0 -= (*c0) * (*a0);
	  w1 -= (*c1) * (*a1);
	}
	  
	*p0 = w0 * (*this)(col + 0, col + 0);
	*p1 = w1 * (*this)(col + 1, col + 1);

	++col; ++p0;

	// last column

	w0 = (*p0) * -v0;

	c0 = (*this)[row + 0]  +  row + 0 + 1;

	a0 = A[col + 0]  +  row + 0 + 1;

	for ( ; c0 + 4 < p0; c0 += 4, c1 += 4, c2 += 4, c3 += 4, a0 += 4, a1 += 4, a2 += 4, a3 += 4) {
	  w0 -= c0[0]*a0[0] + c0[1]*a0[1] + c0[2]*a0[2] + c0[3]*a0[3]; 
	}

	for ( ; c0 != p0; ++c0, ++a0) {
	  w0 -= (*c0) * (*a0);
	}
	  
	*p0 = w0 * (*this)(col + 0, col + 0);
      }

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

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

	v0 = *(p0++);

	for (col = row + 1; col != this->size(); ++col, ++p0) {

	  w0 = *p0;

	  w0 *= -v0;

	  c0 = (*this)[row + 0]  +  row + 1;

	  a0 = A[col + 0]  +  row + 0 + 1;

	  for ( ; c0 != p0; ++c0, ++a0) {
	    w0 -= (*c0) * (*a0);
	  }
	  
	  *p0 = w0 * (*this)(col + 0, col + 0);
	}
      }


      // solve A^-1 x L = U^-1

      double* buffer = A.data();

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

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

          buffer[i] = (*this)(i, col - 1);

	  (*this)(i, col - 1) = 0.0;
	}

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

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

	  c0 = p0 + 1;
	  c1 = p1 + 1;
	  c2 = p2 + 1;
	  c3 = p3 + 1;

	  a0 = buffer + col;

	  w0 = 0.0;
	  w1 = 0.0;
	  w2 = 0.0;
	  w3 = 0.0;

	  for (i = col; i + 4 <= this->size(); i += 4, c0 += 4, c1 += 4, c2 += 4, c3 += 4, a0 += 4) {
	    w0 += c0[0]*a0[0] + c0[1]*a0[1] + c0[2]*a0[2] + c0[3]*a0[3];
	    w1 += c1[0]*a0[0] + c1[1]*a0[1] + c1[2]*a0[2] + c1[3]*a0[3];
	    w2 += c2[0]*a0[0] + c2[1]*a0[1] + c2[2]*a0[2] + c2[3]*a0[3];
	    w3 += c3[0]*a0[0] + c3[1]*a0[1] + c3[2]*a0[2] + c3[3]*a0[3];
	  }

	  for ( ; i != this->size(); ++i, ++c0, ++c1, ++c2, ++c3, ++a0) {
	    w0 += (*c0) * (*a0);
	    w1 += (*c1) * (*a0);
	    w2 += (*c2) * (*a0);
	    w3 += (*c3) * (*a0);
	  }

	  *p0 -= w0;
	  *p1 -= w1;
	  *p2 -= w2;
	  *p3 -= w3;
	}

	for ( ; row < this->size(); ++row) {

	  p0 = (*this)[row] + col - 1;

	  c0 = p0 + 1;

	  a0 = buffer + col;

	  w0 = 0.0;

	  for (i = col; i + 4 <= this->size(); i += 4, c0 += 4, c1 += 4, c2 += 4, c3 += 4, a0 += 4) {
	    w0 += c0[0]*a0[0] + c0[1]*a0[1] + c0[2]*a0[2] + c0[3]*a0[3];
	  }

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

	  *p0 -= w0;
	}
      }


      for (JPermutationMatrix::reverse_iterator p = P.rbegin(); p != P.rend(); ++p) {
	cswap(p->first, p->second);
      }
    }


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

#endif