/**********************************************************

This software is part of J.-S. Caux's ABACUS library.

Copyright (c) J.-S. Caux.

-----------------------------------------------------------

File:  ABACUS_Matrix.h

Purpose:  Declares square and rectangular matrix classes.

***********************************************************/

#ifndef ABACUS_MATRIX_H
#define ABACUS_MATRIX_H

namespace ABACUS {


  // CLASS DEFINITIONS

  template <class T>
    class SQMat {

  private:
    int dim;
    T** M;

  public:
    SQMat (int N);             // initializes all elements of this n by n matrix to zero
    SQMat (const SQMat& rhs);  // copy constructor
    SQMat (const T& a, int N);  // initialize to diagonal matrix with value a (NOT like in NR !!!)
    SQMat (const SQMat& a, const SQMat& b);  // initialize to tensor product of a and b
    SQMat (const SQMat& a, int row_id, int col_id);  // init by cutting row row_id and col col_id
    void Print ();
    SQMat& operator= (const SQMat& rhs);  // assignment
    SQMat& operator= (const T& a);        // assign 1 to diagonal elements (NOT like in NR !!!)
    inline T* operator[] (const int i);  // subscripting: pointer to row i
    inline const T* operator[] (const int i) const;
    SQMat& operator+= (const T& a);
    SQMat& operator+= (const SQMat& a);
    SQMat& operator-= (const T& a);
    SQMat& operator-= (const SQMat& a);
    SQMat& operator*= (const T& a);
    SQMat& operator*= (const SQMat& a);
    inline int size() const;
    ~SQMat();

  };

  template <class T>
    SQMat<T>::SQMat (int N) : dim(N) , M(new T*[N])
    {
      M[0] = new T[N*N];
      for (int i = 1; i < N; i++) M[i] = M[i-1] + N;
    }

  template <class T>
    SQMat<T>::SQMat (const SQMat& rhs) : dim(rhs.dim) , M(new T*[dim])
    {
      int i,j;
      M[0] = new T[dim*dim];
      for (i = 1; i < dim; i++) M[i] = M[i-1] + dim;
      for (i = 0; i < dim; i++)
	for (j = 0; j < dim; j++) M[i][j] = rhs[i][j];
    }

  template <class T>
    SQMat<T>::SQMat (const T& a, int N) : dim(N) , M(new T*[dim])
    {
      int i, j;
      M[0] = new T[dim*dim];
      for (i = 1; i < dim; i++) M[i] = M[i-1] + dim;
      for (i = 0; i < dim; i++) {
	for (j = 0; j < dim; j++) M[i][j] = T(0);
	M[i][i] = a;
      }
    }

  template <class T>
    SQMat<T>::SQMat (const SQMat& a, const SQMat& b) : dim (a.dim * b.dim) , M(new T*[a.dim * b.dim])
    {
      M[0] = new T[a.dim * b.dim * a.dim * b.dim];

      for (int i = 1; i < a.dim * b.dim; ++i) M[i] = M[i-1] + a.dim * b.dim;

      for (int i1 = 0; i1 < a.dim; ++i1) {

	for (int i2 = 0; i2 < a.dim; ++i2) {

	  for (int j1 = 0; j1 < b.dim; ++j1) {

	    for (int j2 = 0; j2 < b.dim; ++j2) {

    	      M[i1 * (b.dim) + j1][i2 * (b.dim) + j2] = a[i1][i2] * b[j1][j2];
	    }
	  }
	}
      }
    }

  template <class T>
    SQMat<T>::SQMat (const SQMat&a, int row_id, int col_id) : dim (a.dim - 1) , M(new T*[dim])
    {
      if (dim == 0) {
	ABACUSerror("Error:  chopping a row and col from size one matrix.");
	exit(1);
      }

      M[0] = new T[dim * dim];

      for (int i = 1; i < dim; ++i) M[i] = M[i-1] + dim;

      for (int i = 0; i < row_id; ++i)
	for (int j = 0; j < col_id; ++j) M[i][j] = a[i][j];
      for (int i = row_id; i < dim; ++i)
	for (int j = 0; j < col_id; ++j) M[i][j] = a[i+1][j];
      for (int i = 0; i < row_id; ++i)
	for (int j = col_id; j < dim; ++j) M[i][j] = a[i][j+1];
      for (int i = row_id; i < dim; ++i)
	for (int j = col_id; j < dim; ++j) M[i][j] = a[i+1][j+1];

    }

