/* routine does nonlinear fit on data in file specified on */
/* command line. Writes fit to second file */
/* on command line. Best fit parameters plus chisquared written to */
/* standard output. Can fit to only a subset of data points */ 

#include <stdio.h>
#include <math.h>
#include <stdlib.h>

static double sqrarg;
#define SQR(a) ((sqrarg=(a)) == 0.0 ? 0.0 : sqrarg*sqrarg)

static double maxarg1,maxarg2;
#define FMAX(a,b) (maxarg1=(a),maxarg2=(b),(maxarg1) > (maxarg2) ?\
(maxarg1) : (maxarg2))

static int iminarg1,iminarg2;
#define IMIN(a,b) (iminarg1=(a),iminarg2=(b),(iminarg1) < (iminarg2) ?\
(iminarg1) : (iminarg2))

#define SIGN(a,b) ((b) > 0.0 ? fabs(a) : -fabs(a)) 

void nrerror(char error_text[]);
double *vector(long nl, long nh);
int *ivector(long nl, long nh);
double **matrix(long nrl, long nrh, long ncl, long nch);
void free_vector(double *v, long nl, long nh);
void free_ivector(int *v, long nl, long nh);
void free_matrix(double **m, long nrl, long nrh, long ncl, long nch);

void lubksb(double **a, int n, int *indx, double b[]);
void ludcmp(double **a, int n, int *indx, double *d);
void hessian(double *w,  int n, double *corr,
double *ix,
double *p, double lambda, double **a, double *b);
double chisq(double *w, int n, double *corr,
double *ix,
double *p);
double func(double x, double *p);
double dfunc(int i, double x, double *p);

int PARAM=2;
int L;

double func(double x, double *p)
{
double dum,dum2;
dum2=(x-L/2)*p[2];
dum=p[1]*(exp(-dum2)+exp(dum2));
return(dum);
}

double dfunc(int i, double x, double *p)
{
double dum,dum2;
dum2=(x-L/2)*p[2];
if(i==1) dum=exp(-dum2)+exp(dum2);
if(i==2) dum=p[1]*(x-L/2)*(exp(dum2)-exp(-dum2));
return(dum);
}

double chisq(double *e, int TIME, double *corr,
double *ix,
double *p)
{
int i;
double add;

add=0.0;
for(i=1;i<=TIME;i++)
add+=(corr[i]-func(ix[i],p))*(corr[i]-func(ix[i],p))/(e[i]*e[i]);

return(add);
}


