#include <math.h>
#include <stdio.h>
#include <string.h>
#ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif

#include "symbol.h"
#include "code.h"
#ifdef VMS
#include "math_tab.h"
#else
#include "math.tab.h"
#endif
#include "fudgit.h"
#include "setshow.h"
#include "head.h"

static double *Pvec[4];
static void yfit(double **der, double *par, double *vec, int ndata);
static void fexpo(double *x, double *a, double *y, double **dyda, int ndata);
static void fcosine(double *X, double **du, int ndata);
static void fsine(double *X, double **du, int ndata);
static void fgauss(double *x, double *a, double *y, double **dyda, int ndata);
static void florentz(double *x, double *a, double *y, double **dyda, int ndata);
static void fpoly(double *X, double **du, int ndata);
static void fleg(double *X, double **du, int ndata);

extern double *Ft_X2;
extern double *Ft_Param;
extern void Ft_mathuser(void);


extern void Ft_fit (double *x, double *y, int ndata, double *sig, int mwt, double *a, double *b, double *siga, double *sigb, double *chi2, double *q);
extern void Ft_pearsn (double *x, double *y, int n, double *r, double *prob, double *z);
extern void Ft_medfit (double *x, double *y, int ndata, double *a, double *b, double *abdev);
extern int Ft_svdvar (double **v, int ma, double *w, double **cvm);