  // operators
  template <class T>
    void SQMat<T>::Print ()
    {
      std::cout << std::endl;
      for (int i = 0; i < dim; ++i) {
	for (int j = 0; j < dim; ++j) std::cout << M[i][j] << " ";
	std::cout << std::endl;
      }
      std::cout << std::endl;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator= (const SQMat<T>& rhs)
    {
      if (this != &rhs) {
	if (dim != rhs.dim) {
	  ABACUSerror("Assignment between matrices of different dimensions.  Bailing out.");
	  exit(1);
	}

	for (int i = 0; i < dim; ++i)
	  for (int j = 0; j < dim; ++j) M[i][j] = rhs[i][j];
      }
      return *this;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator= (const T& a)
    {
      for (int i = 0; i < dim; ++i) {
	for (int j = 0; j < dim; ++j)
	  M[i][j] = T(0);
	M[i][i] = a;
      }
      return *this;
    }

  template <class T>
    inline T* SQMat<T>::operator[] (const int i)
    {
      return M[i];
    }

  template <class T>
    inline const T* SQMat<T>::operator[] (const int i) const
    {
      return M[i];
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator+= (const T& a)
    {

      for (int i = 0; i < dim; ++i) M[i][i] += a;

      return *this;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator+= (const SQMat<T>& a)
    {
      if (dim != a.dim) {
	ABACUSerror("Incompatible matrix sizes in matrix operator +.");
	exit(1);
      }
      for (int i = 0; i < dim; ++i) {
	for (int j = 0; j < dim; ++j) {
	  M[i][j] += a[i][j];
	}
      }
      return *this;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator-= (const T& a)
    {

      for (int i = 0; i < dim; ++i) M[i][i] -= a;

      return *this;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator-= (const SQMat<T>& a)
    {
      if (dim != a.dim) {
	ABACUSerror("Incompatible matrix sizes in matrix operator +.");
	exit(1);
      }
      for (int i = 0; i < dim; ++i) {
	for (int j = 0; j < dim; ++j) {
	  M[i][j] -= a[i][j];
	}
      }
      return *this;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator*= (const T& a)
    {

      for (int i = 0; i < dim; ++i) for (int j = 0; j < dim; ++j) M[i][j] *= a;

      return *this;
    }

  template <class T>
    SQMat<T>& SQMat<T>::operator*= (const SQMat<T>& a)
    {

      if (dim != a.dim) {
	ABACUSerror("Incompatible matrix sizes in matrix operator *.");
	exit(1);
      }

      SQMat<T> leftarg(*this);  // use copy constructor.

      for (int i = 0; i < dim; ++i) {

	for (int j = 0; j < dim; ++j) {

	  M[i][j] = 0.0;

	  for (int k = 0; k < dim; ++k) {

	    M[i][j] += leftarg[i][k] * a[k][j];
	  }
	}
      }

      return *this;
    }

  template <class T>
    inline int SQMat<T>::size() const
    {
      return dim;
    }

  template <class T>
    SQMat<T>::~SQMat()
    {
      if (M != 0) {
	delete[] (M[0]);
	delete[] (M);
      }
    }


  //*****************************

  template <class T>
    class RecMat {

  private:
    int nrows;
    int ncols;
    T** M;

  public:
    RecMat (int Nrows, int Ncols);             // initializes all elements of this n by n matrix to zero
    RecMat (const T& a, int Nrows, int Ncols);
    RecMat (const RecMat& rhs);  // copy constructor
    void Print ();
    RecMat& operator= (const RecMat& rhs);  // assignment
    inline T* operator[] (const int i);  // subscripting: pointer to row i
    inline const T* operator[] (const int i) const;
    inline int nr_rows() const;
    inline int nr_cols() const;
    ~RecMat();

  };

  template <class T>
    RecMat<T>::RecMat (int Nrows, int Ncols) : nrows(Nrows), ncols(Ncols), M(new T*[Nrows])
    {
      M[0] = new T[Nrows*Ncols];
      for (int i = 1; i < Nrows; i++) M[i] = M[i-1] + Ncols;

      for (int i = 0; i < Nrows; i++) for (int j = 0; j < Ncols; j++) M[i][j] = T(0);
    }

  template <class T>
    RecMat<T>::RecMat (const T& a, int Nrows, int Ncols) : nrows(Nrows), ncols(Ncols), M(new T*[Nrows])
    {
      M[0] = new T[Nrows*Ncols];
      for (int i = 1; i < Nrows; i++) M[i] = M[i-1] + Ncols;

      for (int i = 0; i < Nrows; i++) for (int j = 0; j < Ncols; j++) {
	  if (i == j) M[i][i] = a;
	  else M[i][j] = T(0);
	}
    }

  template <class T>
    RecMat<T>::RecMat (const RecMat& rhs) : nrows(rhs.nrows), ncols(rhs.ncols), M(new T*[nrows])
    {
      int i,j;
      M[0] = new T[nrows*ncols];
      for (i = 1; i < nrows; i++) M[i] = M[i-1] + ncols;
      for (i = 0; i < nrows; i++)
	for (j = 0; j < ncols; j++) M[i][j] = rhs[i][j];
    }

  // operators
  template <class T>
    void RecMat<T>::Print ()
    {
      std::cout << std::endl;
      for (int i = 0; i < nrows; ++i) {
	for (int j = 0; j < ncols; ++j) std::cout << M[i][j] << " ";
	std::cout << std::endl;
      }
      std::cout << std::endl;
    }

  template <class T>
    RecMat<T>& RecMat<T>::operator= (const RecMat<T>& rhs)
    {
      if (this != &rhs) {
	if (nrows != rhs.nrows || ncols != rhs.ncols) {
	  if (M != 0) {
	    delete[] (M[0]);
	    delete[] (M);
	  }
	  nrows = rhs.nrows;
	  ncols = rhs.ncols;
	  M = new T*[nrows];
	  M[0] = new T[nrows * ncols];
	}

	for (int i = 0; i < nrows; ++i)
	  for (int j = 0; j < ncols; ++j) M[i][j] = rhs[i][j];
      }
      return *this;
    }

  template <class T>
    inline T* RecMat<T>::operator[] (const int i)
    {
      return M[i];
    }

  template <class T>
    inline const T* RecMat<T>::operator[] (const int i) const
    {
      return M[i];
    }

  template <class T>
    inline int RecMat<T>::nr_rows() const
    {
      return nrows;
    }

  template <class T>
    inline int RecMat<T>::nr_cols() const
    {
      return ncols;
    }

  template <class T>
    inline std::ostream& operator<< (std::ostream& s, const RecMat<T>& matrix)
    {
      for (int i = 0; i < matrix.nr_rows(); ++i) {
	for (int j = 0; j < matrix.nr_cols(); ++j) s << matrix[i][j] << " ";
	s << std::endl;
      }

      return (s);
    }

  template <class T>
    RecMat<T>::~RecMat()
    {
      if (M != 0) {
	delete[] (M[0]);
	delete[] (M);
      }
    }

  // TYPEDEFS:

  typedef ABACUS::SQMat<DP> SQMat_DP;
  typedef ABACUS::SQMat<std::complex<double> > SQMat_CX;


  // FUNCTION DEFINITIONS

  // Functions in src/MATRIX directory

  DP det_LU (SQMat_DP a);
  DP lndet_LU (SQMat_DP a);
  std::complex<DP> lndet_LU_dstry (SQMat_DP& a);
  std::complex<DP> det_LU_CX (SQMat_CX a);
  std::complex<DP> lndet_LU_CX (SQMat_CX a);
  std::complex<DP> lndet_LU_CX_dstry (SQMat_CX& a);
  void eigsrt (Vect_DP& d, SQMat_DP& v);
  void balanc (SQMat_DP& a);
  void elmhes (SQMat_DP& a);
  void gaussj (SQMat_DP& a, SQMat_DP& b);
  void hqr (SQMat_DP& a, Vect_CX& wri);
  void jacobi (SQMat_DP& a, Vect_DP& d, SQMat_DP& v, int& nrot);
  void lubksb (SQMat_DP& a, Vect_INT& indx, Vect_DP& b);
  void lubksb_CX (SQMat_CX& a, Vect_INT& indx, Vect_CX& b);
  void ludcmp (SQMat_DP& a, Vect_INT& indx, DP& d);
  void ludcmp_CX (SQMat_CX& a, Vect_INT& indx, DP& d);
  DP pythag(DP a, DP b);
  void tqli(Vect_DP& d, Vect_DP& e, SQMat_DP& z);
  void tred2 (SQMat_DP& a, Vect_DP& d, Vect_DP& e);

} // namespace ABACUS

#endif