int main(int argc, char *argv[])
{
FILE *fp1,*fp2;
int i,j,k,count,ERRORS,yes;
int TIME,*label;
double *corr,dum,*w,**a,*b,*e,*ix,dum2,dum2err;
double *p,*ep,lambda,oldchi,newchi,idum,err;

/* read in correlation function data */

ERRORS=1;

(void)printf("Input lattice length\n");
scanf("%d",&L);

/* don't include t=0 data point */

TIME=L-1;

fp1=fopen(argv[1],"r");
fp2=fopen(argv[2],"w");

corr=vector(1,TIME);
ix=vector(1,TIME);
e=vector(1,TIME);
label=ivector(1,PARAM);

k=1;
for(j=1;j<=L;j++){
(void)fscanf(fp1,"%lg%lg%lg",&idum,&dum,&err);
if(j>1){
ix[k]=idum;
corr[k]=dum;
e[k]=err;
(void)printf("reading data %lg %lg %lg\n",ix[k],corr[k],e[k]);
k++;
}
}

for(i=1;i<=PARAM;i++)
label[i]=i;

/* read in initial guesses */

p=vector(1,PARAM);

(void)printf("Input initial set of parameters\n");
for(i=1;i<=PARAM;i++)
(void)scanf("%lg",&p[i]);

printf("initial parameters are:\n");
for(i=1;i<=PARAM;i++)
printf("%lg ",p[i]);
printf("\n");

/*initialize*/

lambda=0.001;

newchi=chisq(e,TIME,corr,ix,p);

count=0;

a=matrix(1,PARAM,1,PARAM);
b=vector(1,PARAM);

do{
 
count++;
oldchi=newchi;

printf("chi is %lg\n",newchi);

/* form hessian matrix */

hessian(e,TIME,corr,ix,p,lambda,a,b);

/* solve linear system ax=b */

ludcmp(a,PARAM,label,&dum);
lubksb(a,PARAM,label,b);

/* update best fit parameters */

for(i=1;i<=PARAM;i++)
p[i]+=b[i];

newchi=chisq(e,TIME,corr,ix,p);

if((newchi-oldchi)<0.0) lambda*=0.1;
else lambda*=10.0;

}
while((fabs(1.0-newchi/oldchi)>0.0001) && (count<1000));

/* inverse hessian matrix gives errors */

ep=vector(1,TIME);

lambda=0.0;
hessian(e,TIME,corr,ix,p,lambda,a,b);
ludcmp(a,PARAM,label,&dum);

for(i=1;i<=PARAM;i++){
for(j=1;j<=PARAM;j++) b[j]=0.0;

b[i]=1.0;
lubksb(a,PARAM,label,b);

ep[i]=sqrt(b[i]);
}

for(i=1;i<=PARAM;i++)
(void)printf("p[%d] = %lg +/- %lg \n",i,p[i],ep[i]);
(void)printf("chi square per dof is %lg\n",newchi/(TIME-PARAM));

for(i=1;i<=(TIME*100);i++)
(void)fprintf(fp2,"%lg %lg\n",i/100.0,func(i/100.0,p)); 

free_vector(e,1,TIME);
free_matrix(a,1,PARAM,1,PARAM);
free_vector(corr,1,TIME);
free_vector(ix,1,TIME);
free_ivector(label,1,PARAM);
free_vector(b,1,PARAM);
free_vector(p,1,PARAM);
free_vector(ep,1,PARAM);

fclose(fp2);
fclose(fp1);

return(0);
}


void hessian(double *w, int TIME, double *corr,
double *ix,
double *p, double lambda, double **a, double *b)
{
int i,j,k;
extern int PARAM;
double **df,*r;

df=matrix(1,PARAM,1,TIME);
r=vector(1,TIME);

for(i=1;i<=TIME;i++){
r[i]=corr[i]-func(ix[i],p);}

for(i=1;i<=PARAM;i++)
for(j=1;j<=TIME;j++)
df[i][j]=dfunc(i,ix[j],p);

for(i=1;i<=PARAM;i++)
for(j=1;j<=PARAM;j++)
{
a[i][j]=0.0;
for(k=1;k<=TIME;k++)
a[i][j]+=df[i][k]*df[j][k]/(w[k]*w[k]);}

for(i=1;i<=PARAM;i++)
a[i][i]=(1.0+lambda)*a[i][i];

for(i=1;i<=PARAM;i++){
b[i]=0.0;
for(k=1;k<=TIME;k++)
b[i]+=r[k]*df[i][k]/(w[k]*w[k]);}

free_matrix(df,1,PARAM,1,TIME);
free_vector(r,1,TIME);

return;
}


void lubksb(double **a, int n, int *indx, double b[])

/* uses back substitution to solve linear LU system */
{
int i,ii=0,ip,j;
double sum;

for(i=1;i<=n;i++){
ip=indx[i];
sum=b[ip];
b[ip]=b[i];
if (ii)
for(j=ii;j<=i-1;j++) sum -=a[i][j]*b[j];
else if (sum) ii=i;
b[i]=sum;
}
for(i=n;i>=1;i--){
sum=b[i];
for(j=i+1;j<=n;j++) sum -=a[i][j]*b[j];
b[i]=sum/a[i][i];
}
}




#define TINY 1.0e-20

void ludcmp(double **a, int n, int *indx, double *d)

/* LU decomposition of matrix */

