#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <f2c.h>
#include "gpklib.h"

extern "C"
/* Subroutine */ int dgels_(char *trans, integer *m, integer *n, integer *
	nrhs, doublereal *a, integer *lda, doublereal *b, integer *ldb, 
	doublereal *work, integer *lwork, integer *info);


static double xx(double x) { return x*x; }

/* Solve A[m][n] * X[n][nrhs] = B[m][mrhs], and return X.
   Returns either the LMS solution (if overdetermined),
	or the minimum norm solution (if underdetermined).
   If errs is non-nil,
   it writes the least-squares errors into errs[0 ... nrhs-1]. */

double ** eqn_solve(int m, int n, double **a,
					int nrhs, double **b,
					double *errs)
{
	assert(m>=1 && n>=1 && nrhs>=1);
	double *aa = double1alloc(m*n);
	char trans = 'N';
	int i;
	long lda = m, ldb = max(m, n), lwork = (m+n+nrhs)*8, info;
	long ml = m, nl = n, nrhsl = nrhs;
	double *bb = double1alloc(ldb*nrhs);
	double *work = double1alloc(lwork);
	assert(aa && bb && work);
	for(i=0; i<m; i++) {
		int j;
		for(j=0; j<n; j++)
			aa[i+j*lda] = a[i][j];
	}
	for(i=0; i<m; i++) {
		int j;
		for(j=0; j<nrhs; j++)
			bb[i+j*ldb] = b[i][j];
	}

	dgels_(&trans, &ml, &nl, &nrhsl, aa, &lda, bb, &ldb, work, &lwork, &info);


	double1free(aa);
	double1free(work);

	if(info != 0)
		return 0;

	double **rv = double2alloc(n, nrhs);
	assert(rv);
	for(i=0; i<n; i++) {
		int j;
		for(j=0; j<nrhs; j++)
			rv[i][j] = bb[i+j*ldb];
	}
	if(errs) {
		int j;
		for(j=0; j<nrhs; j++) {
			int k;
			errs[j] = 0;
			for(k=n; k<m; k++) {
				errs[j] += xx(bb[j*ldb+k]);
			}
		}
	}
	double1free(bb);
	assert(double2test(rv) == 0);

	return  rv;
}

#if 0

int main(int argc, char **argv)
{
#if 0
	double **a = double2alloc(3, 1);
	a[0][0] = 2;
	a[1][0] = 2;
	a[2][0] = 2;
	double **b = double2alloc(3, 1);
	b[0][0] = 6;
	b[1][0] = 5;
	b[2][0] = 4;
	double *errs = double1alloc(1);
	double **x = eqn_solve(3, 1, a, 1, b, errs);
	printf("x=%g (should be 2.5), and errs=%g (should be 0.5)\n", x[0][0], errs[0]);
	double1free(errs);
	return 1;
#endif
#if 0
	double **a = double2alloc(1, 3);
	a[0][0] = 1;
	a[0][1] = 1;
	a[0][2] = 1;
	double **b = double2alloc(1, 1);
	b[0][0] = 6;
	double *errs = double1alloc(3);
	double **x = eqn_solve(1, 3, a, 1, b, errs);
	printf("x=%g %g %g , and errs=%g\n", x[0][0], x[1][0], x[2][0], errs[0]);
	double1free(errs);
	return 1;
#endif
#if 1
	double **a = double2alloc(2, 2);
	a[0][0] = 1;
	a[0][1] = 0;
	a[1][0] = 1;
	a[1][1] = -1;
	double **b = double2alloc(2, 1);
	b[0][0] = 6;
	b[1][0] = 2;
	double *errs = double1alloc(1);
	double **x = eqn_solve(2, 2, a, 1, b, errs);
	printf("x=%g %g, and errs=%g\n", x[0][0], x[1][0], errs[0]);
	double1free(errs);
	return 1;
#endif
#if 0
	double **a = double2alloc(1, 3);
	a[0][0] = 1;
	a[0][1] = 1;
	a[0][2] = 1;
	double **b = double2alloc(1, 2);
	b[0][0] = 6;
	b[0][1] = 3;
	double *errs = double1alloc(2);
	double **x = eqn_solve(1, 3, a, 2, b, errs);
	printf("x=%g %g %g then %g %g %g, and errs=%g %g\n", x[0][0], x[1][0], x[2][0],
				x[0][1], x[1][1], x[2][1], errs[0], errs[1]);
	double1free(errs);
	return 1;
#endif
#if 0
	double **a = double2alloc(3, 1);
	a[0][0] = 2;
	a[1][0] = 2;
	a[2][0] = 2;
	double **b = double2alloc(3, 2);
	b[0][0] = 6;
	b[1][0] = 5;
	b[2][0] = 4;
	b[0][1] = 3;
	b[1][1] = 2;
	b[2][1] = 1;
	double *errs = double1alloc(2);
	double **x = eqn_solve(3, 1, a, 2, b, errs);
	printf("x=%g %g, and errs=%g %g\n", x[0][0], x[0][1], errs[0], errs[1]);
	double1free(errs);
	return 1;
#endif
}

