//============================================================================
//
// This file is part of the Shapes Generative Model Library (SGML) project.
//
// Copyright (c) 2010-2014 - Evangelos Kalogerakis, Siddhartha Chaudhuri (authors of the code) / Stanford University, UMass-Amherst
//
// SGML is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// SGML is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with SGML.  If not, see <http://www.gnu.org/licenses/>.
//
//============================================================================

#ifndef __SNL_S12_Matrix_hpp__
#define __SNL_S12_Matrix_hpp__

#include "S12_Vector.hpp"
#include <algorithm>
#include <stdexcept>
#include <cstring>
#include <iterator>
#include <cstdlib>
#include <cmath>
#include <iostream>
#include <boost/type_traits.hpp>

namespace SNL {
namespace S12 {

// Adapted from the Boost type_traits examples.

namespace Internal {

template <typename I1, typename I2, bool b>
I2
fastCopyImpl(I1 first, I1 last, I2 out, boost::integral_constant<bool, b> const &)
{
  while (first != last) *(out++) = *(first++);
  return out;
}

template <typename T>
T *
fastCopyImpl(T const * first, T const * last, T * out, boost::true_type const &)
{
   memcpy(out, first, (last - first) * sizeof(T));
   return out + (last - first);
}

// Same semantics as std::copy, calls memcpy where appropriate.
template <typename I1, typename I2, bool b>
I2
fastCopyBackwardImpl(I1 first, I1 last, I2 out, boost::integral_constant<bool, b> const &)
{
  while (last != first) *(--out) = *(--last);
  return out;
}

template <typename T>
T *
fastCopyBackwardImpl(T const * first, T const * last, T * out, boost::true_type const &)
{
   memmove(out, first, (last - first) * sizeof(T));
   return out;
}

template <typename MatrixT> void invertMatrix(MatrixT & mat, size_t * col_index, size_t * row_index, size_t * pivot);

} // namespace Internal

/**
 * A version of <tt>std::copy</tt> that calls <tt>memcpy</tt> where appropriate (if the class has a trivial assignment operator)
 * for speed.
 *
 * To take advantage of fast copying, specialize boost::has_trivial_assign to return true for the value type.
 */
template <typename I1, typename I2>
inline I2
fastCopy(I1 first, I1 last, I2 out)
{
  //
  // We can copy with memcpy if T has a trivial assignment operator,
  // and if the iterator arguments are actually pointers (this last
  // requirement we detect with overload resolution):
  //
  typedef typename std::iterator_traits<I1>::value_type value_type;
  return Internal::fastCopyImpl(first, last, out, boost::has_trivial_assign<value_type>());
}

/**
 * A version of <tt>std::copy_backward</tt> that calls <tt>memmove</tt> where appropriate (if the class has a trivial assignment
 * operator) for speed.
 *
 * To take advantage of fast copying, specialize boost::has_trivial_assign to return true for the value type.
 */
template <typename I1, typename I2>
inline I2
fastCopyBackward(I1 first, I1 last, I2 out)
{
  typedef typename std::iterator_traits<I1>::value_type value_type;
  return Internal::fastCopyBackwardImpl(first, last, out, boost::has_trivial_assign<value_type>());
}

enum MatrixLayout
{
  ROW_MAJOR,
  COLUMN_MAJOR
};

#define SNL_S12_alwaysAssertM(cond, msg) { if (!(cond)) throw std::runtime_error(msg); }

/** A standard dense 2D matrix, * either row-major or column-major. */
template <typename T, MatrixLayout L = ROW_MAJOR>
class /* THEA_API */ Matrix
{
  public:
    static MatrixLayout const Layout = L;  ///< The layout of the matrix (row-major or column-major).
    typedef T Value;                       ///< The type of values in the matrix.
    typedef T *       Iterator;            ///< Iterator over matrix elements.
    typedef T const * ConstIterator;       ///< Const iterator over matrix elements.

