// @(#)root/matrix:$Name:  $:$Id: TDecompQRH.cxx,v 1.13 2004/06/13 14:53:15 brun Exp $
// Authors: Fons Rademakers, Eddy Offermann  Dec 2003

/*************************************************************************
 * Copyright (C) 1995-2000, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/

///////////////////////////////////////////////////////////////////////////
//                                                                       //
// QR Decomposition class                                                //
//                                                                       //
// Decompose  a general (m x n) matrix A into A = fQ fR H   where        //
//                                                                       //
//  fQ : (m x n) - orthogonal matrix                                     //
//  fR : (n x n) - upper triangular matrix                               //
//  H  : HouseHolder matrix which is stored through                      //
//  fUp: (n) - vector with Householder up's                              //
//  fW : (n) - vector with Householder beta's                            //
//                                                                       //
// Errors arise from formation of reflectors i.e. singularity .          //
// Note it attempts to handle the cases where the nRow <= nCol .         //
//                                                                       //
///////////////////////////////////////////////////////////////////////////

#include "TDecompQRH.h"

ClassImp(TDecompQRH)

//______________________________________________________________________________
TDecompQRH::TDecompQRH(Int_t nrows,Int_t ncols)
{
  if (nrows < ncols) {
    Error("TDecompQRH(Int_t,Int_t","matrix rows should be >= columns");
    return;
  }

  fQ.ResizeTo(nrows,ncols);
  fR.ResizeTo(ncols,ncols);
  if (nrows <= ncols) {
    fW.ResizeTo(nrows);
    fUp.ResizeTo(nrows);
  } else {
    fW.ResizeTo(ncols);
    fUp.ResizeTo(ncols);
  }
}

//______________________________________________________________________________
TDecompQRH::TDecompQRH(Int_t row_lwb,Int_t row_upb,Int_t col_lwb,Int_t col_upb)
{
  const Int_t nrows = row_upb-row_lwb+1;
  const Int_t ncols = col_upb-col_lwb+1;

  if (nrows < ncols) {
    Error("TDecompQRH(Int_t,Int_t,Int_t,Int_t","matrix rows should be >= columns");
    return;
  }

  fRowLwb = row_lwb;
  fColLwb = col_lwb;

  fQ.ResizeTo(nrows,ncols);
  fR.ResizeTo(ncols,ncols);
  if (nrows <= ncols) {
    fW.ResizeTo(nrows);
    fUp.ResizeTo(nrows);
  } else {
    fW.ResizeTo(ncols);
    fUp.ResizeTo(ncols);
  }
}

//______________________________________________________________________________
TDecompQRH::TDecompQRH(const TMatrixD &a,Double_t tol)
{
  Assert(a.IsValid());
  if (a.GetNrows() < a.GetNcols()) {
    Error("TDecompQRH(const TMatrixD &","matrix rows should be >= columns");
    return;
  }

  SetBit(kMatrixSet);
  fCondition = a.Norm1();
  fTol = a.GetTol();
  if (tol > 0.0)
    fTol = tol;

  fRowLwb = a.GetRowLwb();
  fColLwb = a.GetColLwb();
  const Int_t nRow = a.GetNrows();
  const Int_t nCol = a.GetNcols();

  fQ.ResizeTo(nRow,nCol);
  memcpy(fQ.GetMatrixArray(),a.GetMatrixArray(),nRow*nCol*sizeof(Double_t));
  fR.ResizeTo(nCol,nCol);
  if (nRow <= nCol) {
    fW.ResizeTo(nRow);
    fUp.ResizeTo(nRow);
  } else {
    fW.ResizeTo(nCol);
    fUp.ResizeTo(nCol);
  }
}

//______________________________________________________________________________
TDecompQRH::TDecompQRH(const TDecompQRH &another) : TDecompBase(another)
{
  *this = another;
}

//______________________________________________________________________________
Bool_t TDecompQRH::Decompose()
{
// QR decomposition of matrix a by Householder transformations,
//  see Golub & Loan first edition p41 & Sec 6.2.
// First fR is returned in upper triang of fQ and diagR. fQ returned in
// 'u-form' in lower triang of fQ and fW, the latter containing the
//  "Householder betas".

  if ( !TestBit(kMatrixSet) )
    return kFALSE;

  const Int_t nRow   = this->GetNrows();
  const Int_t nCol   = this->GetNcols();
  const Int_t rowLwb = this->GetRowLwb();
  const Int_t colLwb = this->GetColLwb();

  TVectorD diagR;
  Double_t work[kWorkMax];
  if (nCol > kWorkMax) diagR.ResizeTo(nCol);
  else                 diagR.Use(nCol,work);

  if (QRH(fQ,diagR,fUp,fW,fTol)) {
    for (Int_t i = 0; i < nRow; i++) {
      const Int_t ic = (i < nCol) ? i : nCol;
      for (Int_t j = ic ; j < nCol; j++)
        fR(i,j) = fQ(i,j);
    }
    TMatrixDDiag diag(fR); diag = diagR;

    fQ.Shift(rowLwb,0);
    fR.Shift(0,colLwb);

    SetBit(kDecomposed);
  }

  return kTRUE;
}

//______________________________________________________________________________
Bool_t TDecompQRH::QRH(TMatrixD &q,TVectorD &diagR,TVectorD &up,TVectorD &w,Double_t tol)
{
  const Int_t nRow = q.GetNrows();
  const Int_t nCol = q.GetNcols();

  const Int_t n = (nRow <= nCol) ? nRow-1 : nCol;

  for (Int_t k = 0 ; k < n ; k++) {
    const TVectorD qc_k = TMatrixDColumn_const(q,k);
    if (!DefHouseHolder(qc_k,k,k+1,up(k),w(k),tol))
      return kFALSE;
    diagR(k) = qc_k(k)-up(k);
    if (k < nCol-1) {
      // Apply HouseHolder to sub-matrix
      for (Int_t j = k+1; j < nCol; j++) {
        TMatrixDColumn qc_j = TMatrixDColumn(q,j);
        ApplyHouseHolder(qc_k,up(k),w(k),k,k+1,qc_j);
      }
    }
  }

  if (nRow <= nCol) {
    diagR(nRow-1) = q(nRow-1,nRow-1);
    up(nRow-1) = 0;
    w(nRow-1)  = 0;
  }

  return kTRUE;
}

//______________________________________________________________________________
void TDecompQRH::SetMatrix(const TMatrixD &a)
{
  Assert(a.IsValid());
  if (a.GetNrows() < a.GetNcols()) {
    Error("TDecompQRH(const TMatrixD &","matrix rows should be >= columns");
    return;
  }

  ResetStatus();
  SetBit(kMatrixSet);
  fCondition = a.Norm1();

  fRowLwb = a.GetRowLwb();
  fColLwb = a.GetColLwb();
  const Int_t nRow = a.GetNrows();
  const Int_t nCol = a.GetNcols();

  fQ.ResizeTo(nRow,nCol);
  memcpy(fQ.GetMatrixArray(),a.GetMatrixArray(),nRow*nCol*sizeof(Double_t));
  fR.ResizeTo(nCol,nCol);
  if (nRow <= nCol) {
    fW.ResizeTo(nRow);
    fUp.ResizeTo(nRow);
  } else {
    fW.ResizeTo(nCol);
    fUp.ResizeTo(nCol);
  }
}

//______________________________________________________________________________
Bool_t TDecompQRH::Solve(TVectorD &b)
{
// Solve Ax=b assuming the QR form of A is stored in fR,fQ and fW, but assume b
// has *not* been transformed.  Solution returned in b.

  Assert(b.IsValid());
  if (TestBit(kSingular))
    return kFALSE;
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose())
      return kFALSE;
  }

  if (fQ.GetNrows() != b.GetNrows() || fQ.GetRowLwb() != b.GetLwb()) { 
    Error("Solve(TVectorD &","vector and matrix incompatible");
    b.Invalidate();
    return kFALSE;
  } 

  const Int_t nQRow = fQ.GetNrows();
  const Int_t nQCol = fQ.GetNcols();

  // Calculate  Q^T.b
  const Int_t nQ = (nQRow <= nQCol) ? nQRow-1 : nQCol;
  for (Int_t k = 0; k < nQ; k++) {
    const TVectorD qc_k = TMatrixDColumn_const(fQ,k);
    ApplyHouseHolder(qc_k,fUp(k),fW(k),k,k+1,b);
  }

  const Int_t nRCol = fR.GetNcols();

  const Double_t *pR = fR.GetMatrixArray();
        Double_t *pb = b.GetMatrixArray();

  // Backward substitution
  for (Int_t i = nRCol-1; i >= 0; i--) {
    const Int_t off_i = i*nRCol;
    Double_t r = pb[i];
    for (Int_t j = i+1; j < nRCol; j++)
      r -= pR[off_i+j]*pb[j];
    if (TMath::Abs(pR[off_i+i]) < fTol)
    {
      Error("Solve(TVectorD &)","R[%d,%d]=%.4e < %.4e",i,i,pR[off_i+i],fTol);
      return kFALSE;
    }
    pb[i] = r/pR[off_i+i];
  }

  return kTRUE;
}

//______________________________________________________________________________
TVectorD TDecompQRH::Solve(const TVectorD &b,Bool_t &ok)
{    
// Solve A x=b assuming the QR form of A is stored in fR,fQ and fW, but assume b
// has *not* been transformed.  
    
  TVectorD x = b; 
  ok = Solve(x);
  if (GetNrows() != GetNcols())
    x.ResizeTo(GetNcols());
      
  return x;
}

//______________________________________________________________________________
Bool_t TDecompQRH::Solve(TMatrixDColumn &cb)
{ 
  const TMatrixDBase *b = cb.GetMatrix();
  Assert(b->IsValid());
  if (TestBit(kSingular))
    return kFALSE;
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose())
      return kFALSE;
  }

  if (fQ.GetNrows() != b->GetNrows() || fQ.GetRowLwb() != b->GetRowLwb())
  {
    Error("Solve(TMatrixDColumn &","vector and matrix incompatible");
    return kFALSE; 
  }
  
  const Int_t nQRow = fQ.GetNrows();
  const Int_t nQCol = fQ.GetNcols();
    
  // Calculate  Q^T.b
  const Int_t nQ = (nQRow <= nQCol) ? nQRow-1 : nQCol;
  for (Int_t k = 0; k < nQ; k++) {
    const TVectorD qc_k = TMatrixDColumn_const(fQ,k);
    ApplyHouseHolder(qc_k,fUp(k),fW(k),k,k+1,cb);
  }   
    
  const Int_t nRCol = fR.GetNcols();
  
  const Double_t *pR  = fR.GetMatrixArray();
        Double_t *pcb = cb.GetPtr();
  const Int_t     inc = cb.GetInc();

  // Backward substitution
  for (Int_t i = nRCol-1; i >= 0; i--) {
    const Int_t off_i  = i*nRCol;
    const Int_t off_i2 = i*inc;
    Double_t r = pcb[off_i2];
    for (Int_t j = i+1; j < nRCol; j++)
      r -= pR[off_i+j]*pcb[j*inc];
    if (TMath::Abs(pR[off_i+i]) < fTol)
    {
      Error("Solve(TMatrixDColumn &)","R[%d,%d]=%.4e < %.4e",i,i,pR[off_i+i],fTol);
      return kFALSE;
    }
    pcb[off_i2] = r/pR[off_i+i];
  }

  return kTRUE;
}

//______________________________________________________________________________
Bool_t TDecompQRH::TransSolve(TVectorD &b)
{
// Solve A^T x=b assuming the QR form of A is stored in fR,fQ and fW, but assume b
// has *not* been transformed.  Solution returned in b.

  Assert(b.IsValid());
  if (TestBit(kSingular))
    return kFALSE;
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose())
      return kFALSE;
  }

  if (fQ.GetNrows() != fQ.GetNcols() || fQ.GetRowLwb() != fQ.GetColLwb()) {
    Error("TransSolve(TVectorD &","matrix should be square");
    b.Invalidate();
    return kFALSE;
  }

  if (fR.GetNrows() != b.GetNrows() || fR.GetRowLwb() != b.GetLwb()) {   
    Error("TransSolve(TVectorD &","vector and matrix incompatible");
    b.Invalidate();
    return kFALSE;
  } 

  const Double_t *pR = fR.GetMatrixArray();
        Double_t *pb = b.GetMatrixArray();

  const Int_t nRCol = fR.GetNcols();

  // Backward substitution
  for (Int_t i = 0; i < nRCol; i++) {
    const Int_t off_i = i*nRCol;
    Double_t r = pb[i];
    for (Int_t j = 0; j < i; j++) {
      const Int_t off_j = j*nRCol;
      r -= pR[off_j+i]*pb[j];
    }
    if (TMath::Abs(pR[off_i+i]) < fTol)
    {
      Error("TransSolve(TVectorD &)","R[%d,%d]=%.4e < %.4e",i,i,pR[off_i+i],fTol);
      return kFALSE;
    }
    pb[i] = r/pR[off_i+i];
  }

  const Int_t nQRow = fQ.GetNrows();

  // Calculate  Q.b; it was checked nQRow == nQCol
  for (Int_t k = nQRow-1; k >= 0; k--) {
    const TVectorD qc_k = TMatrixDColumn_const(fQ,k);
    ApplyHouseHolder(qc_k,fUp(k),fW(k),k,k+1,b);
  }

  return kTRUE;
}

//______________________________________________________________________________
TVectorD TDecompQRH::TransSolve(const TVectorD &b,Bool_t &ok)
{   
// Solve A^T x=b assuming the QR form of A is stored in fR,fQ and fW, but assume b
// has *not* been transformed.

  TVectorD x = b;
  ok = TransSolve(x);

  return x;
}

//______________________________________________________________________________
Bool_t TDecompQRH::TransSolve(TMatrixDColumn &cb)
{
  const TMatrixDBase *b = cb.GetMatrix();
  Assert(b->IsValid());
  if (TestBit(kSingular))
    return kFALSE;
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose())
      return kFALSE;
  }

  if (fQ.GetNrows() != fQ.GetNcols() || fQ.GetRowLwb() != fQ.GetColLwb()) {
    Error("TransSolve(TMatrixDColumn &","matrix should be square");
    return kFALSE;
  }
  
  if (fR.GetNrows() != b->GetNrows() || fR.GetRowLwb() != b->GetRowLwb()) {
    Error("TransSolve(TMatrixDColumn &","vector and matrix incompatible");
    return kFALSE;
  }
  
  const Double_t *pR  = fR.GetMatrixArray();
        Double_t *pcb = cb.GetPtr();
  const Int_t     inc = cb.GetInc();
  
  const Int_t nRCol = fR.GetNcols();
  
  // Backward substitution
  for (Int_t i = 0; i < nRCol; i++) {
    const Int_t off_i  = i*nRCol;
    const Int_t off_i2 = i*inc;
    Double_t r = pcb[off_i2];
    for (Int_t j = 0; j < i; j++) {
      const Int_t off_j = j*nRCol;
      r -= pR[off_j+i]*pcb[j*inc];
    }
    if (TMath::Abs(pR[off_i+i]) < fTol)
    {
      Error("TransSolve(TMatrixDColumn &)","R[%d,%d]=%.4e < %.4e",i,i,pR[off_i+i],fTol);
      return kFALSE;
    }
    pcb[off_i2] = r/pR[off_i+i];
  }

  const Int_t nQRow = fQ.GetNrows();

  // Calculate  Q.b; it was checked nQRow == nQCol
  for (Int_t k = nQRow-1; k >= 0; k--) {
    const TVectorD qc_k = TMatrixDColumn_const(fQ,k);
    ApplyHouseHolder(qc_k,fUp(k),fW(k),k,k+1,cb);
  }

  return kTRUE;
}

//______________________________________________________________________________
void TDecompQRH::Det(Double_t &d1,Double_t &d2)
{
  // This routine calculates the absolute (!) value of the determinant

  if ( !TestBit(kDetermined) ) {
    if ( !TestBit(kDecomposed) )
      Decompose();
    if (TestBit(kSingular)) {
      fDet1 = 0.0;
      fDet2 = 0.0;
    } else
      TDecompBase::Det(d1,d2);
    SetBit(kDetermined);
  }
  d1 = fDet1;
  d2 = fDet2;
}

//______________________________________________________________________________
void TDecompQRH::Print(Option_t *opt) const
{
  TDecompBase::Print(opt);
  fQ.Print("fQ");
  fR.Print("fR");
  fUp.Print("fUp");
  fW.Print("fW");
}

//______________________________________________________________________________
TDecompQRH &TDecompQRH::operator=(const TDecompQRH &source)
{ 
  if (this != &source) {
    TDecompBase::operator=(source);
    fQ.ResizeTo(source.fQ);
    fR.ResizeTo(source.fR);
    fUp.ResizeTo(source.fUp);
    fW.ResizeTo(source.fW);
    fQ  = source.fQ;
    fR  = source.fR;
    fUp = source.fUp;
    fW  = source.fW;
  }
  return *this;
}