int Ft_fits(int argc, char **argv, int ndata)
{
	int i, j;
	Symbol *sym[4];
	char name[TOKENSIZE+10];

	extern int Ft_svdfit(double **coef, double *y, double *sig, int ndata, double *a, int ma, double **v, double *w, double *chisq);
	extern int Ft_lfit(double **coef, double *y, double *sig, int ndata, double *a, int ma, int *lista, int mfit, double **covar, double *chisq);
	extern int Ft_mrqmin(double **coef, double *f, double *y, double *sig, int ndata, double *a, int ma, int *lista, int mfit, double **covar, double **alpha, double *chisq, double *alamda);
	extern Symbol *Ft_lookup(char *s);

/* CHECK FOR STUPID MISTAKES  */
	if (ndata == 0) {
		fputs("No data to fit!\n", stderr);
		return(ERRR);
	}
	else if ((int) *Ft_Param == 0) {
		fputs("You must set parameters first!\n", stderr);
		return(ERRR);
	}
	else if (Ft_Methi == NONE) {
		fputs("You must set method first!\n", stderr);
		return(ERRR);
	}
	else if (Ft_Funci == NONE) {
		fputs("You must set function first!\n", stderr);
		return(ERRR);
	}
	else if ((Ft_Methi == LS_FIT || Ft_Methi == ML_FIT) && Ft_Mlist == 0) {
		fprintf(stderr, "You must 'adjust' before fitting with '%s'.\n",
		Ft_Method[Ft_Methi].name);
		return(ERRR);
	}
/* LOOKUP-INSTALL FIT VARIABLES */
	for (i=0;i<argc-1;i++) {
		if ((sym[i] = Ft_lookup(argv[i+1])) == 0) {
			fprintf(stderr, "%s: No such variable.\n", argv[i+1]);
			return(ERRR);
		}
		else if (sym[i]->type != VEC) {
			fprintf(stderr, "%s: Not a vector.\n", argv[i+1]);
			return(ERRR);
		}
		Pvec[i] = sym[i]->u.vec;
	}
	/* make fit vector */
	sprintf(name, "%sFIT", sym[1]->name);
	if ((sym[3] = Ft_lookup(name)) == 0) {
		sym[3] = Ft_install(name, VEC, Ft_Samples);
	}
	Pvec[3] = sym[3]->u.vec;
	for (i=1;i<= (int) *Ft_Param;i++) {
		sprintf(name, "D%sFITD%d", sym[1]->name, i);
		if ((sym[3] = Ft_lookup(name)) == 0) {
			sym[3] = Ft_install(name, VEC, Ft_Samples);
		}
		Ft_Mparxsamp[i] = sym[3]->u.vec;
	}

/* PREPARE PARTIAL DERIVATIVE MATRIX  */
	/* here we start --- going Ft_Functionwise   */
	if (Ft_Funci == POLY) {
		fpoly(Pvec[0], Ft_Mparxsamp, ndata);
	}
	else if (Ft_Funci == SINE) {
		fsine(Pvec[0], Ft_Mparxsamp, ndata);
	}
	else if (Ft_Funci == COSINE) {
		fcosine(Pvec[0], Ft_Mparxsamp, ndata);
	}
	else if (Ft_Funci == LEG) {
		fleg(Pvec[0], Ft_Mparxsamp, ndata);
	}
	else if (Ft_Funci == GAUSS) {
		if ((int) *Ft_Param%3) {
			fputs("Number of parameters not a multiple of 3.\n", stderr);
			return(ERRR);
		}
		if (Ft_Methi != ML_FIT) {
			fprintf(stderr, "'%s' can only be used with '%s'.\n",
			Ft_Function[Ft_Funci].name, Ft_Method[ML_FIT].name);
			return(ERRR);
		}
		fgauss(Pvec[0], Ft_A, Pvec[3], Ft_Mparxsamp, ndata);
	}
    else if (Ft_Funci == LORENTZ) {
        if ((int) *Ft_Param%3) {
            fputs("Number of parameters not a multiple of 3.\n", stderr);
            return(ERRR);
        }
        if (Ft_Methi != ML_FIT) {
            fprintf(stderr, "'%s' can only be used with '%s'.\n",
            Ft_Function[Ft_Funci].name, Ft_Method[ML_FIT].name);
            return(ERRR);
        }
        florentz(Pvec[0], Ft_A, Pvec[3], Ft_Mparxsamp, ndata);
    }
	else if (Ft_Funci == EXPO) {
		if ((int) *Ft_Param%2) {
			fputs("Number of parameters not a multiple of 2.\n", stderr);
			return(ERRR);
		}
		if (Ft_Methi != ML_FIT) {
			fprintf(stderr, "'%s' can only be used with '%s'.\n",
			Ft_Function[Ft_Funci].name, Ft_Method[ML_FIT].name);
			return(ERRR);
		}
		fexpo(Pvec[0], Ft_A, Pvec[3], Ft_Mparxsamp, ndata);
	}
/* IF STRAIGHT LINE DO IT NOW */
	else if (Ft_Funci == STRL) {
		if (Ft_Methi == SVD_FIT || Ft_Methi == LS_FIT) {
			fprintf(stderr, "Use a polynomial for '%s'!\n",
			Ft_Method[Ft_Methi].name);
			return(ERRR);
		}
		else if (Ft_Methi == LS_REG) {
			int dy = 1;

			if (argc == 3) {
				dy = 0;
				Pvec[2] = 0;
			}
			if ((int) *Ft_Param != 2) {
				fprintf(stderr, "2 parameters are needed for '%s'.\n",
				Ft_Method[LS_REG].name);
				return(ERRR);
			}
			fpoly(Pvec[0], Ft_Mparxsamp, ndata);
			Ft_fit(Pvec[0], Pvec[1], ndata, Pvec[2], dy,
			Ft_A+1, Ft_A+2, Ft_DA+1, Ft_DA+2, Ft_X2, &Ft_Q);
			yfit(Ft_Mparxsamp, Ft_A, Pvec[3], ndata);
			Ft_pearsn(Pvec[0], Pvec[1], ndata, Ft_Cortest, Ft_Cortest+1,
			Ft_Cortest+2);
		}
		else if (Ft_Methi == LA_REG) {
			if ((int) *Ft_Param != 2) {
				fprintf(stderr, "2 parameters are needed for '%s'.\n",
				Ft_Method[LA_REG].name);
				return(ERRR);
			}
			if (argc == 4) {
				fprintf(stderr, "Warning: Errors are not used in '%s'.\n",
				Ft_Method[LA_REG].name);
			}
			fpoly(Pvec[0], Ft_Mparxsamp, ndata);
			Ft_medfit(Pvec[0], Pvec[1], ndata, Ft_A+1, Ft_A+2, Ft_X2);
			yfit(Ft_Mparxsamp, Ft_A, Pvec[3], ndata);
			Ft_pearsn(Pvec[0], Pvec[1], ndata, Ft_Cortest,
			Ft_Cortest+1, Ft_Cortest+2);
		}
		else {
			fprintf(stderr, "Ft_Function '%s' is not supported under '%s'.\n",
			Ft_Function[Ft_Funci].name, Ft_Method[Ft_Methi].name);
			return(ERRR);
		}
		return(0);
	}
	else if (Ft_Funci == USER && Ft_Methi != ML_FIT) {
		Ft_mathuser();
	}

/* WHICH OTHER METHOD */
	/* Now going method wise  */
	if (Ft_Methi == SVD_FIT) {
		if (Ft_svdfit(Ft_Mparxsamp, Pvec[1], Pvec[2],
		ndata, Ft_A, (int) *Ft_Param, Ft_M2parxpar, Ft_Mfparx1, Ft_X2) != ERRR) {
			Ft_svdvar(Ft_M2parxpar, (int) *Ft_Param, Ft_Mfparx1, Ft_M1parxpar);
			yfit(Ft_Mparxsamp, Ft_A, Pvec[3], ndata);
			for (i=1;i<= (int) *Ft_Param;i++) {
				Ft_DA[i] = sqrt(Ft_M1parxpar[i][i]);
			}
		}
	}
	else if (Ft_Methi == LS_FIT) {
		if (Ft_lfit(Ft_Mparxsamp, Pvec[1], Pvec[2], ndata,
		Ft_A, (int) *Ft_Param, Ft_Miparx1, Ft_Mlist, Ft_M1parxpar, Ft_X2) != ERRR) {
			yfit(Ft_Mparxsamp, Ft_A, Pvec[3], ndata);
			for (i=1;i<= (int) *Ft_Param;i++) {
				Ft_DA[i] = sqrt(Ft_M1parxpar[i][i]);
			}
		}
	}
	else if (Ft_Methi == ML_FIT) { /* The state of the art */
		int same = 0;
		double ox2 = 0.0;
		int iter = (Ft_Iter > 0? Ft_Iter : -Ft_Iter);
		/* negative Iter forces the number of iteration
		 * even if it converged. Now what do we do
		 * if (Iter == 0)? Let's get out of here!
		 */

		if (iter == 0) {
			fprintf(stderr, "Warning: Iteration number is null.\n");
			return(0);
		}

		Ft_Q = -1;
		if (Ft_mrqmin(Ft_Mparxsamp, Pvec[3], Pvec[1], Pvec[2], ndata, Ft_A,
		(int) *Ft_Param, Ft_Miparx1, Ft_Mlist, Ft_M1parxpar,
		Ft_M2parxpar, Ft_X2, &Ft_Q) == ERRR) {
			return(ERRR);
		}
		for (i=1; i<= iter;i++) {
			if (Ft_mrqmin(Ft_Mparxsamp, Pvec[3], Pvec[1], Pvec[2], ndata, Ft_A,
			(int) *Ft_Param, Ft_Miparx1, Ft_Mlist, Ft_M1parxpar, Ft_M2parxpar,
			Ft_X2, &Ft_Q) == ERRR) {
				break;
			}
			fprintf(stderr, "~~~~~~~~~~ Iteration: %02d ~~~~~~~~~~~~\t", i);
			fprintf(stderr, "%s: ", "Chi^2");
			fprintf(stderr, Ft_Format, *Ft_X2);
			fprintf(stderr, "\n");
			ox2 = *Ft_X2 -ox2;
			fprintf(stderr, "%45s: ", "DChi^2");
			fprintf(stderr, Ft_Format, ox2);
			fputc('\n', stderr);
			for (j=1;j <= (int) *Ft_Param;j++) {
				char buffer[128];
				sprintf(buffer, "%s[%d]", Ft_Pname, j);
				fprintf(stderr, "%20s: ", buffer);
				fprintf(stderr, Ft_Format, Ft_A[j]);
				fputc('\n', stderr);
			}
			if (Ft_Iter > 0) { /* check for convergence if Iter > 0 */
				if (ox2 == (double)0.0)
					same++;
				else
					same = 0;
				if (same == 2) {
					fprintf(stderr, "Convergence after %d iterations.\n", i);
					break;
				}
				ox2 = *Ft_X2;
			}
		}
		Ft_Q = 0.0;
		if (Ft_mrqmin(Ft_Mparxsamp, Pvec[3], Pvec[1], Pvec[2], ndata, Ft_A,
		(int) *Ft_Param, Ft_Miparx1, Ft_Mlist, Ft_M1parxpar,
		Ft_M2parxpar, Ft_X2, &Ft_Q) == ERRR) {
			return(ERRR);
		}
		for (i=1;i<= (int) *Ft_Param;i++) {
			Ft_DA[i] = sqrt(Ft_M1parxpar[i][i]);
		}
	}
	else {
		fprintf(stderr, "Ft_Function '%s' is not supported under '%s'.\n",
		Ft_Function[Ft_Funci].name, Ft_Method[Ft_Methi].name);
		return(ERRR);
	}
	return(0);
}