    /**
     * Constructs a matrix of a given size, filled with a given value. If the number of columns is omitted or zero, a square
     * matrix is created.
     */
    Matrix(size_t num_rows_ = 0, size_t num_cols_ = 0) : num_rows(0), num_cols(0), values(NULL), owns_memory(true)
    {
      resize(num_rows_, num_cols_ ? num_cols_ : num_rows_);
    }

    /** Constructs a matrix of a given size, filled with a given value. */
    Matrix(size_t num_rows_, size_t num_cols_, T const & fill_value) : num_rows(0), num_cols(0), values(NULL), owns_memory(true)
    {
      resize(num_rows_, num_cols_, fill_value);
    }

    /**
     * Constructs this matrix as a wrapper for an existing block of storage. The matrix thus created is <b>not resizable</b>,
     * and the memory block will <b>not be freed</b> when the Matrix object is destroyed.
     */
    Matrix(T * existing_data, size_t num_rows_, size_t num_cols_)
    : num_rows(num_rows_), num_cols(num_cols_ ? num_cols_ : num_rows_), values(existing_data), owns_memory(false)
    {}

    /** Construct a square diagonal matrix, given the values on the diagonal. */
    template <class VectorT>
    static Matrix fromDiagonal(VectorT const & diagonal)
    {
      Matrix m((size_t)diagonal.size(), (size_t)diagonal.size(), static_cast<T>(0));
      for (size_t i = 0; i < diagonal.size(); ++i)
        m((size_t)i, (size_t)i) = diagonal[i];

      return m;
    }

    /** Copy constructor. */
    Matrix(Matrix const & m) : num_rows(0), num_cols(0), values(NULL), owns_memory(true) { *this = m; }

    /** Destructor. */
    ~Matrix()
    {
      if (owns_memory) delete [] values;
    }

    /** Assignment operator. */
    Matrix & operator=(Matrix const & src)
    {
      if (owns_memory)
        resize(src.numRows(), src.numColumns());
      else
      {
        SNL_S12_alwaysAssertM(num_rows == src.numRows() && num_cols == src.numColumns(),
                      "Matrix: A wrapper matrix cannot be assigned a value of different dimensions");
      }

      fastCopy(src.begin(), src.end(), values);
      return *this;
    }

    /** Get the layout of the matrix (row or column major) */
    static MatrixLayout getLayout() { return L; }

    /** Get the number of rows in the matrix. */
    size_t numRows() const { return num_rows; }

    /** Get the number of columns in the matrix. */
    size_t numColumns() const { return num_cols; }

    /** Get the number of elements in the matrix. */
    size_t numElements() const { return num_rows * num_cols; }

    /** Check if the matrix is empty. */
    bool isEmpty() const { return num_rows <= 0 || num_cols <= 0; }

    /** Check if the matrix is square. */
    bool isSquare() const { return num_rows == num_cols; }

    /** Resize the matric to given dimensions. Prior data is discarded. */
    void resize(size_t num_rows_, size_t num_cols_)
    {
      SNL_S12_alwaysAssertM(owns_memory, "Matrix: This matrix does not own its storage block and can't be resized");

      if (num_rows != num_rows_ || num_cols != num_cols_)
      {
        size_t new_num_elems = num_rows_ * num_cols_;
        if (numElements() != new_num_elems)
        {
          delete [] values;

          if (new_num_elems > 0) values = new T[new_num_elems];
          else                   values = NULL;
        }

        num_rows = num_rows_;
        num_cols = num_cols_;
      }
    }

    /** Resize the matrix and set all entries to a given value. */
    void resize(size_t num_rows_, size_t num_cols_, T const & fill_value)
    {
      resize(num_rows_, num_cols_);
      fill(fill_value);
    }

    /** Element access. */
    T const & operator()(size_t row, size_t col) const
    {
      // Compiler should optimize to remove the conditional
      if (L == ROW_MAJOR)
        return values[row * num_cols + col];
      else
        return values[col * num_rows + row];
    }

