#include <stdio.h>

#include "cmpack_common.h"
#include "linsolver.h"
#include "qrdecomp.h"

/*************************    LINEAR SOLVER BASE CLASS       *************************************/

typedef void tSolve(CmpackLinSolver *solver, const double *b, double *x);
typedef void tDestroy(CmpackLinSolver *solver);

struct _CmpackLinSolver
{
	tSolve *solve;
	tDestroy *destroy;
};

/* Find vector x for given vector b that meets Ax = b */
void cmpack_linsolver_solve(CmpackLinSolver *solver, const double *b, double *x)
{
	if (solver)
		solver->solve(solver, b, x);
}

/* Free allocated memory */
void cmpack_linsolver_destroy(CmpackLinSolver *solver)
{
	if (solver)
		solver->destroy(solver);
}

/*************************  SOLVER BASED ON QR DECOMPOSITION *************************************/

struct _CmpackLinSolverQR
{
	struct _CmpackLinSolver p;
	int	n;							/**< Dimension of the matrix - n x n */
	double *q, *r, *y;
};

static void qr_solve(CmpackLinSolver *solver, const double *b, double *x);
static void qr_destroy(CmpackLinSolver *solver);

/* Create a linear solver based on QR decomposition */
CmpackLinSolver *cmpack_linsolver_qr_create(int n, const double *a)
{
	struct _CmpackLinSolverQR *retval = cmpack_malloc(sizeof(struct _CmpackLinSolverQR));

	QRDecomposition decomp;

	/* Memory allocation */
	retval->p.solve = qr_solve;
	retval->p.destroy = qr_destroy;
	retval->n = n;
	retval->y = (double*)cmpack_malloc(n*sizeof(double));
	retval->r = (double*)cmpack_malloc(n*n*sizeof(double));
	retval->q = (double*)cmpack_malloc(n*n*sizeof(double));
	cmpack_qrd_alloc(&decomp, n, n);
	cmpack_qrd_set(&decomp, a, retval->q, retval->r);
	cmpack_qrd_free(&decomp);

	return (CmpackLinSolver*)retval;
}

static void qr_destroy(CmpackLinSolver *solver)
{
	struct _CmpackLinSolverQR *data = (struct _CmpackLinSolverQR *)solver;
	cmpack_free(data->y);
	cmpack_free(data->q);
	cmpack_free(data->r);
	cmpack_free(data);
}

/* Solve equation Ax = b */
static void qr_solve(CmpackLinSolver *solver, const double *b, double *x)
{
	struct _CmpackLinSolverQR *data = (struct _CmpackLinSolverQR *)solver;

	int i, j, k, n = data->n;
	double *y = data->y, *q = data->q, *r = data->r;

	/* Iy = Q'B */
	for (j=0; j<n; j++) {
		double aux = 0.0;
		for (k=0; k<n; k++)
			aux += q[j+k*n] * b[k];
		y[j] = aux;
	}

	/* Rx = y */
	for (j=n-1; j>=0; j--) {
		double aux = y[j]; 
		for (i=j+1; i<n; i++) 
			aux -= x[i] * r[i+j*n]; 
		x[j] = aux / r[j+j*n];
	}
}
