#include "matrix.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <math.h>


double **multiply(double **A, int am, int an, double **B, int bm, int bn){

  if(an!=bm)
    return 0;
  
  double **rst = new double*[am];
  for(int i=0;i<am;i++){
    rst[i] = new double[bn];
    memset(rst[i],0,bn*sizeof(double));
  }

  for(int i=0;i<am;i++)
    for(int j=0;j<bn;j++)
      for(int k=0;k<an;k++)
        rst[i][j] += A[i][k]*B[k][j];
  
  return rst;
}


double **mminus(double **A, double **B, int m, int n){

  double **rst = new double*[m];
  for(int i=0;i<m;i++){
    rst[i] = new double[n];
    for(int j=0;j<n;j++)
      rst[i][j] = A[i][j]-B[i][j];
  }

  return rst;
}


double **mplus(double **A, double **B, int m, int n){

  double **rst = new double*[m];
  for(int i=0;i<m;i++){
    rst[i] = new double[n];
    for(int j=0;j<n;j++)
      rst[i][j] = A[i][j]+B[i][j];
  }

  return rst;
}


void clear_matrix(double **A, int m, int n){

  for(int i=0;i<m;i++)
    delete[] A[i];
  delete[] A;
}
  

void display_matrix(double **A, int m, int n){

  for(int i=0;i<m;i++){
    for(int j=0;j<n;j++)
      printf("%9.4f  ", A[i][j]);
    printf("\n");
  }
  printf("\n");
}


// Compute inverse of matrix
double **inverse(double **A1, int N, double &determ){
  // Input
  //    A    -    Matrix A (N by N)
  // Outputs
  //   Ainv  -    Inverse of matrix A (N by N)
  //  determ -    Determinant of matrix A	(return value)
  
  int i, j, k;
  double *scale = new double[N];
  double **b = new double*[N];	 // Scale factor and work array
  double **Ainv = new double*[N];
  double **A = new double*[N];
  
  for(i=0;i<N;i++){
    A[i] = new double[N];
    Ainv[i] = new double[N];
    b[i] = new double[N];
    memset(b[i],0,N*sizeof(double));
    memset(Ainv[i],0,N*sizeof(double));
    memcpy(A[i],A1[i],N*sizeof(double));
  }

  int *index = new int[N];
  
  //* Matrix b is initialized to the identity matrix
  for(i=0;i<N;i++)
    b[i][i] = 1.0;

  //* Set scale factor, scale(i) = max( |a(i,j)| ), for each row
  for(i=0;i<N;i++){
    index[i] = i;			  // Initialize row index list
    double scalemax = 0;
    for(j=0;j< N;j++)
      scalemax = (scalemax > fabs(A[i][j])) ? scalemax : fabs(A[i][j]);
    scale[i] = scalemax;
  }

  //* Loop over rows k = 0, ..., (N-2)
  int signDet = 1;
  for(k=0;k<N-1;k++){
    //* Select pivot row from max( |a(j,k)/s(j)| )
    double ratiomax = 0;
    int jPivot = k;
    for(i=k;i<N;i++){
      double ratio = fabs(A[index[i]][k])/scale[index[i]];
      if(ratio > ratiomax){
        jPivot = i;
        ratiomax = ratio;
      }
    }
    //* Perform pivoting using row index list
    int indexJ = index[k];
    if(jPivot != k){
      // Pivot
      indexJ = index[jPivot];
      index[jPivot] = index[k];   // Swap index jPivot and k
      index[k] = indexJ;
      signDet *= -1;			  // Flip sign of determinant
    }
    //* Perform forward elimination
    for(i=k+1;i<N;i++){
      double coeff = A[index[i]][k]/A[indexJ][k];
      for(j=k+1;j<N;j++)
        A[index[i]][j] -= coeff*A[indexJ][j];
      A[index[i]][k] = coeff;
      for(j=0;j<N;j++)
        b[index[i]][j] -= A[index[i]][k]*b[indexJ][j];
    }
  }
  //* Compute determinant as product of diagonal elements
  determ = signDet;	   // Sign of determinant
  for(i=0;i<N;i++)
    determ *= A[index[i]][i];
  
  //* Perform backsubstitution
  for(k=0;k<N;k++){
    Ainv[N-1][k] = b[index[N-1]][k]/A[index[N-1]][N-1];
    for(i=N-2;i>=0;i--){
      double sum = b[index[i]][k];
      for(j=i+1;j<N;j++)
        sum -= A[index[i]][j]*Ainv[j][k];
      Ainv[i][k] = sum/A[index[i]][i];
    }
  }

  // Release allocated memory
  delete[] scale; 
  delete[] index;	
  clear_matrix(A,N,N);
  return Ainv;        
}