    /** Element access. */
    T & operator()(size_t row, size_t col)
    {
      // Compiler should optimize to remove the conditional
      if (L == ROW_MAJOR)
        return values[row * num_cols + col];
      else
        return values[col * num_rows + row];
    }

    /** Access to a row (if row-major) or column (if column-major). */
    T const * operator[](std::size_t i) const
    {
      // Compiler should optimize to remove the conditional
      if (L == ROW_MAJOR)
        return &values[i * num_cols];
      else
        return &values[i * num_rows];
    }

    /** Access to a row (if row-major) or column (if column-major). */
    T * operator[](std::size_t i)
    {
      // Compiler should optimize to remove the conditional
      if (L == ROW_MAJOR)
        return &values[i * num_cols];
      else
        return &values[i * num_rows];
    }

    /** Get a row-vector of the matrix. */
    SNL_VECTOR(T) getRow(std::size_t row) const
    {
      SNL_VECTOR(T) r(num_cols);
      if (L == ROW_MAJOR)
      {
        T const * row_start = &values[0] + row * num_cols;
        fastCopy(row_start, row_start + num_cols, &r[0]);
      }
      else
      {
        for (std::size_t i = 0; i < num_cols; ++i)
          r[i] = (*this)(row, i);
      }

      return r;
    }

    /** Get a column-vector of the matrix. */
    SNL_VECTOR(T) getColumn(std::size_t col) const
    {
      SNL_VECTOR(T) c(num_rows);
      if (L == COLUMN_MAJOR)
      {
        T const * col_start = &values[0] + col * num_rows;
        fastCopy(col_start, col_start + num_rows, &c[0]);
      }
      else
      {
        for (std::size_t i = 0; i < num_rows; ++i)
          c[i] = (*this)(i, col);
      }

      return c;
    }

    /** Fill the matrix with a single value. */
    void fill(T const & value)
    {
      if (values)
      {
        for (T * vi = values, * e = values + numElements(); vi != e; ++vi)
          *vi = value;
      }
    }

    /** Get the storage block of the matrix. */
    T const * data() const { return values; }

    /** Get the storage block of the matrix. */
    T * data() { return values; }

    /** Get an iterator pointing to the beginning of the matrix. */
    ConstIterator begin() const { return values; }

    /** Get an iterator pointing to the beginning of the matrix. */
    Iterator begin() { return values; }

    /** Get an iterator pointing to one position past the end of the matrix. */
    ConstIterator end() const { return values + numElements(); }

    /** Get an iterator pointing to one position past the end of the matrix. */
    Iterator end() { return values + numElements(); }

    /**
     * Get the inverse of the matrix (assertion failure if matrix is not square). All computations are done using type T, so do
     * <b>not</b> call this function on integer matrices (built-in (POD) integer types will generate assertion failures).
     * Creates a new matrix for the result, so you might prefer in-place inversion using invert().
     */
    Matrix inverse() const { Matrix result = *this; result.invert(); return result; }

    /**
     * Invert the matrix (assertion failure if matrix is not square). All computations are done using type T, so do <b>not</b>
     * call this function on integer matrices (built-in (POD) integer types will generate assertion failures).
     */
    void invert()
    {
      SNL_S12_alwaysAssertM(this->isSquare(), "Matrix: Only square matrices can be inverted");

      SNL_VECTOR(size_t) col_index((size_t)num_rows), row_index((size_t)num_rows), pivot((size_t)num_rows);
      Internal::invertMatrix(*this, &col_index[0], &row_index[0], &pivot[0]);
    }

    /** Get the transpose of the matrix. Creates a new matrix for the result, so you might prefer transpose(Matrix &). */
    Matrix transpose() const
    {
      Matrix t(num_cols, num_rows);

      for (size_t r = 0; r < num_rows; ++r)
        for (size_t c = 0; c < num_cols; ++c)
          t(c, r) = (*this)(r, c);

      return t;
    }