void Ft_fml_fit(int ndata)
{
	extern void Ft_mathuser(void);

	switch(Ft_Funci) {
		case USER:
			Ft_mathuser();
			return;
		case GAUSS:
			fgauss(Pvec[0], Ft_A, Pvec[3], Ft_Mparxsamp, ndata);
			return;
        case LORENTZ:
            florentz(Pvec[0], Ft_A, Pvec[3], Ft_Mparxsamp, ndata);
            return;
		case EXPO:
			fexpo(Pvec[0], Ft_A, Pvec[3], Ft_Mparxsamp, ndata);
			return;
	}
}

/* linear */
static void fleg(double *X, double **du, int ndata)
{
	int i, j;
	double twox,f2,f1,d;

	if (Ft_Methi != LS_FIT && Ft_Methi != SVD_FIT) {
		fprintf(stderr,
		"Legendre series not supported under '%s' method.\n",
		Ft_Method[Ft_Methi].name);
		Ft_catcher(ERRR);
	}
	for (i=1; i<= ndata; i++) {
		du[1][i]=1.0;
		du[2][i]=X[i];
		if ((int) *Ft_Param > 2) {
			twox=2.0*X[i];
			f2=X[i];
			d=1.0;
			for (j=3;j<=(int) *Ft_Param;j++) {
				f1=d;
				d += 1.0;
				f2 += twox;
				du[j][i]=(f2*du[j-1][i]-f1*du[j-2][i])/d;
			}
		}
	}
	return;
}