#endif

extern "C"
int dgbsv_(integer *n, integer *kl, integer *ku, integer *
	nrhs, doublereal *ab, integer *ldab, integer *ipiv, doublereal *b, integer *ldb, 
	integer *info);


static int mx(int a, int b) { return a>b ? a : b;}

#define AA(i,j)		aa[(i)+(j)*ldab]

/* Solve A[n][n] * X[n][nrhs] = B[n][mrhs], and return X.
 * A is a band-diagonal matrix with a diagonal and w extras on each side.
 * It is, however, stored as a full-size matrix, such as you can get
 * from double2alloc().   It depends on the a[i][j] addressing convention.
   */

double ** eqn_solve_band(int n, double **a, int bw,
				int nrhs, double **b)
{
	fprintf(stderr, "eqsb\n");
	assert(n>=1 && nrhs>=1 && bw>=0 && bw<n);
	long kl = bw, ku = bw;
	long ldab = 2*kl+ku+1;
	long ldb = n, info;
	double *aa = double1alloc(n*ldab);
	long nl = n, nrhsl = nrhs;
	int i;
	double *bb = double1alloc(ldb*nrhs);
	long *ipiv = new long[n];
	assert(aa && bb && ipiv);
	fprintf(stderr, "loop1\n");

	for(i=0; i<n; i++) {
		int j;
		for(j=mx(0, i-kl); j<=i+ku && j<n; j++) {
			AA(kl+ku+i-j, j) = a[i][j];
		}
	}
	for(i=0; i<n; i++) {
		int j;
		for(j=0; j<nrhs; j++) {
			bb[i+j*ldb] = b[i][j];
		}
	}

	dgbsv_(&nl, &kl, &ku, &nrhsl, aa, &ldab, ipiv, bb, &ldb, &info);


	double1free(aa);
	delete [] ipiv;

	if(info != 0)
		return 0;

	double **rv = double2alloc(n, nrhs);
	assert(rv);
	for(i=0; i<n; i++) {
		int j;
		for(j=0; j<nrhs; j++)
			rv[i][j] = bb[i+j*ldb];
	}
	double1free(bb);

	return  rv;
}


#if 0

int main(int argc, char **argv)
{
#if 0
	double **a = double2alloc(2, 2);
	a[0][0] = 1;
	a[0][1] = 0;
	a[1][0] = 1;
	a[1][1] = -1;
	double **b = double2alloc(2, 1);
	b[0][0] = 6;
	b[1][0] = 2;
	double **x = eqn_solve_band(2, a, 1, 1, b);
	printf("x=%g %g\n", x[0][0], x[1][0]);
	return 1;
#endif
#if 0
	double **a = double2alloc(2, 2);
	a[0][0] = 1;
	a[1][1] = -1;
	double **b = double2alloc(2, 1);
	b[0][0] = 6;
	b[1][0] = 2;
	fprintf(stderr, "ready to eqs\n");
	double **x = eqn_solve_band(2, a, 0, 1, b);
	printf("x=%g %g\n", x[0][0], x[1][0]);
	return 1;
#endif
#if 0
	double **a = double2alloc(3, 3);
	a[0][0] = 1;
	a[1][1] = -1;
	a[2][2] = 2;
	double **b = double2alloc(3, 1);
	b[0][0] = 6;
	b[1][0] = 2;
	b[2][0] = 2;
	fprintf(stderr, "ready to eqs\n");
	double **x = eqn_solve_band(3, a, 0, 1, b);
	printf("x=%g %g %g\n", x[0][0], x[1][0], x[2][0]);
	return 1;
#endif
#if 0
	double **a = double2alloc(3, 3);
	a[0][0] = 1;
	a[1][1] = -1;
	a[2][2] = 2;
	a[1][0] = a[0][1] = a[2][1] = a[1][2] = 0;
	double **b = double2alloc(3, 1);
	b[0][0] = 6;
	b[1][0] = 2;
	b[2][0] = 2;
	fprintf(stderr, "ready to eqs\n");
	double **x = eqn_solve_band(3, a, 1, 1, b);
	printf("x=%g %g %g\n", x[0][0], x[1][0], x[2][0]);
	return 1;
#endif
#if 1
	double **a = double2alloc(4, 4);
	a[0][0] = 1;
	a[1][1] = -1;
	a[1][0] = 1;
	a[2][2] = 2;
	a[3][3] = 1;
	a[0][1] = a[2][1] = a[1][2] = 0;
	a[3][2] = a[2][3] = 0;
	a[2][0] = a[0][2] = a[3][1] = a[1][3] = 0;
	double **b = double2alloc(3, 1);
	b[0][0] = 6;
	b[1][0] = 2;
	b[2][0] = 2;
	b[3][0] = 0;
	fprintf(stderr, "ready to eqs\n");
	double **x = eqn_solve_band(4, a, 2, 1, b);
	printf("x=%g %g %g %g\n", x[0][0], x[1][0], x[2][0], x[3][0]);
	return 1;
#endif
}

#endif