    /**
     * Get the transpose of the matrix. No memory allocation/deallocation is required if the input and output matrices have the
     * same number of elements.
     *
     * @note The result matrix cannot be the same object as this matrix.
     */
    template <typename U, MatrixLayout ArgLayout> void transpose(Matrix<U, ArgLayout> & result) const
    {
      result.resize(num_cols, num_rows);

      if (L == ROW_MAJOR)
      {
        for (size_t r = 0; r < num_rows; ++r)
          for (size_t c = 0; c < num_cols; ++c)
            result(c, r) = (*this)(r, c);
      }
      else
      {
        for (size_t c = 0; c < num_cols; ++c)
          for (size_t r = 0; r < num_rows; ++r)
            result(c, r) = (*this)(r, c);
      }
    }

    /** Get the negation of this matrix (the matrix equal to this one except with every element negated). */
    Matrix operator-() const
    {
      Matrix neg(num_rows, num_cols);
      T * v2 = neg.values;
      for (T const * v1 = values, * e1 = values + numElements(); v1 != e1; ++v1, ++v2)
        *v2 = -(*v1);
    }

    /** Addition (per-component). */
    Matrix operator+(Matrix const & rhs) const
    {
      SNL_S12_alwaysAssertM(num_rows == rhs.num_rows && num_cols == rhs.num_cols,
                    "Matrix: Can't add matrices of different dimensions");

      Matrix result(num_rows, num_cols);
      T * u = result.values;
      for (T const * v1 = values, * v2 = rhs.values, * e1 = values + numElements(); v1 != e1; ++v1, ++v2, ++u)
        *u = *v1 + *v2;

      return result;
    }

    /** Add-and-assign. */
    Matrix & operator+=(Matrix const & rhs)
    {
      SNL_S12_alwaysAssertM(num_rows == rhs.num_rows && num_cols == rhs.num_cols,
                    "Matrix: Can't add matrices of different dimensions");

      T const * v2 = rhs.values;
      for (T * v1 = values, * e1 = values + numElements(); v1 != e1; ++v1, ++v2)
        *v1 += *v2;

      return *this;
    }

    /** Subtraction (per-component). */
    Matrix operator-(Matrix const & rhs) const
    {
      SNL_S12_alwaysAssertM(num_rows == rhs.num_rows && num_cols == rhs.num_cols,
                    "Matrix: Can't subtract matrices of different dimensions");

      Matrix result(num_rows, num_cols);
      T * u = result.values;
      for (T const * v1 = values, * v2 = rhs.values, * e1 = values + numElements(); v1 != e1; ++v1, ++v2, ++u)
        *u = *v1 - *v2;

      return result;
    }

    /** Subtract-and-assign. */
    Matrix & operator-=(Matrix const & rhs)
    {
      SNL_S12_alwaysAssertM(num_rows == rhs.num_rows && num_cols == rhs.num_cols,
                    "Matrix: Can't add matrices of different dimensions");

      T const * v2 = rhs.values;
      for (T * v1 = values, * e1 = values + numElements(); v1 != e1; ++v1, ++v2)
        *v1 -= *v2;

      return *this;
    }

    /** Post-multiply by a scalar. */
    Matrix operator*(T const & s) const
    {
      Matrix result(num_rows, num_cols);
      T * u = result.values;
      for (T const * v = values, * e = values + numElements(); v != e; ++v, ++u)
        *u = s * (*v);

      return result;
    }

    /** Post-multiply by a scalar and assign. */
    Matrix & operator*(T const & s)
    {
      for (T * v = values, * e = values + numElements(); v != e; ++v)
        (*v) *= s;

      return *this;
    }

    /** Multiplication. */
    Matrix operator*(Matrix const & rhs)
    {
      SNL_S12_alwaysAssertM(num_cols == rhs.num_rows, "Matrix: Matrices don't have compatible dimensions for multiplication");

      Matrix result(num_rows, rhs.num_cols);
      for (size_t i = 0; i < num_rows; ++i)
        for (size_t k = 0; k < rhs.num_cols; ++k)
        {
          result(i, k) = (*this)(i, 0) * rhs(0, k);

          for (size_t j = 1; j < num_cols; ++j)
            result(i, k) += (*this)(i, j) * rhs(j, k);
        }

      return result;
    }

