// implementation of bandMatrix class

#include "band.h"
#include "util.h"

bandMatrix::bandMatrix( int n, int bandwidth )
  :ncol(n), bw(bandwidth),
   d(n,0.0), a(n*bandwidth-(bandwidth*(bandwidth+1))/2,0.0),
   b(n*bandwidth-(bandwidth*(bandwidth+1))/2,0.0)
{// buidl an nxn matrix 
 if(bw<0) error("negative bandwidth");
}

double& bandMatrix::operator() (int i, int j)
{
 int k = i-j;
 int ka = abs(k);
 int ks = ncol*(ka-1)-((ka-1)*ka)/2;
 if( !(i>=0 && i<ncol && j>=0 && j<ncol && ka<=bw) )
   error("bandmatrix::operator(): bad index");
 if( k==0 ) return d[i];
 else if( k<0 ) return a[ks+i];
 else  return b[ks+j];
}

double bandMatrix::operator() (int i, int j) const
{
 int k = i-j;
 int ka = abs(k);
 int ks = ncol*(ka-1)-((ka-1)*ka)/2;
 if( !(i>=0 && i<ncol && j>=0 && j<ncol ) )
   error("bandmatrix::operator(): bad index");
 if( k==0 ) return d[i];
 else if( ka<= bw && k<0 ) return a[ks+i];
 else if( ka<= bw && k>0 ) return b[ks+j];
 else  return 0.0;

}


std::ostream& operator <<( std::ostream& os, const bandMatrix& A )
{// print small matrices

 int n = A.cols();
 int i, j;
 for(i=0; i<n; i++){
    for(j=0; j<n; j++)
       os << A(i,j) << "\t";
    os << "\n";
   }
 return os;
}


std::ostream& operator <<( std::ostream& os, const vector<double>& x )
{// print small matrices

 int n = x.size();
 int i;
 for(i=0; i<n; i++){
    os << x[i] << "\t";
 os << "\n";
   }
 return os;
}


vector<double> operator*(const bandMatrix& A, const vector<double>& X )
{// return A X
 int bw = A.hbw(), n =A.cols(), i, k;
 vector<double> y(n,0.0);
 for( i=0; i<n; i++ ){
  y[i] += A(i,i)*X[i];
  for( k=1; k<=bw; k++ ){
     if( i-k >=0 )
        y[i] += A(i,i-k)*X[i-k];
     if( i+k <n )
        y[i] += A(i,i+k)*X[i+k];
  }
}

 return y;
}


vector<double> solve(const bandMatrix & A, const vector<double> & Y )
{// return the solution X of A X = Y

 bandMatrix B(A);
 vector<double> r(Y);

 int i, j, k, m;
 int n=B.cols(), bw=B.hbw(); 

 for( i=0; i<n; i++ ){ // matrix factorization
    B(i,i) = 1.0/B(i,i);
    k = min(i+bw,n-1);
    for( j=i+1; j<=k; j++ )
       B(i,j) *= B(i,i);
    for( j=i+1; j<=k; j++ ){
       for( m=i+1; m<=k; m++ )
          B(j,m) -= B(j,i)*B(i,m);
    }
 }

 for( i=0; i<n; i++ ){ // forward soln
    double rv = r[i];
    double bv = B(i,i);
    r[i] *= B(i,i);
    rv = r[i];
    k = min(i+bw,n-1);
    for( j=i+1; j<=k; j++ ){
      r[j] -= B(j,i)*r[i];
    }
 }

 for( i=n-2; i>=0; i-- ){ // back soln
    k = min(i+bw,n-1);
    for( j=i+1; j<=k; j++ ){
      r[i] -= B(i,j)*r[j];
    }
 }
 return r;
}