/* linear */
static void fpoly(double *X, double **du, int ndata)
{
	int i, j;

	if (Ft_Methi == ML_FIT) {
		fprintf(stderr,
		"Power series not supported under '%s' method.\n",
		Ft_Method[Ft_Methi].name);
		Ft_catcher(ERRR);
	}
	for (i=1;i<=ndata;i++) {
		du[1][i]=1.0;
		for (j=2;j<=(int) *Ft_Param;j++) du[j][i]=du[j-1][i]*X[i];
	}
	return;
}

/* non-linear  */
static void fgauss(double *x, double *a, double *y, double **dyda, int ndata)
{
	int i, j;
	double fac, ex, arg;

	for (j=1;j<= ndata;j++) {
		y[j]=0.0;
		for (i=1;i<=(int) *Ft_Param-1;i+=3) {
			arg=(x[j]-a[i+1])/a[i+2];
			ex=exp(-arg*arg);
			fac=a[i]*ex*2.0*arg;
			y[j] += a[i]*ex;
			dyda[i][j]=ex;
			dyda[i+1][j]=fac/a[i+2];
			dyda[i+2][j]=fac*arg/a[i+2];
		}
	}
	return;
}

 /* non-linear  */
static void florentz(double *x, double *a, double *y, double **dyda, int ndata)
{
    int i, j;
    double nenner, eg2;
	double xj2;

    for (j=1;j<= ndata;j++) {
        y[j]=0.0;
        for (i=1;i<=(int) *Ft_Param-1;i+=3) {
			xj2=x[j]*x[j];
            eg2=xj2*a[i+2]*a[i+2];
            nenner=(xj2-a[i]*a[i])*(xj2-a[i]*a[i])+eg2;
            y[j] += a[i+1]*eg2/nenner;
            dyda[i][j]= a[i+1]*eg2*4*a[i]*(xj2-a[i]*a[i])/
                          nenner/nenner;
            dyda[i+1][j]=eg2/nenner;
            dyda[i+2][j]=a[i+1]*(2*eg2/a[i+2]*nenner-2*eg2*eg2/a[i+2])/
                          nenner/nenner;
        }
    }
    return;
}


static void fsine(double *X, double **du, int ndata)  /* linear */
{
	int i, j;

	if (Ft_Methi != LS_FIT  && Ft_Methi != SVD_FIT) {
		fprintf(stderr,
		"Sine series not supported under '%s' method.\n",
		Ft_Method[Ft_Methi].name);
		Ft_catcher(ERRR);
	}
	for (i=1;i<= ndata;i++) {
		for (j=1;j<= (int) *Ft_Param;j++) {
			du[j][i] = sin(j*X[i]);
		}
	}
	return;
}

static void fcosine(double *X, double **du, int ndata)  /* linear */
{
	int i, j;

	if (Ft_Methi != LS_FIT  && Ft_Methi != SVD_FIT) {
		fprintf(stderr,
		"Sine series not supported under '%s' method.\n",
		Ft_Method[Ft_Methi].name);
		Ft_catcher(ERRR);
	}
	for (i=1;i<= ndata;i++) {
		du[1][i] = 1.0;
		for (j=2; j<= (int) *Ft_Param;j++) {
			du[j][i] = cos((j-1)*X[i]);
		}
	}
	return;
}

   /* non-linear  */
static void fexpo(double *x, double *a, double *y, double **dyda, int ndata)
{
	int i, j;
	double ex, arg;

	for (j=1;j <= ndata;j++) {
		y[j] = 0.0;
		for (i=1;i<=(int) *Ft_Param-1;i+=2) {
			arg = (x[j] * a[i+1]);
			ex = exp(arg);
			y[j] += a[i] * ex;
			dyda[i][j] = ex;
			dyda[i+1][j] = a[i] * a[i+1] * ex;
		}
	}
	return;
}

/* rebuild the function from the partial(total) linear derivatives  */
/* not valid for non-linear fits  */
static void yfit(double **der, double *par, double *vec, int ndata)
{
	int i, j;

	for (j=1;j<=ndata;j++) {
		vec[j] = 0.0;
	}
	for (i=1;i<=(int) *Ft_Param;i++) {
		for (j=1;j<=ndata;j++) {
			vec[j] += par[i]*der[i][j];
		}
	}
	return;
}