    /**
     * Multiply a vector \a v by this matrix (M), yielding the product \a result = M * \a v.
     *
     * @param v The vector to be multiplied, must contain num_cols elements.
     * @param result The result vector, must be preallocated to num_rows elements.
     */
    template <typename U> void postmulVector(U const * v, U * result) const
    {
      debugAssertM(v, "Matrix: Vector to be multiplied cannot be null");
      debugAssertM(result, "Matrix: Result vector of multiplication cannot be null");

      if (L == ROW_MAJOR)
      {
        for (size_t r = 0; r < num_rows; ++r)
        {
          result[r] = static_cast<U>(0);
          for (size_t c = 0; c < num_cols; ++c)
            result[r] += static_cast<U>((*this)(r, c) * v[c]);
        }
      }
      else
      {
        for (size_t r = 0; r < num_rows; ++r)
          result[r] = static_cast<U>(0);

        for (size_t c = 0; c < num_cols; ++c)
          for (size_t r = 0; r < num_rows; ++r)
            result[r] += static_cast<U>((*this)(r, c) * v[c]);
      }
    }

    /**
     * Multiply a vector \a v by this matrix (M), yielding the product \a result = \a v * M.
     *
     * @param v The vector to be multiplied, must contain num_rows elements.
     * @param result The result vector, must be preallocated to num_cols elements.
     */
    template <typename U> void premulVector(U const * v, U * result) const
    {
      debugAssertM(v, "Matrix: Vector to be multiplied cannot be null");
      debugAssertM(result, "Matrix: Result vector of multiplication cannot be null");

      if (L == ROW_MAJOR)
      {
        for (size_t c = 0; c < num_cols; ++c)
          result[c] = static_cast<U>(0);

        for (size_t r = 0; r < num_rows; ++r)
          for (size_t c = 0; c < num_cols; ++c)
            result[c] += static_cast<U>(v[r] * (*this)(r, c));
      }
      else
      {
        for (size_t c = 0; c < num_cols; ++c)
        {
          result[c] = static_cast<U>(0);
          for (size_t r = 0; r < num_rows; ++r)
            result[c] += static_cast<U>(v[r] * (*this)(r, c));
        }
      }
    }

    /**
     * Utility function for efficiently computing  a matrix-vector product (w = M * v). This lets this matrix be passed directly
     * to ARPACK++. The implementation simply calls postmulVector().
     *
     * @see postmulVector()
     */
    template <typename U> void MultMv(U const * v, U * w) const { postmulVector(v, w); }

    /** Return a matrix containing the component-wise minima of this matrix and another. */
    Matrix min(Matrix const & other) const
    {
      Matrix result(num_rows, num_cols);
      for (size_t r = 0; r < num_rows; ++r)
        for (size_t c = 0; c < num_cols; ++c)
          result(r, c) = std::min((*this)(r, c), other(r, c));

      return result;
    }

    /** Return a matrix containing the component-wise maxima of this matrix and another. */
    Matrix max(Matrix const & other) const
    {
      Matrix result(num_rows, num_cols);
      for (size_t r = 0; r < num_rows; ++r)
        for (size_t c = 0; c < num_cols; ++c)
          result(r, c) = std::max((*this)(r, c), other(r, c));

      return result;
    }