{
int i,imax,j,k;
double big,dum,sum,temp;
double *vv;

vv=vector(1,n);
*d=1.0;
for(i=1;i<=n;i++){
big=0.0;
for(j=1;j<=n;j++)
if((temp=fabs(a[i][j])) > big) big=temp;
if(big ==0.0 ) nrerror("Singular matrix in routine ludcmp");
vv[i]=1.0/big;
}

for(j=1;j<=n;j++){
for(i=1;i<j;i++){
sum=a[i][j];
for(k=1;k<i;k++) sum-=a[i][k]*a[k][j];
a[i][j]=sum;
}
big=0.0;
for(i=j;i<=n;i++){
sum=a[i][j];
for(k=1;k<j;k++)
sum -=a[i][k]*a[k][j];
a[i][j]=sum;
if((dum=vv[i]*fabs(sum)) >=big ){
big=dum;
imax=i;
}
}
if(j!=imax){
for(k=1;k<=n;k++){
dum=a[imax][k];
a[imax][k]=a[j][k];
a[j][k]=dum;
}
*d=-(*d);
vv[imax]=vv[j];
}
indx[j]=imax;
if(a[j][j]==0.0) a[j][j]=TINY;
if(j!=n){
dum=1.0/a[j][j];
for(i=j+1;i<=n;i++) a[i][j]*=dum;
}
}

free_vector(vv,1,n);

}



#define NR_END 1
#define FREE_ARG char*

void nrerror(char error_text[])
{
fprintf(stderr,"Numerical Recipes run-time error ...\n");
fprintf(stderr,"%s\n",error_text);
fprintf(stderr,"...now exiting to system ...\n");
exit(1);
}

double *vector(long nl, long nh)
/* allocate a double vector with subscript range v[nl..nh] */
{
double *v;

v=(double *)malloc((size_t)((nh-nl+1+NR_END)*sizeof(double)));
if (!v) nrerror("allocation failure in vector()");
return v-nl+NR_END;
}

int *ivector(long nl, long nh)
/* allocate an int vector with subscript range v[nl..nh] */
{
int *v;

v=(int *)malloc((size_t) ((nh-nl+1+NR_END)*sizeof(int)));
if(!v) nrerror("allocation failure in ivector()");
return v-nl+NR_END;
}

double **matrix(long nrl, long nrh, long ncl, long nch)
/* allocate a double matrix with subscript range m[nrl..nrh][ncl..nch] */
{
long i,nrow=nrh-nrl+1,ncol=nch-ncl+1;
double **m;

/* allocate pointers to rows */

m=(double **) malloc((size_t)((nrow+NR_END)*sizeof(double *)));
if(!m) nrerror("allocation failure 1 in matrix()");
m+= NR_END;
m-= nrl;

/* allocate rows and set pointers to them */
m[nrl]=(double *) malloc((size_t)((nrow*ncol+NR_END)*sizeof(double)));
if(!m[nrl]) nrerror("allocation failure 2 in matrix()");
m[nrl]+= NR_END;
m[nrl]-= ncl;

for(i=nrl+1;i<=nrh;i++) m[i]=m[i-1]+ncol;

/* return pointer to array of pointers to rows */
return m;
}

void free_vector(double *v, long nl, long nh)
/* free a double vector allocated with vector() */
{
free((FREE_ARG) (v+nl-NR_END));
}

void free_ivector(int *v, long nl, long nh)
/* free an int vector allocated with ivector() */
{
free((FREE_ARG) (v+nl-NR_END));
}

void free_matrix(double **m, long nrl, long nrh, long ncl, long nch)
/* free a double matrix allocated by matrix () */
{
free((FREE_ARG) (m[nrl]+ncl-NR_END));
free((FREE_ARG) (m+nrl-NR_END));
}



double pythag(double a, double b)
{
double absa,absb;
absa=fabs(a);
absb=fabs(b);
if(absa > absb) return absa*sqrt(1.0+SQR(absb/absa));
else return(absb == 0.0 ? 0.0 : absb*sqrt(1.0+SQR(absa/absb)));
}