double **transpose(double **A, int m, int n){

  double **rst = new double*[n];
  for(int i=0;i<n;i++){
    rst[i] = new double[m];
    for(int j=0;j<m;j++)
      rst[i][j] = A[j][i];
  }

  return rst;
}


int cholesky(double **orig, int N, double **aug, int mcol, double **chol, double **cholaug, int ofs){
     /* 
	Do the augmented cholesky decomposition as described in FA Graybill
	(1976) Theory and Application of the Linear Model. The original matrix
	must be symmetric positive definite. The augmentation matrix, or
	series of column vectors, are multiplied by C^-t, where C is the
	upper triangular cholesky matrix, ie C^t * C = M and M is the original
	matrix. Returns with a value of 0 if M is a non-positive definite 
	matrix. Returns with a value of 1 with succesful completion.

	Arguments:

	orig (input) double N x N array. The matrix to take the Cholesky
	      decomposition of.
	N    (input) integer. Number of rows and columns in orig.
	aug  (input) double N x mcol array. The matrix for the augmented
	      part of the decomposition.
	mcol (input) integer. Number of columns in aug.
	chol (output) double N x N array. Holds the upper triangular matrix
	      C on output. The lower triangular portion remains unchanged.
	      This maybe the same as orig, in which case the upper triangular
	      portion of orig is overwritten.
	cholaug (output) double N x mcol array. Holds the product C^-t * aug.
	         May be the same as aug, in which case aug is over written.
	ofs (input) integer. The index of the first element in the matrices.
	     Normally this is 0, but commonly is 1 (but may be any integer).
     */

  int i, j, k, l;
  int retval = 1;

  for(i=ofs;i<N+ofs;i++){
    chol[i][i] = orig[i][i];
    for(k=ofs;k<i;k++)
      chol[i][i] -= chol[k][i]*chol[k][i];
    if(chol[i][i] <= 0){
     // fprintf(stderr, "\nError: non-positive definite matrix!\n");
     // printf("\nproblem from %d %f\n", i+1,chol[i][i]);
      retval = 0;
      return retval;
    }
    chol[i][i] = sqrt(chol[i][i]);

    /* This portion multiplies the extra matrix by C^-t */
    for(l=ofs;l<mcol+ofs;l++){
      cholaug[i][l] = aug[i][l];
      for(k=ofs;k<i;k++)
        cholaug[i][l] -= cholaug[k][l]*chol[k][i];
      cholaug[i][l] /= chol[i][i];
    }

    for(j=i+1;j<N+ofs;j++){
      chol[i][j] = orig[i][j];
      for(k=ofs;k<i;k++)
        chol[i][j] -= chol[k][i]*chol[k][j];
      chol[i][j] /= chol[i][i];
    }
  }

  return retval;
}


double **self_colmultiply(double **A, int am, int m_start, int m_end, int n_start, int n_end){

  int m = m_end-m_start+1;
  int n = n_end-n_start+1;

  double **rst = new double*[m];
  for(int i=0;i<m;i++){
    rst[i] = new double[n];
    memset(rst[i],0,n*sizeof(double));
  }

  for(int i=0;i<m;i++)
    for(int j=0;j<n;j++)
      for(int k=0;k<am;k++)
        rst[i][j] += A[k][i+m_start-1]*A[k][j+n_start-1];
  
  return rst;
}


double **quadratic_product(double **A, int am, int an, double **B, int bm, int bn, double **C, int cm, int cn){
  // A^T B C

  if(am!=bm || bn!=cm)
    return 0;
  
  double **rst = new double*[an];
  for(int i=0;i<an;i++){
    rst[i] = new double[bn];
    memset(rst[i],0,bn*sizeof(double));
  }

  for(int i=0;i<an;i++)
    for(int j=0;j<bn;j++)
      for(int k=0;k<am;k++)
        rst[i][j] += A[k][i]*B[k][j];

  double **hst = new double*[an];
  for(int i=0;i<an;i++){
    hst[i] = new double[cn];
    memset(hst[i],0,cn*sizeof(double));
  }

  for(int i=0;i<an;i++)
    for(int j=0;j<cn;j++)
      for(int k=0;k<bn;k++)
        hst[i][j] += rst[i][k]*C[k][j];

  for(int i=0;i<an;i++)
    delete[] rst[i];
  delete[] rst;
  
  return hst;
}


double **sub_matrix(double **A, int m, int n, const vector<int> & row_index, const vector<int> & col_index){

  int row = row_index.size();
  int col = col_index.size();

  double **rst = new double*[row];
  for(int i=0;i<row;i++){
    rst[i] = new double[col];
    memset(rst[i],0,col*sizeof(double));
  }

  for(int i=0;i<row;i++)
    for(int j=0;j<col;j++){
      int s = row_index[i];
      int t = col_index[j];
      if(s>m || t>n)
        return 0;
      else
        rst[i][j]=A[s-1][t-1];
    }

  return rst;
}