package math.matrix;

import java.io.Serializable;

import math.Maths;

/**
 * QR Decomposition.
 * <p>
 * For an m-by-n matrix A with m &ge; n, the QR decomposition is an m-by-n
 * orthogonal matrix Q and an n-by-n upper triangular matrix R so that
 * A = Q*R.
 * <p>
 * The QR decompostion always exists, even if the matrix does not have
 * full rank, so the constructor will never fail.  The primary use of the
 * QR decomposition is in the least squares solution of nonsquare systems
 * of simultaneous linear equations.  This will fail if isFullRank()
 * returns false.
 * @author Daniel Weschke
 */
public class QRDecomposition extends Matrix implements Serializable {

   /**
	 * UID
	 */
	private static final long serialVersionUID = 6913030878857169502L;

/**
    * Array for internal storage of decomposition.
    * @serial internal array storage.
    */
   private double[][] QR;

   /**
    * Row dimensions.
    * @serial row dimension.
    */
   private int m;

   /**
    * Column dimensions.
    * @serial column dimension.
    */
   private int n;

   /**
    * Array for internal storage of diagonal of R.
    * @serial diagonal of R.
    */
   private double[] Rdiag;

   /**
    * QR Decomposition, computed by Householder reflections.
    * Structure to access R and the Householder vectors and compute Q.
    * @param A rectangular matrix
    */
   public QRDecomposition(Matrix A){
      int i, j, k;
      // Initialize.
      QR = A.getCopy();
      m = A.getM();
      n = A.getN();
      Rdiag = new double[n];
      // Main loop.
      for(k=0; k<n; k++){
         // Compute 2-norm of k-th column without under/overflow.
         double nrm = 0;
         for(i=k; i<m; i++)
            nrm = Maths.hypot(nrm,QR[i][k]);
         if(nrm != 0.0){
            // Form k-th Householder vector.
            if(QR[k][k] < 0)
               nrm = -nrm;
            for(i=k; i<m; i++)
               QR[i][k] /= nrm;
            QR[k][k] += 1.0;
            // Apply transformation to remaining columns.
            for(j=k+1; j<n; j++){
               double s = 0.0; 
               for(i=k; i<m; i++)
                  s += QR[i][k]*QR[i][j];
               s = -s/QR[k][k];
               for(i=k; i<m; i++)
                  QR[i][j] += s*QR[i][k];
            }
         }
         Rdiag[k] = -nrm;
      }
   }
   
   /**
    * Is the matrix full rank?
    * @return true if R, and hence A, has full rank.
    */
   public boolean isFullRank(){
      int j;
      for(j=0; j<n; j++)
         if(Rdiag[j] == 0)
            return false;
      return true;
   }

   /**
    * Return the Householder vectors
    * @return lower trapezoidal matrix whose columns define the reflections
    */
   public Matrix getH(){
      int i, j;
      Matrix H = new Matrix(m,n);
      for(i=0; i<m; i++)
         for(j=0; j<n; j++){
            if(i >= j)
               H.set(i,j, QR[i][j]);
            else
               H.set(i,j, 0.0);
         }
      return H;
   }

   /**
    * Return the upper triangular factor
    * @return R
    */
   public Matrix getR(){
      int i, j;
      Matrix R = new Matrix(n,n);
      for(i=0; i<n; i++)
         for(j=0; j<n; j++){
            if(i < j)
               R.set(i,j, QR[i][j]);
            else if(i == j)
               R.set(i,j, Rdiag[i]);
            else
               R.set(i,j, 0.0);
         }
      return R;
   }

   /**
    * Generate and return the (economy-sized) orthogonal factor
    * @return Q
    */
   public Matrix getQ(){
      int i, j, k;
      Matrix Q = new Matrix(m,n);
      for(k=n-1; k>=0; k--) {
         for(i=0; i<m; i++)
            Q.set(i,k, 0.0);
         Q.set(k,k, 1.0);
         for(j=k; j<n; j++)
            if(QR[k][k] != 0){
               double s = 0.0;
               for(i=k; i<m; i++)
                  s += QR[i][k]*Q.get(i, j);
               s = -s/QR[k][k];
               for(i=k; i<m; i++)
                  Q.plus(i,j, s*QR[i][k]);
            }
      }
      return Q;
   }

   /**
    * Least squares solution of A*X = B
    * @param B a Matrix with as many rows as A and any number of columns.
    * @return X that minimizes the two norm of Q*R*X-B.
    * @exception IllegalArgumentException  Matrix row dimensions must agree.
    * @exception RuntimeException  Matrix is rank deficient.
    */
   public Matrix solve(Matrix B){
      if(B.getM() != m)
         throw new IllegalArgumentException("Matrix row dimensions must agree.");
      if(!this.isFullRank())
         throw new RuntimeException("Matrix is rank deficient.");
      int i, j, k;
      // Copy right hand side
      int nx = B.getN();
      double[][] X = B.getCopy();
      // Compute Y = transpose(Q)*B
      for(k=0; k<n; k++)
         for(j=0; j<nx; j++){
            double s = 0.0; 
            for(i=k; i<m; i++)
               s += QR[i][k]*X[i][j];
            s = -s/QR[k][k];
            for(i=k; i<m; i++)
               X[i][j] += s*QR[i][k];
         }
      // Solve R*X = Y;
      for(k=n-1; k>=0; k--){
         for(j=0; j<nx; j++)
            X[k][j] /= Rdiag[k];
         for(i=0; i<k; i++)
            for(j=0; j<nx; j++)
               X[i][j] -= X[k][j]*QR[i][k];
      }
      return (new Matrix(X).get(0,n-1,0,nx-1));
   }
}
