#include <stdio.h>
#include <string.h>
#include <time.h>
#include <math.h>
#include "include.h"
#include "run.h"
doublereal dot2(const doublereal *,const doublereal *,const doublereal *,int);

#define PRINTIT 1==1 /*  Print timing and convergence information  */

/** This file runs the lanczos algorithm for finding the lowest
    Eigenvalues of the Hamiltonian with extrapolation for nonzero P_1.
    The results are then analyzed for consistant dispersion relations,
    and a best fit is found based on this criterion.
**/

#define NJPC 2 /* Number of JPC sectors */
#define TOTALSTATES 8 /*Total number of states in all NJPC sectors */
#define NP1 2  /* Number of different values of transverse momentum */
#if 1==0  /* Truncations used in paper */
#define NS 7
  const static int kt[NS]={20,20,20,22,24,26,28}, np[NS]={4,6,8,6,6,6,6};
#elif 1==1
#define NS 5
  const static int kt[NS]={14,10,10,12,14}, np[NS]={4,6,8,6,6};
#else
#define NS 5
  const static int kt[NS]={20,20,22,22,32}, np[NS]={6,8,6,8,6};
#endif
const int nvals[NJPC]={4,4};
static struct sectors s[NJPC][NS],t[NJPC][NS];
#if 1==0
  element par[NPARAMS]={0.1,1.0,-0.2,-0.23,10.0}, rescale=1.0;
#elif 1==1 /* Best fit couplings from paper  */
  element par[NPARAMS]={0.0649839, 1.,-0.201606821,-0.550384068,116.657},
	  rescale=3.46911443596316;
#elif 1==0 /*  This is the new best fit using shuffle */
  element par[NPARAMS]={0.064984, 1.000000,-0.147802,-0.611392,116.657000},
	  rescale=3.410446;
#else /*  This is the new best fit using larger basis. */
  element par[NPARAMS]={0.064984, 1.000000, -0.152583, -0.582787, 116.657000},
	     rescale=3.552178;
#endif
static int wpar[NPARAMS];
static doublereal longvals[NP1][TOTALSTATES],
  p1a[NP1][HINDEX]={{0.0},{0.25}};
static const doublereal 
  err2[TOTALSTATES]={1.0,10000.0,10000.0,1.0,1.0,1.0,10000.0,10000.0};
static const doublereal teper[TOTALSTATES]={
       16.63019968814394, 39.77813397880572, 45.7690369882667,
          45.7690369882667,
       35.73301963006418, 55.454956994115, 61.83048213326445, 
          61.83048213326445},
	 specterr2[TOTALSTATES]={
        1.020792692835208, 7.16636029411156, 7.995737368805813,
	     7.995737368805813, 
       16.58584836076752, 65.06810610315358, 67.44195640455192,
	    67.44195640455192};

   
