/* Homework 2, problem 3.  Compile with
 *
 *	make hw2_3
 *
 * Run with
 *
 *	hw2_3  N  <  hw2.3.dat  |  gpl -p
 *
 * where N is the number of coefficients in the fitting polynomial
 * (i.e. the polynomial is of degree N-1).
 *
 * The program will print out (to stdout) the original data points,
 * followed by the fitting polynomial evaluated at 1000 points
 * spanning the fitting range.  It also prints (to stderr) the
 * fitting coefficients of 1 (a0), x (a1), x^2 (a2), etc.
 *
 * NOTE: The data were actually generated from a simple polynomial
 * of degree 6, with coefficients a0 = -1, a1 = 3, a2 = 18, a3 = -8,
 * a4 = -48, a5 = 8, and a6 = 32, to which was added noise described
 * by a gaussian random distribution (see Numerical Recipes 7.2) with
 * variable standard deviation randomly chosen between 0.5 and 1.5.
 */

#include <stdio.h>
#include <math.h>
#include "nrutil.h"

#define real double

real poly(real x, int m, real c[])	/* c indices run 1 to m */
{
    int j = m;
    real p = c[j];
    while (j > 1) p = p*x + c[--j];

    return p;
}

#define TOL 1.0e-5

void get_funcs(real x, real *func, int ma)
{
    int i;
    func[1] = 1;
    for (i = 2; i <= ma; i++) func[i] = x*func[i-1];
}

/* This version of svdfit is similar to but (very) slightly modified
 * from the NR version.  See Numerical Recipes, Sec. 15.4. */

void svdfit(real x[], real y[], real sig[], int ndata,
	    real a[], int ma)
{
    int i, j;
    real wmax, tmp, thresh;

    /* Allocate temporary space. */

    real *b = dvector(1, ndata);
    real *afunc = dvector(1, ma);
    real **u = dmatrix(1, ndata, 1, ma);
    real **v = dmatrix(1, ma, 1, ma);
    real *w = dvector(1, ma);

    /* Create the design matrix. */

    for (i = 1; i <= ndata;i ++) {
	get_funcs(x[i], afunc, ma);
	tmp = 1.0/sig[i];
	for (j = 1;j <= ma;j++) u[i][j] = afunc[j]*tmp;
	b[i] = y[i]*tmp;
    }

    svdcmp(u, ndata, ma, w, v);

    wmax = 0.0;
    for (j = 1; j <= ma; j++)
	if (w[j] > wmax) wmax = w[j];
    thresh = TOL*wmax;
    for (j = 1; j <= ma; j++)
	if (w[j] < thresh) w[j] = 0.0;

    svbksb(u, w, v, ndata, ma, b, a);

    /* Clean up. */

    free_dvector(b, 1, ma);
    free_dvector(afunc, 1, ma);
    free_dmatrix(u, 1, ndata, 1, ma);
    free_dmatrix(v, 1, ma, 1, ma);
    free_dvector(w, 1, ma);
}

#define N 500		/* length of the input data file */
#define NFIT 7		/* default number of terms in the fit */
#define NPLOT 1000	/* how many points to plot */

main(int argc, char *argv[])
{
    int k, l;

    /* NR-style declarations of matrices and vectors: */

    real *x  = dvector(1,N);
    real *y  = dvector(1,N);
    real *sig  = dvector(1,N);
    real chi2 = 0;

    int nfit = NFIT;
    if (argc > 1) nfit = atoi(argv[1]);
    real *coef = dvector(1, nfit);

    /* Read the data. */

    for (k = 1; k <= N; k++)
	scanf("%lf %lf %lf", x+k, y+k, sig+k);

    /* Fit the data. */

    svdfit(x, y, sig, N, coef, nfit);

    /* Print out the array of best-fitting coefficients. */

    fprintf(stderr, "N = %d\n", nfit);
    for (k = 1; k <= nfit; k++)
	fprintf(stderr, "a%d =  %f\n", k-1, coef[k]);

    /* Compute chi^2 */

    for (k = 1; k <= N; k++) {
	real yfit = poly(x[k], nfit, coef);
	real error = (y[k] - yfit) / sig[k];
	chi2 += error*error;
    }
    fprintf(stderr, "chi^2 = %f\n", chi2);    

    /* Output the results in a form suitable for plotting with gpl.
       First the data points... */

    for (k = 1; k <= N; k++) printf("%f %f\n", x[k], y[k]);

    /* ... then the fit. */

    for (k = 0; k <= NPLOT; k++) {
	real xx = -1 + k/(0.5*NPLOT);
	printf("%f %f\n", xx, poly(xx, nfit, coef));
    }
}