  private:
    size_t num_rows, num_cols;
    T * values;
    bool owns_memory;

}; // class Matrix

/** Pre-multiply by a scalar. */
template <typename T, MatrixLayout L>
Matrix<T, L>
operator*(T const & s, Matrix<T, L> const & m)
{
  return m * s;
}

namespace Internal {

/**
 * <b>[Internal]</b> Invert a square matrix in place. This function is called by various matrix classes and should not be used
 * directly. Code borrowed from G3D::Matrix::inverseInPlaceGaussJordan(), originally from Numerical Recipes in C (gaussj). The
 * integer arrays \a pivot, \a row_index, and \a col_index are used for bookkeeping on the pivoting and must be
 * preallocated to n elements each, where n is the size (number of rows or number of columns) in the matrix to be inverted.
 *
 * Depending on which preprocessor flags are defined, this function inverts addressable matrices (in which case it uses virtual
 * accessors), or matrices with non-virtual element access via operator(). The latter is expected to be faster.
 *
 * @param mat Input matrix.
 * @param col_index A book-keeping array. Should be preallocated to n elements.
 * @param row_index A book-keeping array. Should be preallocated to n elements.
 * @param pivot A book-keeping array. Should be preallocated to n elements.
 */
template <typename MatrixT> /* THEA_DLL_LOCAL */
void
invertMatrix(MatrixT & mat, size_t * col_index, size_t * row_index, size_t * pivot)
{
  typedef typename MatrixT::Value Value;

  SNL_S12_alwaysAssertM(mat.isSquare(), "Can't invert non-square matrices");
  SNL_S12_alwaysAssertM(!boost::is_integral<Value>::value, "Can't invert integer matrices");

  size_t n = mat.numRows();

  // Initialize the pivot array to default values.
  static size_t const NO_PIVOT = -1;
  for (size_t i = 0; i < n; ++i)
    pivot[i] = NO_PIVOT;

  size_t col = 0, row = 0;

  // This is the main loop over the columns to be reduced
  //
  // Loop over the columns
  for (size_t c = 0; c < n; ++c)
  {
    // Find the largest element and use that as a pivot
    Value largest_magnitude = 0;

    // This is the outer loop of the search for a pivot element
    for (size_t r = 0; r < n; ++r)
    {
      // Unless we've already found the pivot, keep going
      if (pivot[r] != 0)
      {
        // Find the largest pivot
        for (size_t k = 0; k < n; ++k)
        {
          if (pivot[k] == NO_PIVOT)
          {
            Value mag = std::abs(mat(r, k));
            if (mag >= largest_magnitude)
            {
              largest_magnitude = mag;
              row = r; col = k;
            }
          }
        }
      }
    }

    pivot[col]++;

    // Interchange columns so that the pivot element is on the diagonal (we'll have to undo this at the end)
    if (row != col)
      for (size_t k = 0; k < n; ++k)
        std::swap(mat(row, k), mat(col, k));

    // The pivot is now at [row, col]
    row_index[c] = row;
    col_index[c] = col;

    Value piv = mat(col, col);
    if (piv == 0)  // FIXME: A more robust comparison?
    {
      std::cerr << "invertMatrix: Singular matrix:" << std::endl;
      for (size_t r = 0; r < n; ++r)
      {
        for (size_t c = 0; c < n; ++c)
          std::cerr << "  " << mat(r, c);

        std::cerr << std::endl;
      }

      throw std::runtime_error("invertMatrix: Matrix is singular");
    }

    // Divide everything by the pivot (avoid computing the division multiple times).
    Value pivot_inverse = 1 / piv;
    mat(col, col) = 1;

    for (size_t k = 0; k < n; ++k)
      mat(col, k) *= pivot_inverse;

    // Reduce all rows
    for (size_t r = 0; r < n; ++r)
    {
      // Skip over the pivot row
      if (r != col)
      {
        Value old_value = mat(r, col);
        mat(r, col) = 0;

        for (size_t k = 0; k < n; ++k)
          mat(r, k) -= mat(col, k) * old_value;
      }
    }
  }

  // Put the columns back in the correct locations
  for (int i = (int)(n - 1); i >= 0; --i)  // int to handle case when i is -1
    if (row_index[i] != col_index[i])
      for (size_t k = 0; k < n; ++k)
        std::swap(mat(k, row_index[i]), mat(k, col_index[i]));
}

} // namespace Internal

#undef SNL_S12_alwaysAssertM

#define SNL_MATRIX(T) SNL::S12::Matrix<T, SNL::S12::ROW_MAJOR>

} // namespace S12
} // namespace SNL

#endif