int MAIN(){
  char out[]="disperse3.out";
  integer totalstates=TOTALSTATES,npar=1,maxfev,mode=1,
    ifail=0,info,nfev,ldfjac=TOTALSTATES;
  integer ipvt[NPARAMS];
  /*  fcntol should be a little less than the Lanczos accuracy
      tol is the desired precision of the coupling constants  */
  doublereal tol=0.001,fcntol=1.0e-7,factor=100.0;
  doublereal vpar[NPARAMS],fvec[TOTALSTATES],diag[NPARAMS],
    fjac[TOTALSTATES*NPARAMS],qtf[NPARAMS],
    wa1[NPARAMS],wa2[NPARAMS],wa3[NPARAMS],wa4[TOTALSTATES];
  
  clock_t t1;
  int i,j,o[NJPC]={1,1},o2[NJPC]={-1,-1},ht[HINDEX]={0},
	     multi[NJPC]={1,2},multi2[NJPC]={2,1};
  int new=(1==1);
  element chi2;
  FILE *fp;

#ifdef BABBAGE  /* Routine to initialize Fortran on Babbage */
  if(hf_fint((char*)NULL)) exit(1);
#endif

  loadelements();
  if((fp = fopen(out,"w")) == NULL){
     fprintf(stderr,__FILE__":  Can't open file %s in monster, exiting\n",out);
    return;
  }
  for(i=0, j=0; i<NJPC; i++)j+=nvals[i];
  if(j!=TOTALSTATES){
     fprintf(stderr,__FILE__":  Number of states don't match in Jessica\n");
    return;
  }
  fprintf(fp,"*** Eigenvalues for the 2+1 transverse lattice  ***\n");
  fprintf(fp,"*** Best fit to Teper's lattice spectrum        ***\n");
  fprintf(fp,"%i starting couplings: (m^2, G^2 N, lambda_i/g^2 N,...):\n  ",
	  NIN);
  for(i=0;i<NIN;i++)fprintf(fp," %f",par[i]);
  fprintf(fp,"\nExtrapolation using (K,p)=");
  for(i=0;i<NS;i++)fprintf(fp,"(%i/2,%i) ",kt[i],np[i]);
  fprintf(fp,"\n");
  fprintf(fp,"Multiplets (o,multiplet, # states):  ");
  for(i=0;i<NJPC;i++)fprintf(fp,"(%i,%i,%i) ",o[i],multi[i],nvals[i]);
  fprintf(fp,"\nThe following is chi^2, # steps, rescaling factor\n");
  fprintf(fp,"The %i couplings (without rescaling) and which were fit\n\n",
	  NIN);
  fprintf(fp,"Then the values of P_1*a are shown along with the spectra\n");

  printf("Constructing %i bases",2*NJPC*NS);
  t1=clock();
  for(j=0; j<NJPC; j++)
    for(i=0; i<NS; i++){
      initbasis(s[j]+i,o[j],ht,kt[i],np[i],multi[j],new);
      initbasis(t[j]+i,o2[j],ht,kt[i],np[i],multi2[j],new);
    }
  new=(1==0);
  printf(" in %f CPU seconds\n",
       (float)(clock()-t1)/(float) CLOCKS_PER_SEC);

  for(;;){
    printf("How many to fit (0 to %i, -1=stop)\n",NIN-2);
    scanf("%li",&npar);
    if(npar<0)
      break;
    else if(npar==0)
      fcn(&totalstates,&npar,vpar,fvec,&info);
    else{
      printf("Input which %li couplings (0=m^2, ...)\n",npar);
      for(i=0;i<npar;i++)scanf("%i",wpar+i);
      for(i=0;i<npar;i++)vpar[i]=par[wpar[i]];
      maxfev=200*(npar+1);

/*      subroutine lmdif(fcn,m,n,x,fvec,ftol,xtol,gtol,maxfev,epsfcn,
  *                 diag,mode,factor,nprint,info,nfev,fjac,ldfjac,
  *                 ipvt,qtf,wa1,wa2,wa3,wa4)  */
 
    
      LMDIF(FCN_TO_FORTRAN,&totalstates,&npar,vpar,fvec,&tol,&tol,
	    &fcntol,&maxfev,&fcntol,
	    diag,&mode,&factor,&ifail,&info,&nfev,fjac,&ldfjac,
	    ipvt,qtf,wa1,wa2,wa3,wa4);
      if(info<1||info>4) fprintf(stderr,__FILE__":  Error in lmdif info=%li\n",info);
      for(i=0; i<npar; i++)par[wpar[i]]=vpar[i];
    }
    chi2=0.0;
    for(i=0; i<totalstates; i++)chi2+=fvec[i]*fvec[i];
    printf("chi2=%e in %li steps with rescale=%e\n",chi2,npar?nfev:1,rescale);
    fprintf(fp,"%e %li %e\n",chi2,nfev,rescale);
    printf("Final couplings: ");
    for(i=0; i<NIN; i++){
      printf(" %f",par[i]); fprintf(fp,"%f ",par[i]);
    }
    if(npar>0)printf("\nwhile fitting");
    for(i=0; i<npar; i++){
      printf(" %i",wpar[i]); fprintf(fp," %i",wpar[i]);
    }
    printf("\nMomenta: "); fprintf(fp,"\n");
    for(j=0; j<NP1; j++){
      printf(" %f",p1a[j][0]); fprintf(fp," %f",p1a[j][0]);
    }
    printf("\n"); fprintf(fp,"\n");
    for(i=0; i<totalstates; i++){
      for(j=0; j<NP1; j++){
	printf(" %f",longvals[j][i]); fprintf(fp," %f",longvals[j][i]);
      }
      printf("\n"); fprintf(fp,"\n");
    }
    fprintf(fp,"\n"); printf("\n");
  }
  fclose(fp);
  return 0;
}

extern int fcn(integer *m, integer *npar, doublereal *xp, 
	       doublereal *fvecc, integer *iflag){
  int i,k,j,l;
  int ipt[HINDEX]={0};
  clock_t tf;
  element *vals,sum=0.0,chi2;
  tf=clock();
  for(i=0; i<*npar; i++)par[wpar[i]]=xp[i];

#if PRINTIT
  printf("  Parameters: ");
  for(i=0; i<NIN; i++)printf(" %f",par[i]);
#endif

  for(l=0; l<NP1; l++)
    for(i=0, j=0; i<NJPC; i++){
      vals=extrapolate(NS,s[i],t[i],nvals[i],NIN,par,ipt,p1a[l]);
      for(k=0; k<nvals[i]; k++, j++)longvals[l][j]=vals[k];  
    }
  if(j!=*m) fprintf(stderr,__FILE__":  fcn: j too big\n");

#if 1==1 /*Find scale by looking only at lowest state */
  rescale=teper[0]/longvals[0][0];
#else  /* Find scale by fitting all P_1=0 states to Teper */
  rescale=dot2(teper,longvals[0],specterr2,*m)/
    dot2(longvals[0],longvals[0],specterr2,*m);
#endif
  if(rescale<0.0)*iflag = -1;

  /*  Calculate the desired dispersion as either a consistancy condition
      or in terms of a fit to Teper's spectrum to determine
      the overall scale.  The chi^2 criterion is calculated in units of 
      G^2 N.   */
#if 1==0 /* find error based on consistancy */
  for(i=0, chi2=0.0; i<*m; i++){
    sum += fvecc[i] = (longvals[1][i]-longvals[0][i])/sqrt(err2[i]);
    chi2 += 1.0/sqrt(err2[i]);
  }
  sum/=chi2;
  for(i=0; i<TOTALSTATES; i++)
    fvecc[i] = fvecc[i]/sum-1.0; 
#else  /*find error based on speed of light from comparing with Teper*/
  sum=rescale*pow(0.1974,2)*(pow(p1a[1][0],2)-pow(p1a[0][0],2));
  for(i=0; i<*m; i++)
    fvecc[i] = ((longvals[1][i]-longvals[0][i])/sum-1.0)/sqrt(err2[i]);
#endif
  /* Calculate chi^2 and rescale spectrum */
  for(i=0, chi2=0.0; i<*m; i++){
    for(l=0; l<NP1; l++)longvals[l][i] *= rescale;
    chi2 += fvecc[i]*fvecc[i];
  }
#if PRINTIT
  printf("\n ");
  for(i=0; i<TOTALSTATES; i++)printf(" %f",fvecc[i]);
  printf("\n  Chi^2=%f rescale=%f;",chi2,rescale);
  printf(" %i extrapolations in %f CPU seconds\n\n",
	 NJPC*NP1,(float)(clock()-tf)/(float) CLOCKS_PER_SEC);
#endif
  return 0;
}


doublereal dot2(const doublereal *x,const doublereal *y,
		const doublereal *z,int n){
  int i;
  doublereal t=0;
  for(i=0; i<n; i++)t+=x[i]*y[i]/z[i];
  return t;
}